"""Visualize: load a dataset (a run from this page, a sample, a Hugging Face dataset or a file), ask questions, get charts."""
from __future__ import annotations

import csv
import io
import json
import math
import os
import re
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeout
from dataclasses import dataclass, field
from typing import Optional

from fastapi import APIRouter, HTTPException
from fastapi.responses import Response
from pydantic import BaseModel

from datainventor import visualize as V
from datainventor.llm import EndpointError
from datainventor.sources import _get, inspect

router = APIRouter(prefix="/api/vis")
MAX_ROWS = int(os.environ.get("VIS_MAX_ROWS", "200"))
MAX_UPLOAD = 6_000_000
MAX_RUNNING = int(os.environ.get("VIS_MAX_RUNNING", "4"))
WRITE_SECONDS = float(os.environ.get("VIS_WRITE_SECONDS", "20"))
CALL_SECONDS = int(os.environ.get("VIS_CALL_SECONDS", "30"))
WRITERS = ThreadPoolExecutor(max_workers=4)
SAMPLES = [("support-general", "Customer support replies"), ("hotel-reviews", "Hotel reviews"), ("medical-qa", "Medical questions"),
           ("legal-qa", "Legal questions"), ("support-banking", "Banking support"), ("african-qa", "Questions about Africa"),
           ("hindi-news", "Hindi news requests"), ("serbian-car-ads", "Serbian car ads")]
JOB_ORDER = ["prompt", "completion", "chosen", "rejected", "area", "subtopic", "topic", "persona", "shape", "difficulty", "language"]


@dataclass
class VisSet:
    id: str
    name: str
    source: str
    rows: list
    cols: dict
    created: float = field(default_factory=time.time)
    lock: threading.Lock = field(default_factory=threading.Lock)
    busy: bool = False
    presets: dict = field(default_factory=dict)


@dataclass
class Ask:
    id: str
    set_id: str
    question: str
    scope: list
    scope_label: str
    stage: str = "plan"   # plan | label | chart | write | done | failed
    done: int = 0
    total: int = 0
    result: Optional[dict] = None
    error: str = ""
    started: float = field(default_factory=time.time)


SETS: dict[str, VisSet] = {}
ASKS: dict[str, Ask] = {}
GUARD = threading.Lock()


def _ep():
    from main import make_endpoint  # the host's endpoint; visitors never supply one
    return make_endpoint()


def _token() -> str:
    return os.environ.get("HF_TOKEN", "")


