import os
import asyncio
import json
import logging
from contextlib import suppress
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Tuple, Union

import ccxt
import ccxt.async_support as ccxt_async
import pandas as pd
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query, HTTPException, status
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware

try:
    from technical.util import resample_to_interval  # type: ignore
except ImportError:  # pragma: no cover - 环境缺失依赖时自动降级
    resample_to_interval = None  # type: ignore

from .storage import (
    ensure_storage,
    read_candles,
    upsert_candles,
    get_last_timestamp,
)


LOG_LEVEL = os.environ.get("LOG_LEVEL", "INFO").upper()
logging.basicConfig(
    level=LOG_LEVEL,
    format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
)
logger = logging.getLogger("datasvc")
RESAMPLE_AVAILABLE = resample_to_interval is not None
RESAMPLE_WARNING_EMITTED = False

AGGREGATION_PLAN: Dict[str, List[str]] = {
    "1m": ["2m", "3m", "4m", "5m", "10m", "15m", "20m", "25m", "30m"],
    "1h": ["2h", "3h", "4h", "6h", "8h", "12h", "16h"],
    "1d": ["2d", "3d", "4d", "5d", "6d"],
    "1w": ["2w"],
    "1M": ["2M", "3M", "6M"],
}

CandleRow = List[Union[int, float]]


def _split_env_list(value: str) -> List[str]:
    return [item.strip() for item in value.split(",") if item.strip()]


def _unique_preserve(values: List[str]) -> List[str]:
    seen = set()
    ordered: List[str] = []
    for item in values:
        if item not in seen:
            ordered.append(item)
            seen.add(item)
    return ordered


def timeframe_to_minutes(tf: str) -> Optional[int]:
    if not tf:
        return None
    unit = tf[-1]
    try:
        value = int(tf[:-1])
    except ValueError:
        return None
    multiplier = {
        "m": 1,
        "h": 60,
        "d": 1440,
        "w": 10080,
        "M": 43200,  # 30 天近似
    }.get(unit)
    if multiplier is None:
        return None
    return value * multiplier

DATA_DIR = os.environ.get("DATA_DIR", "/data")
EXCHANGE = os.environ.get("EXCHANGE", "binance")
SYMBOLS = _split_env_list(os.environ.get("SYMBOLS", "BTC/USDT:USDT,ETH/USDT:USDT"))
if not SYMBOLS:
    SYMBOLS = ["BTC/USDT:USDT"]

_default_timeframes = ["1m", "1h", "1d", "1w", "1M"]
requested_timeframes = _split_env_list(os.environ.get("TIMEFRAMES", ",".join(_default_timeframes)))
if not requested_timeframes:
    requested_timeframes = _default_timeframes

FETCH_TIMEFRAMES = _unique_preserve(requested_timeframes)
AVAILABLE_TIMEFRAMES = list(FETCH_TIMEFRAMES)
for base_tf in FETCH_TIMEFRAMES:
    for derived_tf in AGGREGATION_PLAN.get(base_tf, []):
        if derived_tf not in AVAILABLE_TIMEFRAMES:
            AVAILABLE_TIMEFRAMES.append(derived_tf)
DERIVED_TIMEFRAMES = [tf for tf in AVAILABLE_TIMEFRAMES if tf not in FETCH_TIMEFRAMES]
AGGREGATION_TARGETS = {tf: AGGREGATION_PLAN.get(tf, []) for tf in FETCH_TIMEFRAMES}

START_FROM = os.environ.get("START_FROM", "2022-01-01")  # 首次启动拉取起始日期（UTC）
POLL_FACTOR = float(os.environ.get("POLL_FACTOR", "0.5"))  # 轮询间隔 = tf_ms * factor
BACKOFF_BASE = float(os.environ.get("BACKOFF_BASE", "2.0"))
BACKOFF_MAX = float(os.environ.get("BACKOFF_MAX", "30.0"))

VALID_SYMBOLS = set(SYMBOLS)
VALID_TIMEFRAMES = set(AVAILABLE_TIMEFRAMES)

ensure_storage(DATA_DIR)

app = FastAPI(title="Local Data Service", version="0.1.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


def tf_to_ms(tf: str) -> int:
    minutes = timeframe_to_minutes(tf)
    if minutes is None:
        logger.warning("无法解析时间周期，默认使用 60 秒", extra={"timeframe": tf})
        return 60_000
    return minutes * 60_000


def ensure_symbol_timeframe(symbol: str, timeframe: str) -> None:
    if symbol not in VALID_SYMBOLS:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"symbol 必须为 {sorted(VALID_SYMBOLS)} 之一。",
        )
    if timeframe not in VALID_TIMEFRAMES:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"tf 必须为 {sorted(VALID_TIMEFRAMES)} 之一。",
        )


