from __future__ import annotations

from datetime import datetime, timedelta
from typing import Any

from backend.data.numbers import finite_number as _number
from backend.data.providers.tushare_transport import TushareError


class IndexMixin:
    def market_indices(self, requested_date: str, lookback_days: int = 45) -> dict[str, Any]:
        trade_date, _ = self.resolve_trade_context(requested_date)
        end = datetime.strptime(trade_date, "%Y%m%d")
        start_date = (end - timedelta(days=max(30, lookback_days * 2))).strftime("%Y%m%d")
        index_names = {
            "000001.SH": "上证指数",
            "399001.SZ": "深证成指",
            "399006.SZ": "创业板指",
        }
        indices = []
        for ts_code, name in index_names.items():
            rows = self.query(
                "index_daily",
                {"ts_code": ts_code, "start_date": start_date, "end_date": trade_date},
                "ts_code,trade_date,close,pct_chg,vol,amount",
            )
            rows.sort(key=lambda item: str(item.get("trade_date") or ""))
            if not rows:
                continue
            latest = rows[-1]
            close = _number(latest.get("close"))
            close_5d = _number(rows[-6].get("close")) if len(rows) >= 6 else _number(rows[0].get("close"))
            close_20d = _number(rows[-21].get("close")) if len(rows) >= 21 else _number(rows[0].get("close"))
            indices.append(
                {
                    "ts_code": ts_code,
                    "name": name,
                    "trade_date": str(latest.get("trade_date") or trade_date),
                    "close": close,
                    "pct_chg": round(_number(latest.get("pct_chg")), 3),
                    "return_5d": round((close / close_5d - 1) * 100, 3) if close_5d else 0,
                    "return_20d": round((close / close_20d - 1) * 100, 3) if close_20d else 0,
                    "amount_billion": round(_number(latest.get("amount")) / 100000, 2),
                }
            )
        if not indices:
            raise TushareError(f"No index data returned for {trade_date}")
        return {
            "trade_date": trade_date,
            "source": "tushare",
            "realtime": False,
            "precise": all(item["trade_date"] == trade_date for item in indices),
            "indices": indices,
            "aggregate": {
                "average_pct_chg": round(sum(item["pct_chg"] for item in indices) / len(indices), 3),
                "average_return_5d": round(sum(item["return_5d"] for item in indices) / len(indices), 3),
                "average_return_20d": round(sum(item["return_20d"] for item in indices) / len(indices), 3),
            },
        }

    def realtime_market_indices(self, requested_date: str) -> dict[str, Any]:
        trade_date, _ = self.resolve_trade_context(requested_date)
        index_names = {
            "000001.SH": "上证指数",
            "399001.SZ": "深证成指",
            "399006.SZ": "创业板指",
        }
        rows = self.query("rt_idx_k", {"ts_code": ",".join(index_names)}, "")
        row_map = {str(row.get("ts_code") or ""): row for row in rows}
        indices = []
        for ts_code, name in index_names.items():
            row = row_map.get(ts_code)
            if not row:
                continue
            close = _number(row.get("close"))
            previous_close = _number(row.get("pre_close"))
            if close <= 0 or previous_close <= 0:
                continue
            history = self.query(
                "index_daily",
                {
                    "ts_code": ts_code,
                    "start_date": (datetime.strptime(trade_date, "%Y%m%d") - timedelta(days=20)).strftime("%Y%m%d"),
                    "end_date": trade_date,
                },
                "ts_code,trade_date,close,pct_chg",
            )
            history.sort(key=lambda item: str(item.get("trade_date") or ""))
            previous_closes = [
                _number(item.get("close")) for item in history
                if str(item.get("trade_date") or "") < trade_date and _number(item.get("close")) > 0
            ]
            close_5d = previous_closes[-5] if len(previous_closes) >= 5 else previous_closes[0] if previous_closes else previous_close
            indices.append(
                {
                    "ts_code": ts_code,
                    "name": str(row.get("name") or name).strip(),
                    "trade_date": trade_date,
                    "close": close,
                    "pct_chg": round((close / previous_close - 1) * 100, 3),
                    "return_5d": round((close / close_5d - 1) * 100, 3) if close_5d else 0,
                    "amount_billion": round(_number(row.get("amount")) / 100000000, 2),
                }
            )
        if len(indices) != len(index_names):
            raise TushareError("Realtime index quotes are incomplete")
        return {
            "trade_date": trade_date,
            "source": "tushare_rt_idx_k",
            "realtime": True,
            "precise": True,
            "indices": indices,
            "aggregate": {
                "average_pct_chg": round(sum(item["pct_chg"] for item in indices) / len(indices), 3),
                "average_return_5d": round(sum(item["return_5d"] for item in indices) / len(indices), 3),
                "average_return_20d": 0,
            },
        }