# ------------------------------------------------------------------ loading
def _hub_rows(dataset: str, n: int) -> list[dict]:
    """A spread of up to n full rows through the dataset viewer: no download, every column kept."""
    if not re.fullmatch(r"[\w.-]+/[\w.-]+", dataset):
        raise HTTPException(400, "Give a dataset id like owner/name.")
    try:
        info = inspect(dataset, _token())
    except Exception as e:
        raise HTTPException(400, f"Could not open {dataset}: {str(e)[:160]}")
    total = info.get("num_rows") or 0
    chunks = 1 if not total or total <= n else min(4, max(3, math.ceil(n / 25)))  # sorted datasets need several slices
    per = math.ceil(n / chunks)
    starts = [0] if chunks == 1 else [max(0, int(total * (i + 0.5) / chunks) - per // 2) for i in range(chunks)]
    out, seen = [], set()
    for s in starts:
        off = s
        while off < s + per and len(out) < n:
            try:
                got = _get("rows", {"dataset": dataset, "config": info["config"], "split": info["split"], "offset": off, "length": min(100, s + per - off)}, _token())
            except Exception as e:
                if not out:
                    raise HTTPException(502, f"The dataset viewer did not return rows for {dataset}: {str(e)[:160]}")
                break
            rows = got.get("rows") or []
            for r in rows:
                idx = r.get("row_idx", off)
                if idx not in seen:
                    seen.add(idx)
                    out.append(r.get("row") or {})
            if len(rows) < min(100, s + per - off):
                break
            off += len(rows)
    if not out:
        raise HTTPException(400, f"{dataset} has no rows the viewer can show.")
    return out[:n]


def _parse_upload(text: str) -> list[dict]:
    t = text.lstrip("")
    s = t.lstrip()
    try:
        if s.startswith("["):
            data = json.loads(s)
        elif s.startswith("{"):
            try:
                one = json.loads(s)
                data = next((v for v in one.values() if isinstance(v, list)), [one]) if isinstance(one, dict) else []
            except ValueError:
                data = [json.loads(line) for line in s.splitlines() if line.strip()]
        else:
            csv.field_size_limit(10_000_000)
            data = list(csv.DictReader(io.StringIO(t)))
    except (ValueError, csv.Error) as e:
        raise HTTPException(400, f"Could not read that file as CSV, JSON or JSONL: {str(e)[:120]}")
    rows = [r for r in data if isinstance(r, dict)]
    if not rows:
        raise HTTPException(400, "The file has no rows. Use CSV with a header row, a JSON list of objects, or JSONL.")
    return rows


def _job_rows(jid: str) -> tuple[list[dict], str]:
    from main import JOBS
    j = JOBS.get(jid)
    if not j:
        raise HTTPException(404, "That run is no longer on this host. Download it and upload the file instead.")
    with j.lock:
        rows = list(j.rows)
    if not rows:
        raise HTTPException(400, "That run has no rows yet.")
    keys = [k for k in JOB_ORDER if k in rows[0]]
    req = (j.config.get("request") or "Your run").strip()
    return [{k: r.get(k) for k in keys} for r in rows], req if len(req) <= 80 else req[:78].rsplit(" ", 1)[0] + "…"


class LoadIn(BaseModel):
    source: str              # job | sample | hub | upload
    job: str = ""
    key: str = ""
    dataset: str = ""
    text: str = ""
    filename: str = ""
    n: int = 50


@router.get("/samples")
def samples():
    return {"samples": [{"key": k, "title": t, "repo": f"LocalLLaMA/datainventor-{k}"} for k, t in SAMPLES], "max_rows": MAX_ROWS}


@router.post("/load")
def load(body: LoadIn):
    n = max(5, min(body.n or 50, MAX_ROWS))
    note = ""
    if body.source == "job":
        raw, name = _job_rows(body.job)
        src = "This page's run"
    elif body.source == "sample":
        if body.key not in dict(SAMPLES):
            raise HTTPException(400, "Unknown sample.")
        raw, name, src = _hub_rows(f"LocalLLaMA/datainventor-{body.key}", n), dict(SAMPLES)[body.key], f"LocalLLaMA/datainventor-{body.key}"
    elif body.source == "hub":
        ds = body.dataset.strip().removeprefix("https://huggingface.co/datasets/").strip("/")
        raw, name, src = _hub_rows(ds, n), ds.split("/")[-1], ds
    elif body.source == "upload":
        if len(body.text) > MAX_UPLOAD:
            raise HTTPException(413, "That file is over 6 MB. Upload a smaller sample.")
        raw, name, src = _parse_upload(body.text), (body.filename or "Your file")[:80], "Uploaded file"
    else:
        raise HTTPException(400, "Unknown source.")
    if len(raw) > MAX_ROWS:
        note = f"Showing the first {MAX_ROWS} of {len(raw)} rows."
        raw = raw[:MAX_ROWS]
    rows, cols = V.prepare(raw)
    if not cols:
        raise HTTPException(400, "None of the columns hold text, categories or numbers this can chart.")
    sid = uuid.uuid4().hex[:10]
    vs = VisSet(id=sid, name=name, source=src, rows=rows, cols=cols, presets=V.presets(rows, cols))
    with GUARD:
        SETS[sid] = vs
        for old in sorted(SETS, key=lambda k: SETS[k].created)[:-30]:
            SETS.pop(old, None)
    return {"id": sid, "name": name, "source": src, "n": len(rows), "note": note, "columns": V.describe(rows, cols),
            "glance": V.glance(rows, cols), "suggestions": list(vs.presets), "rows": rows}


# ------------------------------------------------------------------ asking
class AskIn(BaseModel):
    question: str
    scope: list[int] = []
    scope_label: str = ""


def _column(d: dict, vals: dict, filled: set, prev: Optional[dict] = None) -> dict:
    c = {"type": "num" if d["kind"] in ("score", "count") else "cat", "derived": d["kind"], "from": d["from"],
         "guide": d.get("guide", ""), "filled": set(filled)}
    if d.get("labels"):
        c["labels"] = d["labels"]
    if d.get("pattern"):
        c["pattern"] = d["pattern"]
    if prev:
        c["filled"] |= prev.get("filled", set())
    return c


def run_ask(a: Ask, vs: VisSet, ep):
    # short per-call timeouts: a stuck provider is retried instead of holding the chart for minutes
    chat = lambda prompt, **kw: ep.chat(prompt, retries=3, timeout=CALL_SECONDS, **kw)
    try:
        ids = set(a.scope)
        rows = [r for r in vs.rows if r["_id"] in ids] if ids else list(vs.rows)
        scope = f" (this question is about a selection of {len(rows)} of {len(vs.rows)} rows: {a.scope_label or 'chosen on a chart'})" if ids else ""
        plan, err = None, ""
        pre = vs.presets.get(a.question)
        if pre:  # a starter question: its plan is already written, and labels found before are reused
            plan = {**pre, "derive": [], "chart": dict(pre["chart"])}
            for d in pre["derive"]:
                c = vs.cols.get(d["name"])
                if not (c and c.get("derived") == d["kind"] and c.get("from") == d["from"]):
                    plan["derive"].append(dict(d))
        def plan_once(temp):
            text = chat(V.plan_prompt(a.question, vs.name, rows, vs.cols, scope, err), max_tokens=700, temperature=temp)
            try:
                return V.parse_plan(text, vs.cols)
            except V.PlanError as e:
                errs.append(str(e))
                print(f"[visualize] plan rejected: {e} | reply: {text[:600]!r}", flush=True)
                return None

        errs: list = []
        for _ in range(0 if plan else 2):
            plan = V.hedged(plan_once, tries=1)
            if plan:
                break
            err = errs[-1] if errs else ""
        if not plan:
            plan = V.fallback_plan(a.question, rows, vs.cols)
        # new columns, plus earlier model-made columns this chart needs that do not cover these rows yet
        todo = [(d, rows, False) for d in plan["derive"]]
        for k in ("x", "y", "color"):
            c = vs.cols.get(plan["chart"].get(k) or "")
            if c and c.get("derived") in V.KINDS:
                missing = [r for r in rows if r["_id"] not in c.get("filled", set())]
                if missing:
                    todo.append(({"name": plan["chart"][k], "kind": c["derived"], "from": c["from"], "guide": c.get("guide", ""),
                                  "labels": c.get("labels", []), "pattern": c.get("pattern", "")}, missing, True))
        touched = {}
        if todo:
            a.stage, a.total, a.done = "label", sum(len(rs) for _, rs, _ in todo), 0
            for d, rs, refill in todo:
                base = a.done
                vals = V.fill(d, rs, chat, a.question, progress=lambda k, base=base: setattr(a, "done", base + k))
                with vs.lock:
                    for r in vs.rows:
                        if r["_id"] in vals:
                            r[d["name"]] = vals[r["_id"]]
                        elif not refill or d["name"] not in r:  # a new definition: clear values left by an old one
                            r[d["name"]] = None if d["kind"] in ("score", "count") else ""
                    vs.cols[d["name"]] = _column(d, vals, set(vals), vs.cols.get(d["name"]) if refill else None)
                touched[d["name"]] = {str(i): v for i, v in vals.items()}
        a.stage = "chart"
        with vs.lock:
            c = V.chart(rows, vs.cols, V.check_spec(plan["chart"], vs.cols))
            columns = V.describe(vs.rows, vs.cols)
        made = [{"name": d["name"], "from": d["from"]} for d, _, _ in todo]
        made = list({m["name"]: m for m in made}.values())
        a.result = {"title": plan["title"], "chart": c, "summary": "", "notes": V.notes(c, made, plan.get("fallback", False)),
                    "facts": V.facts(c), "follow_ups": plan["follow_ups"], "derived": touched, "columns": columns,
                    "scope": len(rows) if ids else 0}
        a.stage = "write"
        try:
            # the chart is already on screen; a slow write-up gives way to the exact numbers after a while
            fut = WRITERS.submit(chat, V.summary_prompt(a.question, plan["title"], c, made), max_tokens=320, temperature=0.3)
            text = fut.result(timeout=WRITE_SECONDS)
            text = re.sub(r"[*#_`]+", "", re.sub(r"^\s*[-•]\s+", "", text, flags=re.M)).strip()
            text = re.sub(r"\s*\n\s*", " ", text)
            a.result["summary"] = " ".join(re.split(r"(?<=[.!?])\s+(?=[A-Z0-9])", text)[:4])[:700]
        except (EndpointError, FutureTimeout):
            pass
        a.stage = "done"
    except (EndpointError, V.PlanError) as e:
        a.stage, a.error = "failed", str(e)[:300]
    except Exception as e:  # never leave the page polling a dead ask
        a.stage, a.error = "failed", f"Something went wrong drawing that: {str(e)[:200]}"
    finally:
        vs.busy = False


@router.post("/{sid}/ask")
def ask(sid: str, body: AskIn):
    vs = SETS.get(sid)
    if not vs:
        raise HTTPException(404, "That dataset is no longer loaded here. Load it again.")
    q = re.sub(r"\s+", " ", body.question).strip()
    if not q:
        raise HTTPException(400, "Type a question first.")
    ep = _ep()
    with GUARD:
        if vs.busy:
            raise HTTPException(409, "This dataset is still answering the last question.")
        if sum(1 for a in ASKS.values() if a.stage not in ("done", "failed")) >= MAX_RUNNING:
            raise HTTPException(429, "Several charts are being drawn right now. Try again in a moment.")
        aid = uuid.uuid4().hex[:10]
        a = Ask(id=aid, set_id=sid, question=q[:400], scope=[i for i in body.scope if isinstance(i, int)][:MAX_ROWS], scope_label=body.scope_label[:120])
        ASKS[aid] = a
        vs.busy = True
        for old in sorted(ASKS, key=lambda k: ASKS[k].started)[:-200]:
            ASKS.pop(old, None)
    threading.Thread(target=run_ask, args=(a, vs, ep), daemon=True).start()
    return {"ask": aid}


@router.get("/ask/{aid}")
def ask_status(aid: str):
    a = ASKS.get(aid)
    if not a:
        raise HTTPException(404, "That question is no longer here. Ask it again.")
    out = {"stage": a.stage, "done": a.done, "total": a.total, "error": a.error, "elapsed": round(time.time() - a.started, 1)}
    if a.result and a.stage in ("write", "done"):
        out["result"] = a.result
    return out


@router.get("/{sid}/download")
def download(sid: str, ids: str = "", format: str = "csv"):
    vs = SETS.get(sid)
    if not vs:
        raise HTTPException(404, "That dataset is no longer loaded here.")
    want = {int(x) for x in ids.split(",") if x.strip().isdigit()}
    rows = [r for r in vs.rows if not want or r["_id"] in want]
    keys = ["_id"] + [k for k, c in vs.cols.items() if c.get("derived") not in ("words", "turns")]
    name = re.sub(r"[^\w-]+", "-", vs.name.lower()).strip("-")[:40] or "rows"
    if format == "jsonl":
        body = "\n".join(json.dumps({("row" if k == "_id" else k): r.get(k) for k in keys}, ensure_ascii=False) for r in rows)
        return Response(body, media_type="application/jsonl", headers={"Content-Disposition": f'attachment; filename="{name}-{len(rows)}.jsonl"'})
    buf = io.StringIO()
    w = csv.writer(buf)
    w.writerow(["row" if k == "_id" else k for k in keys])
    for r in rows:
        w.writerow(["" if r.get(k) is None else r.get(k) for k in keys])
    return Response(buf.getvalue().encode("utf-8-sig"), media_type="text/csv", headers={"Content-Disposition": f'attachment; filename="{name}-{len(rows)}.csv"'})