"""Historical backfill pipeline for A-share data.

Pulls full history for daily K-line from Sina and stores in Parquet.
Weekly and monthly K-line are derived from stored daily data.
"""

from __future__ import annotations

import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import date, datetime, timedelta
from pathlib import Path
from typing import Optional

import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
from loguru import logger

from ashare_dp.config import Settings
from ashare_dp.core.codes import to_ts_code
from ashare_dp.core.exceptions import StorageError
from ashare_dp.core.models import (
    BACKFILLABLE_FREQS,
    INTRADAY_FREQS,
    Freq,
)
from ashare_dp.data.sources.akshare_client import AKShareClient
from ashare_dp.data.store.database import get_db
from ashare_dp.data.store.repository import KLineRepository, _standardize_df
from ashare_dp.data.store.partitioning import partition_glob

settings = Settings()

# Sina daily API returns English column names
SINA_DAILY_COLUMN_MAPPING = {
    "date": "trade_time",
    "open": "open",
    "high": "high",
    "low": "low",
    "close": "close",
    "volume": "volume",
    "amount": "amount",
}

# Sina minute API returns these columns
SINA_MIN_COLUMN_MAPPING = {
    "day": "trade_time",
    "open": "open",
    "high": "high",
    "low": "low",
    "close": "close",
    "volume": "volume",
    "amount": "amount",
}

# Legacy Chinese column mappings (for backward compatibility)
LEGACY_COLUMN_MAPPING = {
    "日期": "trade_time",
    "股票代码": "ts_code",
    "开盘": "open",
    "最高": "high",
    "最低": "low",
    "收盘": "close",
    "成交量": "volume",
    "成交额": "amount",
}

LEGACY_MIN_COLUMN_MAPPING = {
    "时间": "trade_time",
    "开盘": "open",
    "最高": "high",
    "最低": "low",
    "收盘": "close",
    "成交量": "volume",
    "成交额": "amount",
}

REQUIRED_COLS = ["ts_code", "trade_time", "open", "high", "low", "close", "volume", "amount"]


def _normalize_hist_df(df: pd.DataFrame, symbol: str, freq: Freq) -> pd.DataFrame:
    """Normalize daily K-line output to standard K-line schema."""
    df = df.copy()

    # Try English mapping first (Sina), then Chinese (legacy)
    for mapping in [SINA_DAILY_COLUMN_MAPPING, LEGACY_COLUMN_MAPPING]:
        rename_map = {}
        for src, dst in mapping.items():
            if src in df.columns:
                rename_map[src] = dst
        if rename_map:
            df = df.rename(columns=rename_map)
            break

    if "ts_code" not in df.columns:
        df["ts_code"] = to_ts_code(symbol)

    if "trade_time" in df.columns:
        df["trade_time"] = pd.to_datetime(df["trade_time"])
        df["trade_date"] = df["trade_time"].dt.date
    elif "trade_date" in df.columns:
        df["trade_date"] = pd.to_datetime(df["trade_date"]).dt.date
        df["trade_time"] = pd.to_datetime(df["trade_date"])

    # Ensure numeric columns
    for col in ["open", "high", "low", "close", "volume", "amount"]:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")

    return df


def _normalize_min_df(df: pd.DataFrame, symbol: str) -> pd.DataFrame:
    """Normalize minute K-line output to standard K-line schema."""
    df = df.copy()

    # Try English mapping first (Sina), then Chinese (legacy)
    for mapping in [SINA_MIN_COLUMN_MAPPING, LEGACY_MIN_COLUMN_MAPPING]:
        rename_map = {}
        for src, dst in mapping.items():
            if src in df.columns:
                rename_map[src] = dst
        if rename_map:
            df = df.rename(columns=rename_map)
            break

    if "ts_code" not in df.columns:
        df["ts_code"] = to_ts_code(symbol)

    if "trade_time" in df.columns:
        df["trade_time"] = pd.to_datetime(df["trade_time"])
        df["trade_date"] = df["trade_time"].dt.date

    for col in ["open", "high", "low", "close", "volume", "amount"]:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")

    return df


