from __future__ import annotations

import base64
import io
from datetime import date, datetime
from typing import Literal

import pyarrow as pa
from fastapi import APIRouter, Depends, Query, Request, Response
from fastapi.responses import StreamingResponse
from findata_core.db import connection
from findata_ingest.reader import apply_adjustment, read_bars
from pyarrow import ipc

from findata_api import auth as A
from findata_api import metrics as M
from findata_api.errors import ApiError
from findata_api.ratelimit import record_usage

router = APIRouter(prefix="/market", tags=["market"])
INTERVAL_DS = {"1m": "market_bars.{a}.1m.raw", "5m": "market_bars.{a}.5m.raw", "1d": "market_bars.{a}.1d.raw"}
MAX_SPAN_DAYS = {"1m": 31, "5m": 186, "1d": 3660}
PAGE_MAX = 10_000


def _dataset_for(symbol: str, interval: str) -> tuple[str, str]:
    with connection() as conn:
        inst = conn.execute("SELECT asset_class FROM instruments WHERE symbol=%s ORDER BY (asset_class='stock') DESC LIMIT 1", (symbol,)).fetchone()
    asset = inst["asset_class"] if inst else "stock"
    if interval not in INTERVAL_DS:
        raise ApiError(400, "bad_interval", "interval must be one of 1m, 5m, 1d")
    return INTERVAL_DS[interval].format(a=asset), asset


@router.get("/bars")
def bars(request: Request, response: Response,
         symbol: str = Query(..., min_length=1, max_length=20), interval: Literal["1m", "5m", "1d"] = "1m",
         start: date = Query(...), end: date = Query(...), adjustment: Literal["raw", "split", "split_div"] = "raw",
         limit: int = Query(5000, ge=1, le=PAGE_MAX), cursor: str | None = None,
         format: Literal["json", "arrow", "csv"] = "json", version: int | None = None,
         p: A.Principal = Depends(A.require("read:market"))):
    sym = symbol.upper().strip()
    if end < start:
        raise ApiError(400, "bad_range", "end must be >= start")
    if (end - start).days > MAX_SPAN_DAYS[interval]:
        raise ApiError(413, "range_too_large", f"{interval} requests are limited to {MAX_SPAN_DAYS[interval]} days per call; "
                       "use POST /api/v1/exports for larger ranges (async Parquet).", max_days=MAX_SPAN_DAYS[interval])
    dataset_id, asset = _dataset_for(sym, interval)
    table, meta = read_bars(dataset_id, [sym], start, end, version=version)
    notes: list[str] = []
    if adjustment != "raw":
        table, notes = apply_adjustment(table, adjustment)
    # cursor = base64(last timestamp_utc isoformat)
    if cursor:
        try:
            after = datetime.fromisoformat(base64.urlsafe_b64decode(cursor.encode()).decode())
        except Exception:
            raise ApiError(400, "bad_cursor", "invalid cursor") from None
        import pyarrow.compute as pc
        table = table.filter(pc.greater(table.column("timestamp_utc"), pa.scalar(after, pa.timestamp("us", tz="UTC"))))
    total = table.num_rows
    page = table.slice(0, limit)
    next_cursor = None
    if total > limit:
        last = page.column("timestamp_utc")[-1].as_py()
        next_cursor = base64.urlsafe_b64encode(last.isoformat().encode()).decode()
    response.headers["X-Dataset-Id"] = dataset_id
    response.headers["X-Dataset-Version"] = str(meta.get("version"))
    M.ROWS.labels("bars").inc(page.num_rows)
    record_usage(p.label, rows=page.num_rows)
    if format == "arrow":
        sink = io.BytesIO()
        with ipc.new_stream(sink, page.schema) as w:
            w.write_table(page)
        headers = {"X-Dataset-Id": dataset_id, "X-Dataset-Version": str(meta.get("version")), "X-Next-Cursor": next_cursor or ""}
        return Response(sink.getvalue(), media_type="application/vnd.apache.arrow.stream", headers=headers)
    if format == "csv":
        import pyarrow.csv as pcsv
        sink = io.BytesIO()
        pcsv.write_csv(page.drop_columns(["quality_flags"]) if "quality_flags" in page.column_names else page, sink)
        return StreamingResponse(iter([sink.getvalue()]), media_type="text/csv",
                                 headers={"Content-Disposition": f'attachment; filename="{sym}_{interval}_{start}_{end}.csv"', "X-Next-Cursor": next_cursor or ""})
    rows = page.to_pylist()
    for r in rows:
        r["timestamp_utc"] = r["timestamp_utc"].isoformat().replace("+00:00", "Z")
        r["session_date"] = r["session_date"].isoformat()
        r["ingested_at"] = r["ingested_at"].isoformat().replace("+00:00", "Z")
    return {"data": rows, "meta": {**meta, "symbol": sym, "interval": interval, "adjustment": adjustment, "start": str(start), "end": str(end),
                                   "rows": len(rows), "rows_matching": total, "adjustment_notes": notes, "timestamp_convention": "bar_open",
                                   "license": "frd-private: no redistribution"}, "next_cursor": next_cursor}


@router.get("/sessions")
def sessions(symbol: str, interval: Literal["1m", "5m"] = "1m", start: date | None = None, end: date | None = None,
             p: A.Principal = Depends(A.require("read:market"))):
    """Per-session coverage for one symbol (bar counts vs expected regular-hours bars)."""
    sym = symbol.upper()
    dataset_id, _ = _dataset_for(sym, interval)
    with connection() as conn:
        rows = conn.execute("SELECT session_date, bar_count, expected_bars, first_ts, last_ts, flags FROM coverage WHERE dataset_id=%s AND symbol=%s "
                            "AND (%s::date IS NULL OR session_date>=%s) AND (%s::date IS NULL OR session_date<=%s) ORDER BY session_date DESC LIMIT 2000",
                            (dataset_id, sym, start, start, end, end)).fetchall()
    return {"data": rows, "meta": {"dataset_id": dataset_id, "symbol": sym}}