"""Zip archive handling and CSV → canonical Arrow conversion for FirstRate bar files.

Vendor file facts: `{TICKER}_{period}_{timeframe}_{adjustment}.txt` (update archives) or `{TICKER}_{timeframe}.txt`
(full archives), no header, `yyyy-MM-dd HH:mm:ss,O,H,L,C,V` (daily: `yyyy-MM-dd,O,H,L,C,V[,OI]`), FX intraday splits
date and time into two fields, line endings may mix LF/CRLF, nested zips occur in bundle archives.
"""
from __future__ import annotations

import io
import re
import zipfile
from collections.abc import Iterator
from datetime import datetime
from pathlib import Path

import polars as pl
import pyarrow as pa
from findata_schema import BARS_SCHEMA, Adjustment, AssetClass, Interval

VENDOR_TF = {"1min": Interval.m1, "5min": Interval.m5, "30min": Interval.m30, "1hour": Interval.h1, "1day": Interval.d1}
VENDOR_ADJ = {"UNADJUSTED": Adjustment.raw, "adj_split": Adjustment.split, "adj_splitdiv": Adjustment.split_div}
_TICKER_RE = re.compile(r"^([A-Za-z0-9.\-]+?)(?:[_-](?:full|day|week|month|1min|5min|30min|1hour|1day|UNADJUSTED|adj_split|adj_splitdiv|contin.*))*\.(txt|csv)$", re.IGNORECASE)


_TF_SUFFIX_RE = re.compile(r"_(1min|5min|30min|1hour|1day)$", re.I)


def ticker_from_filename(name: str, contract: bool = False) -> str:
    base = Path(name).name
    if contract:  # individual futures contracts: 'A6_F25_1day.txt' -> 'A6_F25'
        stem = Path(base).stem
        return _TF_SUFFIX_RE.sub("", stem).upper()
    m = _TICKER_RE.match(base)
    if m:
        return m.group(1).upper()
    return base.split("_")[0].upper()


def iter_archive_files(zip_path: Path) -> Iterator[tuple[str, bytes]]:
    """Yield (member_name, bytes) for every data file, descending into nested zips.

    Some vendor archives (option chains from 2023 on) use Deflate64, which Python's zipfile cannot decompress
    (`NotImplementedError: That compression method is not supported`). In that case the archive is extracted with
    the system `7z` into a temporary directory next to the archive and the files are streamed from disk."""
    try:
        yield from _iter_zip_members(zip_path)
    except NotImplementedError:
        yield from _iter_via_7z(zip_path)


def _iter_zip_members(zip_path: Path) -> Iterator[tuple[str, bytes]]:
    with zipfile.ZipFile(zip_path) as zf:
        for info in zf.infolist():
            if info.is_dir():
                continue
            name = info.filename
            if name.lower().endswith(".zip"):
                data = zf.read(info)
                with zipfile.ZipFile(io.BytesIO(data)) as inner:
                    for inner_info in inner.infolist():
                        if not inner_info.is_dir() and inner_info.filename.lower().endswith((".txt", ".csv")):
                            yield f"{name}!/{inner_info.filename}", inner.read(inner_info)
            elif name.lower().endswith((".txt", ".csv")):
                yield name, zf.read(info)


def _find_7z() -> str | None:
    import shutil
    for cand in ("7zz", "7z", "7za"):
        p = shutil.which(cand) or shutil.which(cand, path="/opt/homebrew/bin:/usr/local/bin")
        if p:
            return p
    return None