def _resample_daily_to_period(df: pd.DataFrame, freq: Freq) -> pd.DataFrame:
    """Resample daily K-line data to weekly or monthly.

    Args:
        df: DataFrame with standard K-line columns.
        freq: Target frequency (w1 or M1).

    Returns:
        Resampled DataFrame.
    """
    if df.empty:
        return df

    df = df.copy()
    df["trade_time"] = pd.to_datetime(df["trade_time"])
    df = df.set_index("trade_time")

    if freq == Freq.w1:
        group_key = pd.Grouper(freq="W")
    elif freq == Freq.M1:
        group_key = pd.Grouper(freq="ME")
    else:
        raise ValueError(f"Unsupported resample frequency: {freq}")

    grouped = df.groupby(["ts_code", group_key])

    result = grouped.agg({
        "open": "first",
        "high": "max",
        "low": "min",
        "close": "last",
        "volume": "sum",
        "amount": "sum",
    }).reset_index()

    result["trade_date"] = result["trade_time"].dt.date
    result["freq"] = freq.value

    return result


WRITE_BATCH = 500  # Write to disk every N stocks to limit memory


class BackfillPipeline:
    """Orchestrates historical data backfill."""

    def __init__(
        self,
        client: AKShareClient | None = None,
        repo: KLineRepository | None = None,
        max_workers: int | None = None,
    ):
        self.client = client or AKShareClient()
        self.repo = repo or KLineRepository()
        self.max_workers = max_workers or settings.backfill_workers

    def init_db(self):
        """Initialize database schema (tables and views)."""
        from ashare_dp.data.store.schema import DDL_STATEMENTS

        logger.info("Initializing database schema...")
        with get_db(read_only=False) as db:
            for ddl in DDL_STATEMENTS:
                try:
                    db.execute(ddl)
                except Exception as e:
                    logger.warning(f"DDL warning: {e}")
        logger.info("Database schema initialized")

    def load_stock_list(self) -> pd.DataFrame:
        """Fetch stock list and store in DuckDB."""
        logger.info("Loading stock list...")
        df = self.client.get_stock_list()

        if df.empty:
            logger.warning("No stocks returned from akshare")
            return df

        # Normalize columns
        df = df.rename(columns={
            "code": "symbol",
            "name": "name",
        })

        # Build ts_code from code
        def _make_ts_code(code: str) -> str:
            code = str(code).zfill(6)
            if code.startswith(("4", "8")) or code.startswith("92"):
                return f"{code}.BJ"
            elif code.startswith("6") or code.startswith("9"):
                return f"{code}.SH"
            else:
                return f"{code}.SZ"

        if "symbol" in df.columns:
            df["ts_code"] = df["symbol"].apply(_make_ts_code)
            df["exchange"] = df["ts_code"].str[-2:]

        # Upsert into DuckDB (batch insert, not row-by-row)
        if "ts_code" in df.columns and "symbol" in df.columns:
            import duckdb
            try:
                with get_db(read_only=False) as db:
                    # Build column arrays for bulk insert
                    ts_codes = df["ts_code"].tolist()
                    symbols = df["symbol"].astype(str).tolist()
                    names = df["name"].astype(str).tolist()
                    exchanges = df["exchange"].tolist()
                    db.conn.execute("""
                    INSERT OR REPLACE INTO stock_info (ts_code, symbol, name, exchange, updated_at)
                    SELECT ts_code, symbol, name, exchange, now()
                    FROM (SELECT UNNEST($1::VARCHAR[]) AS ts_code,
                                 UNNEST($2::VARCHAR[]) AS symbol,
                                 UNNEST($3::VARCHAR[]) AS name,
                                 UNNEST($4::VARCHAR[]) AS exchange)
                """, [ts_codes, symbols, names, exchanges])
            except Exception as e:
                logger.warning(f"Batch stock_info insert failed ({e}), falling back to row-by-row")
                with get_db(read_only=False) as db2:
                    for _, row in df.iterrows():
                        try:
                            db2.execute("""
                                INSERT OR REPLACE INTO stock_info (ts_code, symbol, name, exchange, updated_at)
                                VALUES (?, ?, ?, ?, CURRENT_TIMESTAMP)
                            """, (
                                row.get("ts_code"),
                                str(row.get("symbol", "")),
                                str(row.get("name", "")),
                                row.get("exchange", ""),
                            ))
                        except Exception:
                            pass

        logger.info(f"Loaded {len(df)} stocks into stock_info")
        return df

    def load_trading_calendar(self):
        """Fetch trading calendar and store in DuckDB."""
        logger.info("Loading trading calendar...")
        df = self.client.get_trading_calendar()

        if df.empty:
            logger.warning("Empty trading calendar returned")
            return

        df["trade_date"] = pd.to_datetime(df["trade_date"]).dt.date

        # Batch insert into DuckDB
        dates = [d for d in df["trade_date"]]
        weekdays = [d.weekday() for d in dates]
        years = [d.year for d in dates]
        months = [d.month for d in dates]
        try:
            with get_db(read_only=False) as db:
                db.conn.execute("""
                    INSERT OR REPLACE INTO trading_calendar (trade_date, is_trading_day, week_day, year, month)
                    SELECT trade_date, 1, week_day, year, month
                    FROM (SELECT UNNEST($1::DATE[]) AS trade_date,
                                 UNNEST($2::INTEGER[]) AS week_day,
                                 UNNEST($3::INTEGER[]) AS year,
                                 UNNEST($4::INTEGER[]) AS month)
                """, [dates, weekdays, years, months])
        except Exception as e:
            logger.warning(f"Batch trading_calendar insert failed ({e}), falling back to row-by-row")
            with get_db(read_only=False) as db2:
                for _, row in df.iterrows():
                    d = row["trade_date"]
                    try:
                        db2.execute("""
                            INSERT OR REPLACE INTO trading_calendar (trade_date, is_trading_day, week_day, year, month)
                            VALUES (?, 1, ?, ?, ?)
                        """, (d, d.weekday(), d.year, d.month))
                    except Exception:
                        pass

        logger.info(f"Loaded {len(df)} trading days into trading_calendar")

    def backfill_daily_weekly_monthly(
        self,
        symbols: list[str] | None = None,
        start_date: date | None = None,
        end_date: date | None = None,
    ) -> dict:
        """Backfill daily K-line, then derive weekly and monthly.

        Daily data is fetched from Sina API. Weekly and monthly are
        derived by resampling the daily data.

        Data is written in batches (every WRITE_BATCH stocks) to avoid
        holding all 5500+ stocks in memory at once.

        Args:
            symbols: List of stock symbols (e.g. ['000001', '600000']).
                     If None, backfills all stocks from stock_info.
            start_date: Start date for backfill (default: 1990-01-01).
            end_date: End date for backfill (default: today).

        Returns:
            Dict with summary stats per frequency.
        """
        if symbols is None:
            stocks = self.client.get_stock_list()
            symbols = [str(c).zfill(6) for c in stocks["code"].tolist()]

        if start_date is None:
            start_date = date(1990, 1, 1)
        if end_date is None:
            end_date = date.today()

        start_str = start_date.strftime("%Y%m%d")
        end_str = end_date.strftime("%Y%m%d")

        logger.info(
            f"Starting backfill: {len(symbols)} stocks, "
            f"{start_str} to {end_str}, {self.max_workers} workers"
        )

        results = {}

        # ---- Step 1: Backfill daily K-line via Sina API ----
        logger.info("Backfilling 1d (daily) from Sina...")
        completed = 0
        failed = 0
        total_records = 0
        batch_frames = []
        temp_dir = Path(settings.parquet_dir) / ".tmp_daily"
        temp_dir.mkdir(parents=True, exist_ok=True)
        batch_idx = 0

        def _flush_batch():
            nonlocal total_records, batch_idx
            if batch_frames:
                chunk = pd.concat(batch_frames, ignore_index=True)
                # Standardize: add freq, keep only standard columns
                chunk["freq"] = Freq.d1.value
                # Normalize ts_code (idempotent — never double-suffixes)
                chunk["ts_code"] = chunk["ts_code"].astype(str).apply(to_ts_code)
                chunk = _standardize_df(chunk)
                total_records += len(chunk)
                table = pa.Table.from_pandas(chunk, preserve_index=False)
                pq.write_table(table, str(temp_dir / f"batch_{batch_idx:04d}.parquet"), compression="zstd", compression_level=3)
                logger.debug(f"  Flushed batch {batch_idx}: {len(chunk)} rows")
                batch_idx += 1
                batch_frames.clear()

        def _backfill_daily(symbol: str):
            try:
                df = self.client.get_hist(
                    symbol=symbol,
                    period="daily",
                    start_date=start_str,
                    end_date=end_str,
                    adjust="qfq",
                )
                if df is not None and not df.empty:
                    return _normalize_hist_df(df, symbol, Freq.d1)
                return None
            except Exception as e:
                logger.error(f"Failed backfill {symbol} 1d: {e}")
                raise

        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            futures = {
                executor.submit(_backfill_daily, sym): sym
                for sym in symbols
            }
            for future in as_completed(futures):
                sym = futures[future]
                try:
                    df = future.result()
                    if df is not None and not df.empty:
                        batch_frames.append(df)
                    completed += 1
                except Exception:
                    failed += 1

                if len(batch_frames) >= WRITE_BATCH:
                    _flush_batch()

                if (completed + failed) % 100 == 0:
                    logger.info(
                        f"  1d: {completed + failed}/{len(symbols)} "
                        f"({completed} ok, {failed} fail)"
                    )

        _flush_batch()

        # Merge all temp batches into final Hive-partitioned output
        if batch_idx > 0:
            logger.info(f"Merging {batch_idx} temp batches into final parquet...")
            t0 = time.monotonic()
            temp_glob = str(temp_dir / "batch_*.parquet")
            import duckdb
            copy_conn = duckdb.connect()
            copy_conn.execute(f"""
                COPY (
                    SELECT *,
                        YEAR(trade_date) AS year,
                        MONTH(trade_date) AS month,
                        DAY(trade_date) AS day
                    FROM read_parquet('{temp_glob}', union_by_name=true)
                    ORDER BY trade_date, ts_code
                ) TO '{settings.parquet_dir}/kline_{Freq.d1.storage_dir}'
                (FORMAT PARQUET, COMPRESSION ZSTD, COMPRESSION_LEVEL 3,
                 PARTITION_BY (year, month, day),
                 OVERWRITE_OR_IGNORE true, FILENAME_PATTERN 'data_{{i}}')
            """)
            copy_conn.close()
            # Clean up temp files
            import shutil
            shutil.rmtree(temp_dir, ignore_errors=True)
            elapsed = time.monotonic() - t0
            logger.info(f"  Merged in {elapsed:.1f}s")

        results[Freq.d1.value] = {
            "records": total_records,
            "completed": completed,
            "failed": failed,
        }
        logger.info(
            f"  1d done: {total_records} records, "
            f"{completed} stocks ok, {failed} failed"
        )

        # ---- Step 2: Derive weekly from stored daily data ----
        logger.info("Deriving 1w (weekly) from daily data...")
        self._derive_from_daily(Freq.w1, symbols, start_str, end_str, results)

        # ---- Step 3: Derive monthly from stored daily data ----
        logger.info("Deriving 1M (monthly) from daily data...")
        self._derive_from_daily(Freq.M1, symbols, start_str, end_str, results)

        return results

    def _derive_from_daily(
        self,
        freq: Freq,
        symbols: list[str],
        start_str: str,
        end_str: str,
        results: dict,
    ):
        """Derive weekly/monthly K-lines from daily data using DuckDB SQL.

        Reads ALL daily parquet files in a single DuckDB scan,
        groups by ts_code + truncated date, and aggregates.
        Much faster than per-stock queries — avoids 5500+ individual reads.
        """
        if freq == Freq.w1:
            trunc = "week"
        elif freq == Freq.M1:
            trunc = "month"
        else:
            raise ValueError(f"Unsupported derive frequency: {freq}")

        parquet_glob = partition_glob(Freq.d1)

        sql = f"""
            SELECT
                ts_code,
                date_trunc('{trunc}', trade_time) AS trade_time,
                CAST(date_trunc('{trunc}', trade_time) AS DATE) AS trade_date,
                FIRST(open ORDER BY trade_time) AS open,
                MAX(high) AS high,
                MIN(low) AS low,
                LAST(close ORDER BY trade_time) AS close,
                SUM(volume) AS volume,
                SUM(amount) AS amount
            FROM read_parquet('{parquet_glob}',
                hive_partitioning=true, union_by_name=true)
            GROUP BY ts_code, date_trunc('{trunc}', trade_time)
            ORDER BY ts_code, trade_time
        """

        logger.info(f"Deriving {freq.value} with single DuckDB scan...")
        t0 = time.monotonic()
        try:
            import duckdb
            derive_conn = duckdb.connect()
            df = derive_conn.execute(sql).fetchdf()
            derive_conn.close()
        except Exception as e:
            logger.error(f"DuckDB derive {freq.value} failed: {e}")
            results[freq.value] = {"records": 0, "completed": 0, "failed": len(symbols)}
            return

        elapsed = time.monotonic() - t0
        logger.info(f"  Aggregated {len(df)} rows in {elapsed:.1f}s")

        if df.empty:
            results[freq.value] = {"records": 0, "completed": 0, "failed": 0}
            return

        total_records = self.repo.write_klines(df, freq)
        results[freq.value] = {
            "records": total_records,
            "completed": len(symbols),
            "failed": 0,
        }
        logger.info(
            f"  {freq.value} done: {total_records} records"
        )

    def backfill_minute(
        self,
        symbols: list[str] | None = None,
        days_back: int = 30,
    ) -> dict:
        """Backfill recent minute data for all (or specified) stocks.

        Uses Sina minute API which returns all recent data.
        We filter to the requested date range after fetching.

        Args:
            symbols: List of stock symbols. If None, backfills all.
            days_back: Number of calendar days to look back.

        Returns:
            Dict with summary stats per frequency.
        """
        if symbols is None:
            stocks = self.client.get_stock_list()
            symbols = [str(c).zfill(6) for c in stocks["code"].tolist()]

        end_date = date.today()
        lookback = end_date - timedelta(days=days_back)

        start_str = lookback.strftime("%Y%m%d")
        end_str = end_date.strftime("%Y%m%d")

        logger.info(
            f"Starting minute backfill: {len(symbols)} stocks, "
            f"{start_str} to {end_str}, {self.max_workers} workers"
        )

        results = {}

        for freq in INTRADAY_FREQS:
            period = freq.akshare_min_period
            logger.info(f"Backfilling {freq.value} ({period}min)...")

            completed = 0
            failed = 0
            total_records = 0
            batch_frames = []

            def _write_min_batch():
                nonlocal total_records
                if batch_frames:
                    chunk = pd.concat(batch_frames, ignore_index=True)
                    total_records += self.repo.write_klines(chunk, freq)
                    batch_frames.clear()

            def _backfill_min(symbol: str):
                try:
                    df = self.client.get_hist_min(
                        symbol=symbol,
                        period=period,
                        start_date=start_str,
                        end_date=end_str,
                    )
                    if df is not None and not df.empty:
                        return _normalize_min_df(df, symbol)
                    return None
                except Exception as e:
                    logger.error(f"Failed minute backfill {symbol} {freq.value}: {e}")
                    raise

            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                futures = {
                    executor.submit(_backfill_min, sym): sym
                    for sym in symbols
                }
                for future in as_completed(futures):
                    sym = futures[future]
                    try:
                        df = future.result()
                        if df is not None and not df.empty:
                            batch_frames.append(df)
                        completed += 1
                    except Exception:
                        failed += 1

                    # Write batch every WRITE_BATCH stocks
                    if len(batch_frames) >= WRITE_BATCH:
                        _write_min_batch()

                    if (completed + failed) % 100 == 0:
                        logger.info(
                            f"  {freq.value}: {completed + failed}/{len(symbols)} "
                            f"({completed} ok, {failed} fail)"
                        )

            # Write remaining for this frequency
            if batch_frames:
                chunk = pd.concat(batch_frames, ignore_index=True)
                total_records += self.repo.write_klines(chunk, freq)

            results[freq.value] = {
                "records": total_records,
                "completed": completed,
                "failed": failed,
            }
            logger.info(
                f"  {freq.value} done: {total_records} records, "
                f"{completed} stocks ok, {failed} failed"
            )

        return results