"""End-of-day option chains: one vendor archive per (year, quarter) or `period=month`; one vendor file per underlying.
Layout: options.eod/base/underlying=<U>/period=<YYYYqN|latest>/part-*.parquet. The monthly archive replaces the
`latest` partition and drops rows already covered by a completed quarter archive (no overlap in the live set)."""
from __future__ import annotations

from datetime import UTC, date, datetime
from typing import Any

import pyarrow as pa
import pyarrow.compute as pc
from findata_core.db import connection, transaction
from findata_market_firstrate import (
    FirstRateMarketConnector,
    FRDNoData,
    iter_archive_files,
    parse_options_file,
    ticker_from_filename,
)
from findata_schema.connector import FetchRequest

from findata_ingest import queue
from findata_ingest.lake import Lake
from findata_ingest.tasks.common import record_raw_archive

DATASET = "options.eod"


def run(job: dict[str, Any], lake: Lake, log) -> dict[str, Any]:
    p = job.get("params") or {}
    run_id = job.get("run_id")
    if p.get("period"):
        params = {"endpoint": "data_file", "type": "options", "period": p["period"]}
        period_label = "latest"
    else:
        params = {"endpoint": "data_file", "type": "options", "year": str(p["year"]), "quarter": p["quarter"]}
        period_label = f"{p['year']}{p['quarter']}"
    conn_fr = FirstRateMarketConnector()
    ck = queue.get_checkpoint(f"job:{job['id']}", "state") or {}
    done = set(ck.get("done") or [])
    if ck.get("raw_path") and lake.abs(ck["raw_path"]).exists():
        from findata_schema.connector import FetchResult
        res = FetchResult(FetchRequest(DATASET, params), lake.abs(ck["raw_path"]), ck["sha256"], ck["bytes"], datetime.fromisoformat(ck["fetched_at"]), None)
        log.info("options.resume", done=len(done))
    else:
        queue.progress(job["id"], {"stage": "download", "period": period_label})
        try:
            res = conn_fr.fetch_range(FetchRequest(DATASET, params), lake.raw_dir("firstrate", DATASET))
        except FRDNoData as e:
            return {"skipped": True, "reason": str(e)[:300]}
        with transaction() as conn:
            record_raw_archive(conn, "firstrate", DATASET, params, lake.rel(res.path), res.sha256, res.bytes, res.file_count, res.vendor_last_update, run_id)
        queue.set_checkpoint(f"job:{job['id']}", "state", {"raw_path": lake.rel(res.path), "sha256": res.sha256, "bytes": res.bytes,
                                                            "fetched_at": res.fetched_at.isoformat(), "done": []})
        log.info("options.downloaded", bytes=res.bytes, files=res.file_count, period=period_label)
    # rows already covered by completed quarter archives are excluded from `latest`
    cutoff: date | None = None
    if period_label == "latest":
        with connection() as conn:
            r = conn.execute("SELECT max(date_max) AS d FROM lake_files WHERE dataset_id=%s AND version_removed IS NULL AND path NOT LIKE %s",
                             (DATASET, "%/period=latest/%")).fetchone()
            cutoff = r["d"]
    ingested_at = datetime.now(UTC)
    written, retire = [], []
    stats = {"underlyings": 0, "rows": 0, "raw_rows": 0, "files": 0, "skipped_files": 0}
    with connection() as conn:
        old_latest = {r["symbol_min"]: r["path"] for r in conn.execute("SELECT symbol_min, path FROM lake_files WHERE dataset_id=%s AND version_removed IS NULL AND path LIKE %s",
                                                                       (DATASET, f"%/period={period_label}/%")).fetchall()}
    batch_rows = 0
    for name, data in iter_archive_files(res.path):
        und = ticker_from_filename(name)
        if und in done or und.startswith("_"):
            continue
        try:
            table, raw_rows = parse_options_file(data, und, "firstrate", ingested_at)
        except Exception as e:
            stats["skipped_files"] += 1
            log.warning("options.skip_file", file=name, error=str(e)[:120])
            continue
        if cutoff is not None and table.num_rows:
            table = table.filter(pc.greater(table.column("trade_date"), pa.scalar(cutoff, pa.date32())))
        stats["raw_rows"] += raw_rows
        if table.num_rows == 0:
            done.add(und)
            continue
        written.append(lake.write_table(DATASET, table, "base", {"underlying": und, "period": period_label}, sort=False))
        if und in old_latest:
            retire.append(old_latest[und])
        stats["underlyings"] += 1
        stats["rows"] += table.num_rows
        done.add(und)
        batch_rows += table.num_rows
        if len(written) >= 1000:  # publish in batches so a crash loses little work
            lake.publish(DATASET, written, retire, run_id, note=f"options {period_label} batch")
            stats["files"] += len(written)
            written, retire = [], []
            queue.set_checkpoint(f"job:{job['id']}", "state", {**ck, "raw_path": lake.rel(res.path), "sha256": res.sha256, "bytes": res.bytes,
                                                                "fetched_at": res.fetched_at.isoformat(), "done": sorted(done)})
            queue.progress(job["id"], {"stage": "publish", **stats})
    if written:
        lake.publish(DATASET, written, retire, run_id, note=f"options {period_label}")
        stats["files"] += len(written)
    with transaction() as conn:
        conn.execute("UPDATE raw_archives SET status='published' WHERE sha256=%s", (res.sha256,))
        conn.execute("UPDATE datasets SET status='active', updated_at=now() WHERE id=%s AND status='planned'", (DATASET,))
    queue.set_checkpoint(f"job:{job['id']}", "state", {"raw_path": lake.rel(res.path), "sha256": res.sha256, "bytes": res.bytes,
                                                        "fetched_at": res.fetched_at.isoformat(), "done": sorted(done)})
    log.info("options.done", period=period_label, **stats)
    return {"period": period_label, **stats}