def parse_start_from_ms(val: str) -> int:
    """将 START_FROM 解析成毫秒级时间戳。
    支持两种格式：
    - YYYY-MM-DD（UTC 00:00:00）
    - 整型毫秒时间戳字符串
    """
    try:
        return int(val)
    except Exception:
        pass
    try:
        dt = datetime.fromisoformat(val)  # 允许 '2022-01-01' 或 '2022-01-01T00:00:00'
    except Exception:
        # 回退到固定日期
        dt = datetime(2022, 1, 1)
    return int(dt.timestamp() * 1000)


class Hub:
    def __init__(self) -> None:
        self.subscribers: Dict[str, List[WebSocket]] = {}

    def topic(self, symbol: str, timeframe: str) -> str:
        return f"candles::{symbol}::{timeframe}"

    async def subscribe(self, ws: WebSocket, symbol: str, timeframe: str):
        topic = self.topic(symbol, timeframe)
        await ws.accept()
        self.subscribers.setdefault(topic, []).append(ws)

    def _clean(self, topic: str):
        conns = self.subscribers.get(topic, [])
        self.subscribers[topic] = [w for w in conns if not w.client_state.name == "DISCONNECTED"]

    async def publish(self, symbol: str, timeframe: str, payload: dict):
        topic = self.topic(symbol, timeframe)
        conns = self.subscribers.get(topic, [])
        if not conns:
            return
        message = json.dumps(payload, ensure_ascii=False)
        dead: List[WebSocket] = []
        for ws in conns:
            try:
                await ws.send_text(message)
            except Exception:
                dead.append(ws)
        if dead:
            self.subscribers[topic] = [w for w in conns if w not in dead]


hub = Hub()

fetch_tasks: List[asyncio.Task] = []