def _iter_via_7z(zip_path: Path) -> Iterator[tuple[str, bytes]]:
    import shutil
    import subprocess
    import tempfile
    tool = _find_7z()
    if not tool:
        raise RuntimeError(f"{zip_path.name}: unsupported zip compression (Deflate64) and no 7z binary found (brew install p7zip)")
    tmp = Path(tempfile.mkdtemp(prefix="unz-", dir=str(zip_path.parent)))
    try:
        r = subprocess.run([tool, "x", "-y", "-bd", "-bso0", "-bsp0", f"-o{tmp}", str(zip_path)], capture_output=True, text=True, timeout=6 * 3600)
        if r.returncode not in (0, 1):  # 1 = warnings
            raise RuntimeError(f"7z extraction failed for {zip_path.name}: {r.stderr[-300:]}")
        for p in sorted(tmp.rglob("*")):
            if p.is_file() and p.suffix.lower() in (".txt", ".csv"):
                yield str(p.relative_to(tmp)), p.read_bytes()
            elif p.is_file() and p.suffix.lower() == ".zip":
                yield from ((f"{p.name}!/{n}", b) for n, b in iter_archive_files(p))
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def parse_bar_file(data: bytes, symbol: str, asset_class: AssetClass, interval: Interval, adjustment: Adjustment,
                   source_id: str, ingested_at: datetime, source_tz: str = "America/New_York",
                   currency: str = "USD", venue: str | None = None, revision_id: int = 1) -> tuple[pa.Table, int]:
    """Parse one vendor CSV into the canonical bars table. Returns (table, raw_row_count).

    Duplicate (symbol, timestamp) rows are *kept* here and flagged downstream; we never drop source records silently.
    Unparseable timestamps (e.g. non-existent local time at DST spring-forward) are dropped from the typed table and
    counted in `raw_row_count - table.num_rows`.
    """
    text = data.decode("utf-8", errors="replace").replace("\r\n", "\n").replace("\r", "\n")
    lines = [ln for ln in text.split("\n") if ln.strip()]
    if lines and lines[0].lower().startswith(("datetime", "date", "timestamp")):
        lines = lines[1:]
    raw_rows = len(lines)
    if raw_rows == 0:
        return pa.table({f.name: pa.array([], type=f.type) for f in BARS_SCHEMA}, schema=BARS_SCHEMA), 0
    # field count = the most common count over a sample (vendor files occasionally contain ragged lines)
    from collections import Counter
    ncols = Counter(len(ln.split(",")) for ln in lines[:500]).most_common(1)[0][0]
    buf = "\n".join(lines).encode()
    if ncols == 7 and re.match(r"^\d{8},\d{2}:\d{2}", lines[0]):
        names = ["d", "t", "open", "high", "low", "close", "volume"]
        df = pl.read_csv(buf, has_header=False, new_columns=names, infer_schema_length=0, truncate_ragged_lines=True)
        df = df.with_columns((pl.col("d") + " " + pl.col("t")).alias("dt_raw")).drop("d", "t")
        fmt = "%Y%m%d %H:%M:%S"
    else:
        if ncols == 5:  # indices: no volume column
            names = ["dt_raw", "open", "high", "low", "close"]
        else:
            names = ["dt_raw", "open", "high", "low", "close", "volume"] + (["open_interest"] if ncols == 7 else [])
        if ncols not in (5, 6, 7):
            raise ValueError(f"unexpected column count {ncols} in {symbol}")
        df = pl.read_csv(buf, has_header=False, new_columns=names, infer_schema_length=0, truncate_ragged_lines=True)
        if ncols == 5:
            df = df.with_columns(pl.lit("0").alias("volume"))
        fmt = "%Y-%m-%d %H:%M:%S" if interval != Interval.d1 else None
    for c in ("open", "high", "low", "close", "volume"):
        df = df.with_columns(pl.col(c).str.strip_chars().cast(pl.Float64, strict=False))
    if interval == Interval.d1:
        local = pl.col("dt_raw").str.strip_chars().str.slice(0, 10).str.to_date("%Y-%m-%d", strict=False).cast(pl.Datetime("us"))
    else:
        local = pl.coalesce(
            pl.col("dt_raw").str.strip_chars().str.to_datetime(fmt, strict=False, time_unit="us"),
            pl.col("dt_raw").str.strip_chars().str.to_datetime("%Y-%m-%d %H:%M", strict=False, time_unit="us"),
        )
    df = df.with_columns(local.alias("local_naive"))
    df = df.filter(pl.col("local_naive").is_not_null() & pl.col("open").is_not_null() & pl.col("close").is_not_null())
    df = df.with_columns(
        pl.col("local_naive").dt.replace_time_zone(source_tz, ambiguous="earliest", non_existent="null").alias("ts_local")
    ).filter(pl.col("ts_local").is_not_null())
    df = df.with_columns(
        pl.col("ts_local").dt.convert_time_zone("UTC").alias("timestamp_utc"),
        pl.col("ts_local").dt.date().alias("session_date"),
    )
    n = df.height
    out = pl.DataFrame({
        "instrument_id": pl.repeat(f"{source_id}:{asset_class.value}:{symbol}", n, eager=True),
        "symbol_as_reported": pl.repeat(symbol, n, eager=True),
        "venue": pl.repeat(venue, n, eager=True, dtype=pl.Utf8),
        "asset_class": pl.repeat(asset_class.value, n, eager=True),
        "interval": pl.repeat(interval.value, n, eager=True),
        "timestamp_utc": df["timestamp_utc"],
        "session_date": df["session_date"],
        "open": df["open"], "high": df["high"], "low": df["low"], "close": df["close"],
        "volume": df["volume"].fill_null(0.0),
        "currency": pl.repeat(currency, n, eager=True),
        "trade_count": pl.repeat(None, n, eager=True, dtype=pl.Int32),
        "vwap": pl.repeat(None, n, eager=True, dtype=pl.Float64),
        "adjustment_type": pl.repeat(adjustment.value, n, eager=True),
        "source_id": pl.repeat(source_id, n, eager=True),
        "ingested_at": pl.repeat(ingested_at, n, eager=True, dtype=pl.Datetime("us", "UTC")),
        "quality_flags": pl.repeat([], n, eager=True, dtype=pl.List(pl.Utf8)),
        "revision_id": pl.repeat(revision_id, n, eager=True, dtype=pl.Int32),
    })
    return out.to_arrow().cast(BARS_SCHEMA), raw_rows


