from __future__ import annotations

from dataclasses import dataclass
from datetime import date
import json
from pathlib import Path
import sys

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

from .config import Settings
from .tdx import is_stock_code


PRICE_SCHEMA_VERSION = 1
DAY_DTYPE = np.dtype([
    ("date", "<u4"), ("open", "<u4"), ("high", "<u4"), ("low", "<u4"),
    ("close", "<u4"), ("amount", "<f4"), ("volume", "<u4"), ("reserved", "<u4"),
])


@dataclass(frozen=True)
class PriceBuildResult:
    rebuilt: bool
    rows: int
    stocks: int
    output_path: Path


def _round_half_up_2(values: np.ndarray) -> np.ndarray:
    scaled = values * 100.0
    return np.where(scaled >= 0, np.floor(scaled + 0.5), np.ceil(scaled - 0.5)) / 100.0


def qfq_coefficients(trade_dates: np.ndarray, events: pd.DataFrame) -> tuple[np.ndarray, np.ndarray, int]:
    """Build affine QFQ coefficients: adjusted_price = A * raw_price + B."""
    if events.empty:
        return np.ones(len(trade_dates)), np.zeros(len(trade_dates)), 0
    ordered = events.sort_values("event_date").copy()
    valid_rows = []
    for row in ordered.itertuples(index=False):
        bonus = float(row.bonus_shares_per_10)
        rights = float(row.rights_shares_per_10)
        dividend = float(row.cash_dividend_per_10)
        rights_price = float(row.rights_price)
        m = (10.0 + bonus + rights) / 10.0
        c = (dividend - rights * rights_price) / 10.0
        if np.isfinite(m) and np.isfinite(c) and m > 0:
            valid_rows.append((np.datetime64(row.event_date, "D"), m, c))
    if not valid_rows:
        return np.ones(len(trade_dates)), np.zeros(len(trade_dates)), 0

    event_dates = np.array([row[0] for row in valid_rows], dtype="datetime64[D]")
    suffix_a = np.ones(len(valid_rows), dtype=float)
    suffix_b = np.zeros(len(valid_rows), dtype=float)
    a, b = 1.0, 0.0
    for index in range(len(valid_rows) - 1, -1, -1):
        _, m, c = valid_rows[index]
        a = a / m
        b = b - a * c
        suffix_a[index] = a
        suffix_b[index] = b
    positions = np.searchsorted(event_dates, trade_dates, side="right")
    output_a = np.ones(len(trade_dates), dtype=float)
    output_b = np.zeros(len(trade_dates), dtype=float)
    mask = positions < len(valid_rows)
    output_a[mask] = suffix_a[positions[mask]]
    output_b[mask] = suffix_b[positions[mask]]
    return output_a, output_b, len(valid_rows)


def _load_gbbq(settings: Settings) -> pd.DataFrame:
    if str(settings.vendor_dir) not in sys.path:
        sys.path.append(str(settings.vendor_dir))
    try:
        from pytdx.reader import GbbqReader
    except ImportError as exc:
        raise RuntimeError("缺少项目内 pytdx 读取器，无法解析通达信 gbbq 文件") from exc
    raw = GbbqReader().get_df(str(settings.gbbq_path))
    raw = raw.loc[raw["category"].eq(1)].copy()
    raw["source_code"] = np.select(
        [raw["market"].eq(0), raw["market"].eq(1)],
        ["sz" + raw["code"].astype(str), "sh" + raw["code"].astype(str)],
        default=None,
    )
    raw = raw[raw["source_code"].fillna("").map(is_stock_code)]
    raw["event_date"] = pd.to_datetime(raw["datetime"].astype(str), format="%Y%m%d", errors="coerce")
    raw = raw.dropna(subset=["event_date"])
    return raw.rename(columns={
        "hongli_panqianliutong": "cash_dividend_per_10",
        "peigujia_qianzongguben": "rights_price",
        "songgu_qianzongguben": "bonus_shares_per_10",
        "peigu_houzongguben": "rights_shares_per_10",
    })[["source_code", "event_date", "cash_dividend_per_10", "rights_price",
        "bonus_shares_per_10", "rights_shares_per_10"]]


def _stock_paths(settings: Settings) -> list[Path]:
    paths = []
    for market in ("sh", "sz"):
        folder = settings.tdx_root / "vipdoc" / market / "lday"
        paths.extend(path for path in folder.glob("*.day") if is_stock_code(path.stem))
    return sorted(paths)


def _read_stock_price(path: Path, events: pd.DataFrame) -> tuple[pd.DataFrame, int]:
    raw = np.fromfile(path, dtype=DAY_DTYPE)
    trade_dates = pd.to_datetime(raw["date"].astype(str), format="%Y%m%d", errors="coerce")
    valid = ~trade_dates.isna()
    raw = raw[valid]
    trade_dates = trade_dates[valid]
    day_values = trade_dates.to_numpy(dtype="datetime64[D]")
    latest = day_values.max()
    applicable = events.loc[events["event_date"].to_numpy(dtype="datetime64[D]") <= latest]
    a, b, used = qfq_coefficients(day_values, applicable)
    result = pd.DataFrame({
        "source_code": path.stem,
        "trade_date": trade_dates,
        "open_raw": raw["open"].astype(float) / 100.0,
        "high_raw": raw["high"].astype(float) / 100.0,
        "low_raw": raw["low"].astype(float) / 100.0,
        "close_raw": raw["close"].astype(float) / 100.0,
        "volume_shares": raw["volume"].astype("int64"),
        "amount_yuan": raw["amount"].astype(float),
        "qfq_a": a,
        "qfq_b": b,
    })
    for column in ("open", "high", "low", "close"):
        adjusted = a * result[f"{column}_raw"].to_numpy() + b
        result[f"{column}_qfq"] = _round_half_up_2(adjusted)
    result["has_adjustment"] = ~np.isclose(a, 1.0) | ~np.isclose(b, 0.0)
    return result, used


