"""Incremental loading of selected NSE symbols directly from the source ZIP."""

from __future__ import annotations

import csv
import json
import zipfile
from collections.abc import Callable, Iterable
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path

from drift_quant.nse_data.audit import (
    EXPECTED_COLUMNS,
    TIMESTAMP_FORMAT,
    ArchiveAuditError,
    _is_valid_ohlcv,
    _select_csv_member,
)


@dataclass(frozen=True, slots=True)
class NseBar:
    """A raw NSE bar preserving its source timestamp and provenance."""

    timestamp: datetime
    symbol: str
    open: float
    high: float
    low: float
    close: float
    volume: float
    source_row: int


@dataclass(frozen=True, slots=True)
class LoadedNseDataset:
    """Selected raw observations and explicit exclusion accounting."""

    archive_path: str
    archive_sha256: str
    symbols: tuple[str, ...]
    bars_by_symbol: dict[str, tuple[NseBar, ...]]
    rows_scanned: int
    rows_selected: int
    invalid_rows_excluded: int

    @property
    def dataset_id(self) -> str:
        selection = ",".join(self.symbols)
        return f"nse15m:{self.archive_sha256[:16]}:{selection}"


def _parse_selected_row(row: dict[str, str | None], row_number: int) -> NseBar | None:
    try:
        timestamp = datetime.strptime(row["date"] or "", TIMESTAMP_FORMAT)
        open_price = float(row["open"] or "")
        high = float(row["high"] or "")
        low = float(row["low"] or "")
        close = float(row["close"] or "")
        volume = float(row["volume"] or "")
        symbol = (row["symbol"] or "").strip().upper()
    except (KeyError, TypeError, ValueError):
        return None
    if not symbol or not _is_valid_ohlcv(open_price, high, low, close, volume):
        return None
    return NseBar(
        timestamp=timestamp,
        symbol=symbol,
        open=open_price,
        high=high,
        low=low,
        close=close,
        volume=volume,
        source_row=row_number,
    )


def load_selected_symbols(
    archive_path: Path,
    archive_sha256: str,
    symbols: Iterable[str],
    *,
    start: datetime | None = None,
    end: datetime | None = None,
    progress: Callable[[str], None] | None = None,
) -> LoadedNseDataset:
    """Read only selected symbols with one streaming pass over the ZIP member.

    `start` is inclusive and `end` is exclusive. Invalid source observations are
    counted and excluded; no prices are filled or repaired.
    """

    requested = tuple(sorted({symbol.strip().upper() for symbol in symbols if symbol.strip()}))
    if not requested:
        raise ArchiveAuditError("At least one symbol is required.")
    if start is not None and end is not None and start >= end:
        raise ArchiveAuditError("Start timestamp must precede end timestamp.")
    if not archive_path.is_file():
        raise ArchiveAuditError(f"Archive does not exist: {archive_path}")

    requested_set = set(requested)
    bars: dict[str, list[NseBar]] = {symbol: [] for symbol in requested}
    rows_scanned = 0
    rows_selected = 0
    invalid_rows_excluded = 0
    with zipfile.ZipFile(archive_path) as archive:
        member = _select_csv_member(archive)
        with archive.open(member) as binary_stream:
            from io import TextIOWrapper

            with TextIOWrapper(binary_stream, encoding="utf-8", newline="") as stream:
                reader = csv.DictReader(stream)
                if tuple(reader.fieldnames or ()) != EXPECTED_COLUMNS:
                    raise ArchiveAuditError(f"Unexpected CSV schema: {reader.fieldnames!r}")
                for row in reader:
                    rows_scanned += 1
                    raw_symbol = (row.get("symbol") or "").strip().upper()
                    if raw_symbol not in requested_set:
                        continue
                    bar = _parse_selected_row(row, rows_scanned)
                    if bar is None:
                        invalid_rows_excluded += 1
                        continue
                    if (start is not None and bar.timestamp < start) or (
                        end is not None and bar.timestamp >= end
                    ):
                        continue
                    bars[bar.symbol].append(bar)
                    rows_selected += 1
                    if progress is not None and rows_scanned % 1_000_000 == 0:
                        progress(
                            f"Scanned {rows_scanned:,} rows; selected "
                            f"{rows_selected:,} valid observations."
                        )

    normalized = {
        symbol: tuple(sorted(symbol_bars, key=lambda bar: bar.timestamp))
        for symbol, symbol_bars in bars.items()
    }
    return LoadedNseDataset(
        archive_path=str(archive_path.resolve()),
        archive_sha256=archive_sha256,
        symbols=requested,
        bars_by_symbol=normalized,
        rows_scanned=rows_scanned,
        rows_selected=rows_selected,
        invalid_rows_excluded=invalid_rows_excluded,
    )


def save_selected_cache(dataset: LoadedNseDataset, output_path: Path) -> None:
    """Persist a compact, explicit cache for repeat runs over the same selection."""

    output_path.parent.mkdir(parents=True, exist_ok=True)
    serializable = {
        "archive_path": dataset.archive_path,
        "archive_sha256": dataset.archive_sha256,
        "symbols": dataset.symbols,
        "rows_scanned": dataset.rows_scanned,
        "rows_selected": dataset.rows_selected,
        "invalid_rows_excluded": dataset.invalid_rows_excluded,
        "bars_by_symbol": {
            symbol: [{**asdict(bar), "timestamp": bar.timestamp.isoformat(sep=" ")} for bar in bars]
            for symbol, bars in dataset.bars_by_symbol.items()
        },
    }
    output_path.write_text(json.dumps(serializable, separators=(",", ":")), encoding="utf-8")


def load_selected_cache(cache_path: Path) -> LoadedNseDataset:
    """Load an explicit cache created by :func:`save_selected_cache`."""

    payload = json.loads(cache_path.read_text(encoding="utf-8"))
    return LoadedNseDataset(
        archive_path=payload["archive_path"],
        archive_sha256=payload["archive_sha256"],
        symbols=tuple(payload["symbols"]),
        bars_by_symbol={
            symbol: tuple(
                NseBar(
                    timestamp=datetime.fromisoformat(bar["timestamp"]),
                    symbol=bar["symbol"],
                    open=bar["open"],
                    high=bar["high"],
                    low=bar["low"],
                    close=bar["close"],
                    volume=bar["volume"],
                    source_row=bar["source_row"],
                )
                for bar in bars
            )
            for symbol, bars in payload["bars_by_symbol"].items()
        },
        rows_scanned=payload["rows_scanned"],
        rows_selected=payload["rows_selected"],
        invalid_rows_excluded=payload["invalid_rows_excluded"],
    )