trade/engine/data/reader.py

"""
K 线数据读取器 —— 从 TimescaleDB 查询历史 K 线供策略分析和回测使用。

所有方法均为只读查询，对应 data 模块 Kline 实体的字段结构。
参考：data/db/entities/kline.entity.ts
"""

from datetime import datetime
from typing import Optional, Sequence

import asyncpg

from engine.data.db import get_pool
from engine.data.models import KlineRecord


class KlineReader:
    """TimescaleDB K 线只读查询器"""

    def __init__(self):
        self._pool: Optional[asyncpg.Pool] = None

    async def _ensure_pool(self) -> asyncpg.Pool:
        if self._pool is None:
            self._pool = await get_pool()
        return self._pool

    async def get_klines(
        self,
        symbol: str,
        interval: str,
        start_time: datetime,
        end_time: datetime,
        exchange: str = "binance",
        limit: int = 1000,
    ) -> list[KlineRecord]:
        """
        查询指定时间范围内的 K 线数据，支持任意周期（1m / 5m / 15m / 1h / 4h / 1d 等）。

        Args:
            symbol: 交易对（如 BTCUSDT）
            interval: K 线周期（1m / 5m / 15m / 30m / 1h / 4h / 1d 等）
            start_time: 起始时间（含）
            end_time: 结束时间（含）
            exchange: 交易所标识
            limit: 最大返回条数
        """
        pool = await self._ensure_pool()

        query = """
            SELECT
                time, exchange, symbol, interval,
                open, high, low, close, volume,
                quote_volume, taker_buy_base_vol, taker_buy_quote_vol,
                trade_count, is_closed, created_at, updated_at
            FROM klines
            WHERE exchange = $1
              AND symbol   = $2
              AND interval = $3
              AND time    >= $4
              AND time    <= $5
            ORDER BY time ASC
            LIMIT $6
        """

        rows: Sequence[asyncpg.Record] = await pool.fetch(
            query,
            exchange,
            symbol,
            interval,
            start_time,
            end_time,
            limit,
        )

        return [KlineRecord.from_record(row) for row in rows]

    async def get_latest_klines(
        self,
        symbol: str,
        interval: str,
        exchange: str = "binance",
        limit: int = 500,
    ) -> list[KlineRecord]:
        """
        获取最近 N 根已闭合的 K 线（策略启动时快速预热）。

        仅查询 is_closed = TRUE 的 K 线，避免使用未闭合的不完整数据。
        """
        pool = await self._ensure_pool()

        query = """
            SELECT
                time, exchange, symbol, interval,
                open, high, low, close, volume,
                quote_volume, taker_buy_base_vol, taker_buy_quote_vol,
                trade_count, is_closed, created_at, updated_at
            FROM klines
            WHERE exchange = $1
              AND symbol   = $2
              AND interval = $3
              AND is_closed = TRUE
            ORDER BY time DESC
            LIMIT $4
        """

        rows = await pool.fetch(query, exchange, symbol, interval, limit)

        records = [KlineRecord.from_record(row) for row in rows]
        records.reverse()
        return records

    async def get_klines_by_count(
        self,
        symbol: str,
        interval: str,
        count: int = 100,
        exchange: str = "binance",
        before_time: Optional[datetime] = None,
    ) -> list[KlineRecord]:
        """
        获取最新的 N 根 K 线（按时间倒序取 N 条再正序返回）。

        适合策略运行时获取最近 N 根 K 线做指标计算，不关心特定时间范围。

        Args:
            symbol: 交易对
            interval: K 线周期
            count: 需要获取的 K 线数量
            exchange: 交易所标识
            before_time: 可选，只获取该时间之前的 K 线
        """
        pool = await self._ensure_pool()

        if before_time:
            query = """
                SELECT
                    time, exchange, symbol, interval,
                    open, high, low, close, volume,
                    quote_volume, taker_buy_base_vol, taker_buy_quote_vol,
                    trade_count, is_closed, created_at, updated_at
                FROM klines
                WHERE exchange = $1
                  AND symbol   = $2
                  AND interval = $3
                  AND time     <= $4
                ORDER BY time DESC
                LIMIT $5
            """
            rows = await pool.fetch(query, exchange, symbol, interval, before_time, count)
        else:
            query = """
                SELECT
                    time, exchange, symbol, interval,
                    open, high, low, close, volume,
                    quote_volume, taker_buy_base_vol, taker_buy_quote_vol,
                    trade_count, is_closed, created_at, updated_at
                FROM klines
                WHERE exchange = $1
                  AND symbol   = $2
                  AND interval = $3
                ORDER BY time DESC
                LIMIT $4
            """
            rows = await pool.fetch(query, exchange, symbol, interval, count)

        records = [KlineRecord.from_record(row) for row in rows]
        records.reverse()
        return records

    async def get_ohlcv_array(
        self,
        symbol: str,
        interval: str,
        exchange: str = "binance",
        limit: int = 200,
        before_time: Optional[datetime] = None,
    ) -> list[tuple[datetime, float, float, float, float, float]]:
        """
        获取 OHLCV 元组数组，直接用于 pandas DataFrame 或 TA-Lib 计算。

        返回格式: [(timestamp, open, high, low, close, volume), ...]
        时间升序排列。
        """
        pool = await self._ensure_pool()

        if before_time:
            query = """
                SELECT time, open, high, low, close, volume
                FROM klines
                WHERE exchange = $1
                  AND symbol   = $2
                  AND interval = $3
                  AND time     <= $4
                ORDER BY time DESC
                LIMIT $5
            """
            rows = await pool.fetch(query, exchange, symbol, interval, before_time, limit)
        else:
            query = """
                SELECT time, open, high, low, close, volume
                FROM klines
                WHERE exchange = $1
                  AND symbol   = $2
                  AND interval = $3
                ORDER BY time DESC
                LIMIT $4
            """
            rows = await pool.fetch(query, exchange, symbol, interval, limit)

        return [
            (row["time"], float(row["open"]), float(row["high"]), float(row["low"]), float(row["close"]), float(row["volume"]))
            for row in reversed(rows)
        ]

    async def get_available_symbols(
        self,
        exchange: str = "binance",
    ) -> list[str]:
        """
        查询已激活的交易对列表。

        从 trading_pairs 配置表读取（仅 active=TRUE 的记录），
        而非扫描 klines 时序表，避免全表扫描。
        对应 data/db/entities/trading-pair.entity.ts。
        """
        pool = await self._ensure_pool()

        rows = await pool.fetch(
            """
            SELECT tp.symbol
            FROM trading_pairs tp
            JOIN exchanges e ON tp.exchange_id = e.id
            WHERE e.name = $1
              AND tp.active = TRUE
            ORDER BY tp.symbol
            """,
            exchange,
        )

        return [row["symbol"] for row in rows]

    async def get_available_intervals(
        self,
        symbol: str,
        exchange: str = "binance",
    ) -> list[str]:
        """
        查询某交易对配置的 K 线周期列表。

        从 trading_pairs.kline_intervals 读取（逗号分隔字符串），
        而非扫描 klines 时序表，避免全表扫描。
        对应 data/db/entities/trading-pair.entity.ts: kline_intervals 字段。
        """
        pool = await self._ensure_pool()

        row = await pool.fetchrow(
            """
            SELECT tp.kline_intervals
            FROM trading_pairs tp
            JOIN exchanges e ON tp.exchange_id = e.id
            WHERE e.name = $1
              AND tp.symbol = $2
              AND tp.active = TRUE
            """,
            exchange,
            symbol,
        )

        if row is None or not row["kline_intervals"]:
            return []

        # kline_intervals 格式: "1m,5m,15m,1h,4h,1d"
        return [
            interval.strip()
            for interval in row["kline_intervals"].split(",")
            if interval.strip()
        ]