def build_daily_price(settings: Settings, force: bool = False) -> PriceBuildResult:
    paths = _stock_paths(settings)
    if not paths:
        raise RuntimeError("没有找到沪深股票日线")
    target = settings.daily_price_path
    metadata_path = settings.report_dir / "price_metadata.json"
    newest_source = max([path.stat().st_mtime for path in paths] + [settings.gbbq_path.stat().st_mtime])
    if target.exists() and target.stat().st_mtime >= newest_source and metadata_path.exists() and not force:
        metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
        if metadata.get("schema_version") == PRICE_SCHEMA_VERSION:
            return PriceBuildResult(False, metadata["rows"], metadata["stocks"], target)

    events = _load_gbbq(settings)
    grouped = {code: frame for code, frame in events.groupby("source_code")}
    target.parent.mkdir(parents=True, exist_ok=True)
    settings.report_dir.mkdir(parents=True, exist_ok=True)
    writer = None
    buffer: list[pd.DataFrame] = []
    buffer_rows = 0
    total_rows = 0
    adjusted_rows = 0
    adjusted_stocks = 0
    events_used = 0
    negative_adjusted_prices = 0
    earliest = None
    latest = None
    try:
        for index, path in enumerate(paths, 1):
            frame, used = _read_stock_price(path, grouped.get(path.stem, events.iloc[0:0]))
            if frame.empty:
                continue
            total_rows += len(frame)
            used_adjustment = bool(frame["has_adjustment"].any())
            adjusted_stocks += int(used_adjustment)
            adjusted_rows += int(frame["has_adjustment"].sum())
            events_used += used
            negative_adjusted_prices += int(frame[["open_qfq", "high_qfq", "low_qfq", "close_qfq"]].le(0).any(axis=1).sum())
            first_date, last_date = frame["trade_date"].min(), frame["trade_date"].max()
            earliest = first_date if earliest is None else min(earliest, first_date)
            latest = last_date if latest is None else max(latest, last_date)
            buffer.append(frame)
            buffer_rows += len(frame)
            if buffer_rows >= 150_000 or index == len(paths):
                batch = pd.concat(buffer, ignore_index=True)
                table = pa.Table.from_pandas(batch, preserve_index=False)
                if writer is None:
                    writer = pq.ParquetWriter(target, table.schema, compression="zstd")
                writer.write_table(table, row_group_size=150_000)
                buffer.clear()
                buffer_rows = 0
    finally:
        if writer is not None:
            writer.close()

    metadata = {
        "schema_version": PRICE_SCHEMA_VERSION,
        "build_date": date.today().isoformat(),
        "rows": total_rows,
        "stocks": len(paths),
        "earliest_trade_date": earliest.date().isoformat(),
        "latest_trade_date": latest.date().isoformat(),
        "gbbq_xrxd_records": int(len(events)),
        "gbbq_events_used": events_used,
        "adjusted_stocks": adjusted_stocks,
        "adjusted_rows": adjusted_rows,
        "non_positive_adjusted_price_rows": negative_adjusted_prices,
        "price_source": "tdx_unadjusted_day",
        "adjustment_source": "tdx_gbbq_category_1",
        "adjustment_model": "affine_qfq_latest_trade_anchor",
        "warning": "长历史现金分红可能使早期前复权价为负；此类区间不得直接计算收益率。",
    }
    metadata_path.write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
    _write_price_report(settings, metadata)
    return PriceBuildResult(True, total_rows, len(paths), target)


def _write_price_report(settings: Settings, metadata: dict[str, object]) -> None:
    report = f"""# M1 行情标准化报告

生成日期：{metadata['build_date']}  
股票数：{metadata['stocks']}  
日线记录：{metadata['rows']}  
日期范围：{metadata['earliest_trade_date']} 至 {metadata['latest_trade_date']}

## 1. 价格口径

- 原始行情：通达信未复权日线。
- 除权除息：通达信 `gbbq` 中 category=1 的事件。
- 前复权模型：仿射变换 `adjusted = A × raw + B`，锚定每只股票最新交易日。
- 除权事件只对事件日前的交易日生效；晚于最新交易日的未来事件不会提前施加。
- 输出同时保留原始 OHLC、前复权 OHLC、成交量、成交额及 A/B 系数。

## 2. 覆盖结果

| 指标 | 数量 |
|---|---:|
| gbbq 沪深除权事件 | {metadata['gbbq_xrxd_records']} |
| 实际参与复权事件 | {metadata['gbbq_events_used']} |
| 发生过价格调整的股票 | {metadata['adjusted_stocks']} |
| 使用非单位系数的日线 | {metadata['adjusted_rows']} |
| 前复权 OHLC 含非正数的日线 | {metadata['non_positive_adjusted_price_rows']} |

## 3. 使用限制

长历史累计现金分红可能令早期前复权价格为零或负数。这是仿射前复权的数学结果，不是原始行情错误；相关区间不得直接计算收益率。后续因子引擎会在价格非正时自动置为缺失。
"""
    (settings.report_dir / "M1_行情标准化报告.md").write_text(report, encoding="utf-8")