"""Incremental sync of an update archive (period=day|week|month) into session_date delta files.

Idempotent: the same archive (sha256) is never published twice; a session_date already live is *superseded*
(old file retired, new file added, dataset version +1) so re-fetching cannot create duplicates.
Checkpoints: dataset:<id> / last_archive, last_vendor_update.
"""
from __future__ import annotations

from collections import defaultdict
from typing import Any

import pyarrow as pa
from findata_core.db import connection, transaction
from findata_market_firstrate import FirstRateMarketConnector
from findata_market_firstrate.client import FRDNoData
from findata_schema import Interval
from findata_schema.connector import FetchRequest

from findata_ingest import queue
from findata_ingest.lake import Lake
from findata_ingest.tasks.common import (
    dataset_for_params,
    live_delta_paths_for_dates,
    record_anomalies,
    record_quality_report,
    record_raw_archive,
    upsert_coverage,
)


def run(job: dict[str, Any], lake: Lake, log) -> dict[str, Any]:
    p = job.get("params") or {}
    asset = p.get("type", "stock")
    period = p.get("period", "day")
    timeframe = p.get("timeframe", "1min")
    adjustment = p.get("adjustment") or ("contin_UNadj" if asset == "futures" else ("UNADJUSTED" if asset in ("stock", "etf") else None))
    dataset_id = job.get("dataset_id") or dataset_for_params({"type": asset, "timeframe": timeframe, "adjustment": adjustment})
    run_id = job.get("run_id")
    conn_fr = FirstRateMarketConnector()
    vendor_last = conn_fr.client.last_update(asset)
    ck = queue.get_checkpoint(f"dataset:{dataset_id}", f"last_vendor_update:{period}") or {}
    if ck.get("vendor_last_update") == vendor_last.isoformat() and not p.get("force"):
        log.info("incremental.skip", reason="vendor not updated since last run", vendor_last_update=str(vendor_last))
        return {"skipped": True, "vendor_last_update": vendor_last.isoformat()}

    req = FetchRequest(dataset_id=dataset_id, params={k: v for k, v in {"endpoint": "data_file", "type": asset, "period": period,
                                                       "timeframe": timeframe, "adjustment": adjustment}.items() if v is not None})
    queue.progress(job["id"], {"stage": "download"})
    try:
        res = conn_fr.fetch_range(req, lake.raw_dir("firstrate", dataset_id))
    except FRDNoData as e:
        queue.set_checkpoint(f"dataset:{dataset_id}", f"last_vendor_update:{period}", {"vendor_last_update": vendor_last.isoformat(), "no_data": True})
        log.info("incremental.no_data", reason=str(e)[:200])
        return {"skipped": True, "reason": str(e)[:300], "vendor_last_update": vendor_last.isoformat()}
    rel_raw = lake.rel(res.path)
    with transaction() as conn:
        raw_id, is_new = record_raw_archive(conn, "firstrate", dataset_id, req.params, rel_raw, res.sha256, res.bytes, res.file_count, res.vendor_last_update, run_id)
    if not is_new:
        res.path.unlink(missing_ok=True)  # identical bytes already stored immutably
        queue.set_checkpoint(f"dataset:{dataset_id}", f"last_vendor_update:{period}", {"vendor_last_update": vendor_last.isoformat(), "sha256": res.sha256})
        log.info("incremental.duplicate_archive", sha256=res.sha256[:12])
        return {"skipped": True, "reason": "archive already ingested", "sha256": res.sha256}
    log.info("incremental.downloaded", bytes=res.bytes, files=res.file_count, vendor_last_update=str(vendor_last))

    # normalize → group rows by session_date (across all symbols)
    queue.progress(job["id"], {"stage": "normalize", "bytes": res.bytes})
    by_date: dict[Any, list[pa.Table]] = defaultdict(list)
    anomalies: list[dict[str, Any]] = []
    symbols = 0
    raw_rows = 0
    dropped = 0
    for batch in conn_fr.normalize(res):
        rep = conn_fr.validate(batch)
        anomalies.extend(rep.anomalies)
        symbols += 1
        raw_rows += batch.raw_rows
        dropped += batch.raw_rows - batch.row_count
        t = batch.table
        for d in set(t.column("session_date").to_pylist()):
            import pyarrow.compute as pc
            by_date[d].append(t.filter(pc.equal(t.column("session_date"), pa.scalar(d, pa.date32()))))
        if symbols % 1000 == 0:
            queue.progress(job["id"], {"stage": "normalize", "symbols": symbols})
    if not by_date:
        raise RuntimeError("archive produced no rows")

    # write one delta file per session_date, then publish atomically (supersede existing dates)
    queue.progress(job["id"], {"stage": "write", "sessions": len(by_date)})
    interval = Interval({"1min": "1m", "5min": "5m", "30min": "30m", "1hour": "1h", "1day": "1d"}[timeframe])
    written = []
    tables = {}
    for d in sorted(by_date):
        t = pa.concat_tables(by_date[d])
        tables[d] = t
        written.append(lake.write_table(dataset_id, t, "delta", {"session_date": d.isoformat()}))
    with connection() as conn:
        retire = live_delta_paths_for_dates(conn, dataset_id, list(tables))
    pub = lake.publish(dataset_id, written, retire, run_id, note=f"incremental {asset} {period} {timeframe} {adjustment}")
    with transaction() as conn:
        conn.execute("UPDATE raw_archives SET status='published' WHERE id=%s", (raw_id,))
        conn.execute("UPDATE datasets SET status='active', updated_at=now() WHERE id=%s AND status='planned'", (dataset_id,))
        cov = sum(upsert_coverage(conn, dataset_id, t, interval) for t in tables.values())
        record_anomalies(conn, dataset_id, anomalies, run_id)
        record_quality_report(conn, dataset_id, run_id, sum(t.num_rows for t in tables.values()), anomalies)
    queue.set_checkpoint(f"dataset:{dataset_id}", f"last_vendor_update:{period}", {"vendor_last_update": vendor_last.isoformat(), "sha256": res.sha256})
    cp = conn_fr.checkpoint(res, len(written))
    queue.set_checkpoint(cp.scope, cp.key, cp.value)
    dates = [d.isoformat() for d in sorted(tables)]
    # downstream: derived 5-minute bars for these sessions
    if interval == Interval.m1 and asset in ("stock", "etf") and dataset_id.endswith(".raw"):
        queue.enqueue("derive", dataset_id.replace(".1m.", ".5m."), {"source_dataset": dataset_id, "target_interval": "5m", "dates": dates},
                      priority=40, dedupe_key=f"derive:{dataset_id}:5m:{dates[0]}:{dates[-1]}:{pub['version']}")
    result = {"vendor_last_update": vendor_last.isoformat(), "archive_sha256": res.sha256, "archive_bytes": res.bytes,
              "symbols": symbols, "raw_rows": raw_rows, "rows_dropped_unparseable": dropped,
              "rows_published": sum(t.num_rows for t in tables.values()), "sessions": dates, "files_added": len(written),
              "files_retired": len(retire), "coverage_rows": cov, "anomalies": len(anomalies), **pub}
    log.info("incremental.done", **{k: v for k, v in result.items() if k != "sessions"})
    return result