def parse_meta_pairs(data: bytes) -> list[tuple[str, float]]:
    """Splits/dividends per-ticker files: `yyyy-MM-dd,value` lines."""
    out = []
    for ln in data.decode("utf-8", errors="replace").splitlines():
        ln = ln.strip()
        if not ln or ln.lower().startswith("date"):
            continue
        parts = ln.split(",")
        if len(parts) >= 2:
            try:
                out.append((parts[0].strip(), float(parts[1])))
            except ValueError:
                continue
    return out


def parse_options_file(data: bytes, underlying: str, source_id: str, ingested_at: datetime) -> tuple[pa.Table, int]:
    """Vendor option-chain CSV (16 columns, no header): Trade Date, Strike, Expiry, Call/Put, Last, Bid, Ask, Bid IV,
    Ask IV, Open Interest, Volume, Delta, Gamma, Vega, Theta, Rho → OPTIONS_SCHEMA."""
    from findata_schema import OPTIONS_SCHEMA
    text = data.decode("utf-8", errors="replace").replace("\r\n", "\n").replace("\r", "\n")
    lines = [ln for ln in text.split("\n") if ln.strip()]
    if lines and lines[0].lower().startswith(("trade", "date")):
        lines = lines[1:]
    raw_rows = len(lines)
    empty = pa.table({f.name: pa.array([], type=f.type) for f in OPTIONS_SCHEMA}, schema=OPTIONS_SCHEMA)
    if raw_rows == 0:
        return empty, 0
    names = ["trade_date", "strike", "expiry", "call_put", "last_price", "bid", "ask", "bid_iv", "ask_iv", "open_interest",
             "volume", "delta", "gamma", "vega", "theta", "rho"]
    from collections import Counter
    ncols = Counter(len(ln.split(",")) for ln in lines[:500]).most_common(1)[0][0]
    if ncols != 16:
        raise ValueError(f"unexpected option-chain column count {ncols} for {underlying}")
    df = pl.read_csv("\n".join(lines).encode(), has_header=False, new_columns=names, infer_schema_length=0, truncate_ragged_lines=True)
    num = [c for c in names if c not in ("trade_date", "expiry", "call_put")]
    df = df.with_columns([pl.col(c).str.strip_chars().cast(pl.Float64, strict=False) for c in num])
    df = df.with_columns(
        pl.col("trade_date").str.strip_chars().str.slice(0, 10).str.to_date("%Y-%m-%d", strict=False),
        pl.col("expiry").str.strip_chars().str.slice(0, 10).str.to_date("%Y-%m-%d", strict=False),
        pl.col("call_put").str.strip_chars().str.to_lowercase(),
    ).filter(pl.col("trade_date").is_not_null() & pl.col("expiry").is_not_null() & pl.col("strike").is_not_null())
    n = df.height
    out = df.with_columns(
        pl.lit(underlying).alias("underlying"), pl.lit(source_id).alias("source_id"),
        pl.lit(ingested_at).cast(pl.Datetime("us", "UTC")).alias("ingested_at"),
    ).select(OPTIONS_SCHEMA.names)
    return out.to_arrow().cast(OPTIONS_SCHEMA), raw_rows