def resample_and_store(symbol: str, base_timeframe: str, derived_timeframes: List[str]) -> List[Tuple[str, List[CandleRow]]]:
    if not RESAMPLE_AVAILABLE or not derived_timeframes:
        return []
    base_df = read_candles(DATA_DIR, symbol, base_timeframe, None, None)
    if base_df.empty:
        return []
    base_df = base_df.copy()
    if "date" not in base_df.columns:
        base_df["date"] = pd.to_datetime(base_df["timestamp"], unit="ms", utc=True)
    base_df = base_df.sort_values("timestamp")

    updates: List[Tuple[str, List[List[float]]]] = []
    for target_tf in derived_timeframes:
        minutes = timeframe_to_minutes(target_tf)
        if minutes is None:
            logger.warning("无法解析聚合周期", extra={"target_timeframe": target_tf})
            continue
        try:
            derived_df = resample_to_interval(base_df, minutes)  # type: ignore[misc]
        except Exception:
            logger.exception(
                "聚合周期计算失败",
                extra={"symbol": symbol, "base_timeframe": base_timeframe, "target_timeframe": target_tf},
            )
            continue
        if derived_df is None or derived_df.empty:
            continue
        derived_df = derived_df.copy()
        if "timestamp" not in derived_df.columns:
            if "date" in derived_df.columns:
                dates = pd.to_datetime(derived_df["date"], utc=True, errors="coerce")
                derived_df["timestamp"] = (dates.view("int64") // 1_000_000)
            elif isinstance(derived_df.index, pd.DatetimeIndex):
                idx = derived_df.index
                if idx.tz is None:
                    idx = idx.tz_localize("UTC")
                else:
                    idx = idx.tz_convert("UTC")
                derived_df["timestamp"] = (idx.view("int64") // 1_000_000)
        if "timestamp" not in derived_df.columns:
            logger.warning(
                "聚合结果缺少 timestamp 列，已跳过",
                extra={"target_timeframe": target_tf},
            )
            continue
        derived_df = derived_df.dropna(subset=["timestamp", "open", "high", "low", "close", "volume"])
        if derived_df.empty:
            continue
        derived_df["timestamp"] = derived_df["timestamp"].astype("int64")
        derived_df = derived_df.sort_values("timestamp")
        last_ts = get_last_timestamp(DATA_DIR, symbol, target_tf)
        if last_ts is not None:
            derived_df = derived_df[derived_df["timestamp"] > last_ts]
        if derived_df.empty:
            continue
        numpy_rows = derived_df[["timestamp", "open", "high", "low", "close", "volume"]].to_numpy()
        records: List[CandleRow] = []
        for ts, o, h, l, c, v in numpy_rows:
            records.append(
                [
                    int(ts),
                    float(o),
                    float(h),
                    float(l),
                    float(c),
                    float(v),
                ]
            )
        if not records:
            continue
        upsert_candles(DATA_DIR, symbol, target_tf, records)
        updates.append((target_tf, records[-3:] if len(records) > 3 else records))
    return updates


@dataclass
class FetchState:
    symbol: str
    timeframe: str
    started_at: datetime = field(default_factory=datetime.utcnow)
    last_fetch_at: Optional[datetime] = None
    last_candle_ts: Optional[int] = None
    consecutive_errors: int = 0
    last_error: Optional[str] = None

    def to_payload(self) -> dict:
        def serialize_dt(dt: Optional[datetime]) -> Optional[str]:
            if not dt:
                return None
            return dt.replace(microsecond=0).isoformat() + "Z"

        return {
            "symbol": self.symbol,
            "timeframe": self.timeframe,
            "started_at": serialize_dt(self.started_at),
            "last_fetch_at": serialize_dt(self.last_fetch_at),
            "last_candle_ts": self.last_candle_ts,
            "consecutive_errors": self.consecutive_errors,
            "last_error": self.last_error,
        }


fetch_states: Dict[Tuple[str, str], FetchState] = {}


def build_exchange():
    if EXCHANGE.lower() == "binance":
        return ccxt_async.binance({"enableRateLimit": True})
    raise RuntimeError(f"Unsupported EXCHANGE: {EXCHANGE}")


async def fetch_loop(symbol: str, timeframe: str):
    """持续增量抓取并广播。"""
    derived_timeframes = AGGREGATION_TARGETS.get(timeframe, [])
    global RESAMPLE_WARNING_EMITTED
    if derived_timeframes and not RESAMPLE_AVAILABLE and not RESAMPLE_WARNING_EMITTED:
        logger.warning(
            "缺少 technical.util.resample_to_interval 模块，聚合时间周期生成已跳过",
            extra={"timeframe": timeframe},
        )
        RESAMPLE_WARNING_EMITTED = True

    exchange = build_exchange()
    tf_ms = tf_to_ms(timeframe)
    start_since = parse_start_from_ms(START_FROM)
    last_ts = get_last_timestamp(DATA_DIR, symbol, timeframe)
    since = max(start_since, (last_ts + tf_ms) if last_ts else start_since)
    backoff = 1.0
    state_key = (symbol, timeframe)
    fetch_states[state_key] = FetchState(symbol=symbol, timeframe=timeframe, last_candle_ts=last_ts)

    logger.info("启动拉取任务", extra={"symbol": symbol, "timeframe": timeframe})

    try:
        while True:
            try:
                candles = await exchange.fetch_ohlcv(symbol, timeframe, since=since, limit=1000)
                if candles:
                    upsert_candles(DATA_DIR, symbol, timeframe, candles)
                    derived_updates: List[Tuple[str, List[CandleRow]]] = []
                    if derived_timeframes and RESAMPLE_AVAILABLE:
                        derived_updates = await asyncio.to_thread(
                            resample_and_store,
                            symbol,
                            timeframe,
                            derived_timeframes,
                        )
                    for row in candles[-3:]:
                        payload = {
                            "topic": f"candles.{symbol}.{timeframe}",
                            "type": "upsert",
                            "data": {
                                "t": row[0],
                                "o": row[1],
                                "h": row[2],
                                "l": row[3],
                                "c": row[4],
                                "v": row[5],
                            },
                        }
                        await hub.publish(symbol, timeframe, payload)
                    for target_tf, rows in derived_updates:
                        if not rows:
                            continue
                        for row in rows:
                            ts = int(row[0])
                            o, h, l, c, v = map(float, row[1:])
                            payload = {
                                "topic": f"candles.{symbol}.{target_tf}",
                                "type": "upsert",
                                "data": {
                                    "t": ts,
                                    "o": o,
                                    "h": h,
                                    "l": l,
                                    "c": c,
                                    "v": v,
                                },
                            }
                            await hub.publish(symbol, target_tf, payload)
                    since = candles[-1][0] + tf_ms
                    backoff = 1.0
                    state = fetch_states[state_key]
                    state.last_fetch_at = datetime.utcnow()
                    state.last_candle_ts = candles[-1][0]
                    state.consecutive_errors = 0
                    state.last_error = None
                await asyncio.sleep(max(1.0, tf_ms * POLL_FACTOR / 1000.0))
            except asyncio.CancelledError:
                raise
            except (ccxt.NetworkError, ccxt.ExchangeNotAvailable, ccxt.RequestTimeout) as exc:
                logger.warning(
                    "网络异常，准备重试",
                    extra={"symbol": symbol, "timeframe": timeframe, "error": str(exc)},
                )
                state = fetch_states[state_key]
                state.last_error = str(exc)
                state.consecutive_errors += 1
                backoff = min(backoff * BACKOFF_BASE, BACKOFF_MAX)
                await asyncio.sleep(backoff)
            except Exception as exc:
                logger.exception(
                    "抓取循环发生异常，重建客户端后重试",
                    extra={"symbol": symbol, "timeframe": timeframe},
                )
                state = fetch_states[state_key]
                state.last_error = str(exc)
                state.consecutive_errors += 1
                await asyncio.sleep(backoff)
                with suppress(Exception):
                    await exchange.close()
                exchange = build_exchange()
                backoff = min(backoff * BACKOFF_BASE, BACKOFF_MAX)
    except asyncio.CancelledError:
        logger.info("取消拉取任务", extra={"symbol": symbol, "timeframe": timeframe})
        state = fetch_states.get(state_key)
        if state:
            state.last_error = "cancelled"
        raise
    finally:
        with suppress(Exception):
            await exchange.close()
        logger.info("拉取任务退出", extra={"symbol": symbol, "timeframe": timeframe})


@app.on_event("startup")
async def on_start():
    ensure_storage(DATA_DIR)
    fetch_tasks.clear()
    for s in SYMBOLS:
        for tf in FETCH_TIMEFRAMES:
            task = asyncio.create_task(fetch_loop(s, tf), name=f"fetch::{s}::{tf}")
            fetch_tasks.append(task)


@app.on_event("shutdown")
async def on_shutdown():
    if not fetch_tasks:
        return
    logger.info("正在停止拉取任务")
    tasks = list(fetch_tasks)
    for task in tasks:
        task.cancel()
    results = await asyncio.gather(*tasks, return_exceptions=True)
    for result in results:
        if isinstance(result, Exception) and not isinstance(result, asyncio.CancelledError):
            logger.warning("任务停止时出现异常：%s", result)
    fetch_tasks.clear()


@app.get("/health")
async def health():
    now = datetime.utcnow().replace(microsecond=0).isoformat() + "Z"
    return {
        "status": "ok",
        "time": now,
        "exchange": EXCHANGE,
        "symbols": SYMBOLS,
        "base_timeframes": FETCH_TIMEFRAMES,
        "derived_timeframes": DERIVED_TIMEFRAMES,
        "timeframes": AVAILABLE_TIMEFRAMES,
        "tasks": [state.to_payload() for state in fetch_states.values()],
    }


@app.get("/api/candles")
def api_candles(
    symbol: str = Query(..., description="如 BTC/USDT:USDT"),
    tf: str = Query("1m", description="时间周期"),
    start: Optional[int] = Query(None, description="开始时间戳(ms)"),
    end: Optional[int] = Query(None, description="结束时间戳(ms)"),
):
    try:
        ensure_symbol_timeframe(symbol, tf)
        df = read_candles(DATA_DIR, symbol, tf, start, end)
        records = df.to_dict("records") if not df.empty else []
        return JSONResponse(records)
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)


@app.websocket("/ws")
async def ws_endpoint(websocket: WebSocket, symbol: str, tf: str, since: Optional[int] = None):
    if symbol not in VALID_SYMBOLS or tf not in VALID_TIMEFRAMES:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="invalid symbol/timeframe")
        return
    await hub.subscribe(websocket, symbol, tf)
    try:
        snap = read_candles(DATA_DIR, symbol, tf, since, None)
        await websocket.send_text(
            json.dumps(
                {
                    "topic": f"candles.{symbol}.{tf}",
                    "type": "snapshot",
                    "data": [
                        {"t": int(r["timestamp"]), "o": r["open"], "h": r["high"], "l": r["low"], "c": r["close"], "v": r["volume"]}
                        for _, r in snap.iterrows()
                    ],
                },
                ensure_ascii=False,
            )
        )
    except Exception:
        pass
    try:
        while True:
            await asyncio.sleep(30)
            await websocket.send_text(json.dumps({"type": "ping", "ts": int(datetime.utcnow().timestamp() * 1000)}))
    except WebSocketDisconnect:
        return


@app.get("/")
def root():
    return {
        "service": "Local Data Service",
        "exchange": EXCHANGE,
        "symbols": SYMBOLS,
        "base_timeframes": FETCH_TIMEFRAMES,
        "derived_timeframes": DERIVED_TIMEFRAMES,
        "timeframes": AVAILABLE_TIMEFRAMES,
    }