"""Wizard runs: save drafts, finalize with the forecaster's own ratings, export.

The server always recomputes the engine output from the inputs (never trusting
the browser's numbers) and records which config version and engine version
produced the suggestion.
"""

from __future__ import annotations

import math
from datetime import date

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import PlainTextResponse
from psycopg import Connection
from pydantic import BaseModel, Field

from .. import engine
from ..auth import User, current_user, require_center
from ..configs import active_config, config_by_id
from ..db import Jsonb, fetchall, fetchone, get_conn
from ..mapping import AFP_PROBLEM_NAMES, AFP_TO_DRUPAL_PROBLEM, LIKELIHOOD_NAMES

router = APIRouter(tags=["runs"])


class RunIn(BaseModel):
    center_id: int
    zone_id: int | None = None
    valid_date: date
    inputs: dict
    justifications: dict[str, str] = {}
    notes: str | None = None
    bottom_line: str | None = None
    client_computed: dict | None = None
    draft_ratings: dict[str, int | None] | None = None


class RunUpdate(BaseModel):
    zone_id: int | None = None
    valid_date: date | None = None
    inputs: dict | None = None
    justifications: dict[str, str] | None = None
    notes: str | None = None
    bottom_line: str | None = None
    client_computed: dict | None = None
    draft_ratings: dict[str, int | None] | None = None


class FinalizeIn(BaseModel):
    upper: int = Field(ge=0, le=5)
    middle: int = Field(ge=0, le=5)
    lower: int = Field(ge=0, le=5)
    overall: int | None = Field(default=None, ge=0, le=5)
    bottom_line: str | None = None
    notes: str | None = None


# --- validation ------------------------------------------------------------------


def _range(value, lo: float, hi: float, label: str) -> dict:
    value = value or {}
    out = {}
    for k in ("min", "max"):
        v = value.get(k)
        if v in (None, ""):
            out[k] = None
            continue
        try:
            f = float(v)
        except (TypeError, ValueError) as e:
            raise HTTPException(422, f"{label}.{k} must be a number") from e
        if math.isnan(f) or not lo <= f <= hi:
            raise HTTPException(422, f"{label}.{k} must be between {lo} and {hi}")
        out[k] = f
    if out["min"] is not None and out["max"] is not None and out["min"] > out["max"]:
        out["min"], out["max"] = out["max"], out["min"]
    return out


MAX_DISCUSSION = 5000


def _discussion(value, i: int) -> str:
    if value is None:
        return ""
    if not isinstance(value, str):
        raise HTTPException(422, f"Problem {i + 1}: discussion must be text")
    if len(value) > MAX_DISCUSSION:
        raise HTTPException(422, f"Problem {i + 1}: discussion is longer than {MAX_DISCUSSION} characters")
    return value


def validate_inputs(inputs: dict, config: dict) -> dict:
    """Normalize wizard inputs to the canonical engine shape; 422 on bad data."""
    condition_ids = {c["id"] for c in config["conditions"]}
    type_ids = {p["id"] for p in config["problemTypes"]}
    max_size = max(s["value"] for s in config["sizes"])
    conditions = [c for c in (inputs.get("conditions") or []) if c in condition_ids]
    problems = inputs.get("problems") or []
    if len(problems) > config.get("maxProblems", 3):
        raise HTTPException(422, f"At most {config.get('maxProblems', 3)} avalanche problems")
    seen = set()
    clean = []
    for i, p in enumerate(problems):
        try:
            t = int(p.get("type"))
        except (TypeError, ValueError) as e:
            raise HTTPException(422, f"Problem {i + 1}: choose a problem type") from e
        if t not in type_ids:
            raise HTTPException(422, f"Problem {i + 1}: unknown problem type {t}")
        if t in seen:
            raise HTTPException(422, f"Problem {i + 1}: each problem type can be listed once")
        seen.add(t)
        locations = []
        for loc in p.get("locations") or []:
            parsed = engine.parse_location(loc)
            if not parsed:
                raise HTTPException(422, f"Problem {i + 1}: invalid location {loc!r}")
            key = f"{parsed[0]} {parsed[1]}"
            if key not in locations:
                locations.append(key)
        choice = p.get("likelihoodChoice")
        clean.append(
            {
                "type": t,
                "locations": locations,
                "sensitivity": _range(p.get("sensitivity"), 0, len(config["sensitivity"]) - 1, f"problems[{i}].sensitivity"),
                "distribution": _range(p.get("distribution"), 0, len(config["distribution"]) - 1, f"problems[{i}].distribution"),
                "size": _range(p.get("size"), 1, max_size, f"problems[{i}].size"),
                "likelihoodChoice": int(choice) if engine.is_num(choice) and 1 <= int(choice) <= 5 else None,
                "discussion": _discussion(p.get("discussion"), i),
            }
        )
    return {"conditions": conditions, "problems": clean}


def _differs(client: dict | None, server: dict) -> bool:
    if not client:
        return False
    try:
        if client.get("suggestions") != server["suggestions"]:
            return True
        return [p.get("likelihood") for p in client.get("problems", [])] != [p["likelihood"] for p in server["problems"]]
    except (AttributeError, TypeError):
        return True


def _clean_ratings(ratings: dict | None) -> Jsonb | None:
    """In-progress danger per band on a draft: {upper, middle, lower} each 0-5 or None."""
    if not ratings:
        return None
    out = {}
    for b in engine.BANDS:
        v = ratings.get(b)
        out[b] = v if isinstance(v, int) and 0 <= v <= 5 else None
    return Jsonb(out)


def _suggest_cols(computed: dict) -> dict:
    s = computed["suggestions"]
    return {"su": s["upper"], "sm": s["middle"], "sl": s["lower"], "so": s["overall"]}


def _load_run(conn: Connection, run_id: int, user: User) -> dict:
    run = fetchone(
        conn,
        """SELECT r.*, u.name AS user_name, cc.version AS config_version, z.name AS zone_name, c.name AS center_name
           FROM tool_runs r JOIN users u ON u.id = r.user_id JOIN center_configs cc ON cc.id = r.config_id
           JOIN centers c ON c.id = r.center_id LEFT JOIN zones z ON z.id = r.zone_id WHERE r.id = %s""",
        (run_id,),
    )
    if not run:
        raise HTTPException(404, "Run not found")
    require_center(user, run["center_id"])
    return run


def _check_zone(conn: Connection, center_id: int, zone_id: int | None) -> None:
    if zone_id is not None and not fetchone(conn, "SELECT 1 FROM zones WHERE id = %s AND center_id = %s", (zone_id, center_id)):
        raise HTTPException(422, "Zone does not belong to this center")


# --- endpoints ---------------------------------------------------------------------


@router.post("/runs")
def create_run(body: RunIn, conn: Connection = Depends(get_conn), user: User = Depends(current_user)):
    require_center(user, body.center_id)
    _check_zone(conn, body.center_id, body.zone_id)
    cfg = active_config(conn, body.center_id)
    inputs = validate_inputs(body.inputs, cfg["config"])
    computed = engine.compute_run(inputs, cfg["config"])
    row = fetchone(
        conn,
        """INSERT INTO tool_runs (center_id, zone_id, user_id, config_id, engine_version, valid_date, inputs, computed,
               justifications, suggested_upper, suggested_middle, suggested_lower, suggested_overall, bottom_line, notes,
               client_mismatch, draft_ratings)
           VALUES (%(c)s, %(z)s, %(u)s, %(cfg)s, %(ev)s, %(d)s, %(i)s, %(comp)s, %(j)s, %(su)s, %(sm)s, %(sl)s, %(so)s,
               %(bl)s, %(n)s, %(mm)s, %(dr)s) RETURNING id""",
        {
            "c": body.center_id, "z": body.zone_id, "u": user.id, "cfg": cfg["id"], "ev": engine.ENGINE_VERSION,
            "d": body.valid_date, "i": Jsonb(inputs), "comp": Jsonb(computed), "j": Jsonb(body.justifications),
            "bl": body.bottom_line, "n": body.notes, "mm": _differs(body.client_computed, computed),
            "dr": _clean_ratings(body.draft_ratings), **_suggest_cols(computed),
        },
    )
    return get_run(row["id"], conn, user)


@router.get("/runs")
def list_runs(center_id: int, status: str | None = None, limit: int = 50, conn: Connection = Depends(get_conn), user: User = Depends(current_user)):
    require_center(user, center_id)
    return fetchall(
        conn,
        """SELECT r.id, r.valid_date, r.status, r.created_at, r.finalized_at, r.zone_id, z.name AS zone_name, u.name AS user_name,
                  r.suggested_upper, r.suggested_middle, r.suggested_lower, r.final_upper, r.final_middle, r.final_lower,
                  cc.version AS config_version,
                  (SELECT array_agg((p->>'type')::int) FROM jsonb_array_elements(r.inputs->'problems') p) AS problem_types
           FROM tool_runs r JOIN users u ON u.id = r.user_id JOIN center_configs cc ON cc.id = r.config_id
           LEFT JOIN zones z ON z.id = r.zone_id
           WHERE r.center_id = %s AND (%s::text IS NULL OR r.status = %s)
           ORDER BY r.valid_date DESC, r.id DESC LIMIT %s""",
        (center_id, status, status, min(max(limit, 1), 500)),
    )


@router.get("/runs/{run_id}")
def get_run(run_id: int, conn: Connection = Depends(get_conn), user: User = Depends(current_user)):
    run = _load_run(conn, run_id, user)
    run["readiness"] = engine.readiness(run["inputs"], run["computed"], run["justifications"])
    run["can_edit"] = run["status"] == "draft" and (run["user_id"] == user.id or user.role_in(run["center_id"]) == "center_admin")
    return run


@router.put("/runs/{run_id}")
def update_run(run_id: int, body: RunUpdate, conn: Connection = Depends(get_conn), user: User = Depends(current_user)):
    run = _load_run(conn, run_id, user)
    if run["status"] != "draft":
        raise HTTPException(409, "Finalized runs cannot be edited")
    if run["user_id"] != user.id and user.role_in(run["center_id"]) != "center_admin":
        raise HTTPException(403, "Only the author or a center admin can edit this run")
    fields = body.model_dump(exclude_unset=True)
    if "zone_id" in fields:
        _check_zone(conn, run["center_id"], fields["zone_id"])
    cfg = active_config(conn, run["center_id"])
    inputs = validate_inputs(fields.get("inputs", run["inputs"]), cfg["config"])
    computed = engine.compute_run(inputs, cfg["config"])
    conn.execute(
        """UPDATE tool_runs SET zone_id = %(z)s, valid_date = %(d)s, inputs = %(i)s, computed = %(comp)s, justifications = %(j)s,
               notes = %(n)s, bottom_line = %(bl)s, config_id = %(cfg)s, engine_version = %(ev)s,
               suggested_upper = %(su)s, suggested_middle = %(sm)s, suggested_lower = %(sl)s, suggested_overall = %(so)s,
               client_mismatch = %(mm)s, draft_ratings = %(dr)s, updated_at = now()
           WHERE id = %(id)s""",
        {
            "id": run_id, "z": fields.get("zone_id", run["zone_id"]), "d": fields.get("valid_date", run["valid_date"]),
            "i": Jsonb(inputs), "comp": Jsonb(computed), "j": Jsonb(fields.get("justifications", run["justifications"])),
            "n": fields.get("notes", run["notes"]), "bl": fields.get("bottom_line", run["bottom_line"]),
            "cfg": cfg["id"], "ev": engine.ENGINE_VERSION, "mm": _differs(fields.get("client_computed"), computed),
            "dr": _clean_ratings(fields["draft_ratings"]) if "draft_ratings" in fields else (
                Jsonb(run["draft_ratings"]) if run["draft_ratings"] else None),
            **_suggest_cols(computed),
        },
    )
    return get_run(run_id, conn, user)


@router.delete("/runs/{run_id}")
def delete_run(run_id: int, conn: Connection = Depends(get_conn), user: User = Depends(current_user)):
    run = _load_run(conn, run_id, user)
    if run["status"] != "draft":
        raise HTTPException(409, "Finalized runs are kept as history")
    if run["user_id"] != user.id and user.role_in(run["center_id"]) != "center_admin":
        raise HTTPException(403, "Only the author or a center admin can delete this run")
    conn.execute("DELETE FROM tool_runs WHERE id = %s", (run_id,))
    return {"ok": True}


@router.post("/runs/{run_id}/finalize")
def finalize_run(run_id: int, body: FinalizeIn, conn: Connection = Depends(get_conn), user: User = Depends(current_user)):
    run = _load_run(conn, run_id, user)
    if run["status"] != "draft":
        raise HTTPException(409, "Already finalized")
    if run["user_id"] != user.id and user.role_in(run["center_id"]) != "center_admin":
        raise HTTPException(403, "Only the author or a center admin can finalize this run")
    config = config_by_id(conn, run["config_id"])["config"]
    computed = engine.compute_run(run["inputs"], config)
    issues = engine.readiness(run["inputs"], computed, run["justifications"])
    if issues:
        raise HTTPException(422, {"message": "Finish the run before finalizing", "issues": issues})
    final = {"upper": body.upper, "middle": body.middle, "lower": body.lower}
    overall = body.overall if body.overall is not None else max(final.values())
    computed["finalChecks"] = engine.guidance_checks(run["inputs"], config, final)
    final_problems = [
        {
            "rank": i + 1,
            "type": p["type"],
            "locations": p["locations"],
            "likelihood": computed["problems"][i]["likelihood"],
            "size_min": p["size"]["min"],
            "size_max": p["size"]["max"],
            "discussion": p.get("discussion") or "",
        }
        for i, p in enumerate(run["inputs"]["problems"])
    ]
    conn.execute(
        """UPDATE tool_runs SET status = 'finalized', finalized_at = now(), updated_at = now(), computed = %s,
               final_upper = %s, final_middle = %s, final_lower = %s, final_overall = %s, final_problems = %s,
               bottom_line = coalesce(%s, bottom_line), notes = coalesce(%s, notes)
           WHERE id = %s""",
        (Jsonb(computed), body.upper, body.middle, body.lower, overall, Jsonb(final_problems), body.bottom_line, body.notes, run_id),
    )
    return get_run(run_id, conn, user)


# --- exports ---------------------------------------------------------------------

ABBR = dict(zip(engine.ASPECTS, ["N", "NE", "E", "SE", "S", "SW", "W", "NW"], strict=True))
WIZARD_BAND = {"upper": "above", "middle": "near", "lower": "below"}


def _fmt_size(v) -> str:
    if v is None:
        return "?"
    f = float(v)
    if f >= 4:
        return "D4–5"
    return f"D{f:g}"


def _locations_text(locations: list[str], band_labels: dict) -> str:
    by_band: dict[str, list[str]] = {}
    for loc in locations:
        aspect, band = loc.split()
        by_band.setdefault(band, []).append(ABBR[aspect])
    return "; ".join(f"{band_labels.get(b, b)}: {', '.join(by_band[b])}" for b in engine.BANDS if b in by_band)


def _problems_for_export(run: dict) -> list[dict]:
    if run.get("final_problems"):
        return run["final_problems"]
    return [
        {
            "rank": i + 1, "type": p["type"], "locations": p["locations"],
            "likelihood": run["computed"]["problems"][i]["likelihood"],
            "size_min": p["size"]["min"], "size_max": p["size"]["max"], "discussion": p.get("discussion") or "",
        }
        for i, p in enumerate(run["inputs"]["problems"])
    ]


@router.get("/runs/{run_id}/export")
def export_run(run_id: int, fmt: str = "text", conn: Connection = Depends(get_conn), user: User = Depends(current_user)):
    run = _load_run(conn, run_id, user)
    config = config_by_id(conn, run["config_id"])["config"]
    problems = _problems_for_export(run)
    band_labels = config.get("bandLabels") or {}
    danger_names = {d["level"]: d["name"] for d in config["dangerLevels"]}

    if fmt == "afp":
        return {
            "product_type": "forecast",
            "bottom_line": run.get("bottom_line") or "",
            "danger": [{"valid_day": "current", "upper": run["final_upper"], "middle": run["final_middle"], "lower": run["final_lower"]}],
            "forecast_avalanche_problems": [
                {
                    "avalanche_problem_id": p["type"],
                    "name": AFP_PROBLEM_NAMES.get(p["type"]),
                    "rank": p["rank"],
                    "location": p["locations"],
                    "likelihood": LIKELIHOOD_NAMES.get(p["likelihood"]),
                    "size": [f"{float(p['size_min']):.1f}", f"{float(p['size_max']):.1f}"] if p["size_min"] is not None else [],
                    "discussion": p.get("discussion") or "",
                }
                for p in problems
            ],
            "meta": {"tool_run_id": run["id"], "valid_date": run["valid_date"], "config_version": run["config_version"], "engine_version": run["engine_version"]},
        }

    if fmt == "backdrop":
        # Same payload shape the Backdrop module's prefill.js consumes.
        return {
            "problems": [
                {
                    "type": str(AFP_TO_DRUPAL_PROBLEM[p["type"]]),
                    "likelihood": str(p["likelihood"] or ""),
                    "sizeMin": str(max(1, min(5, math.floor(float(p["size_min"] or p["size_max"]))))),
                    "sizeMax": str(max(1, min(5, math.ceil(float(p["size_max"] or p["size_min"]))))),
                    "rose": [
                        f"{WIZARD_BAND[loc.split()[1]]}:{engine.ASPECTS.index(loc.split()[0])}" for loc in p["locations"]
                    ],
                }
                for p in problems
            ],
            "suggestions": {
                "above": danger_names.get(run["suggested_upper"]),
                "near": danger_names.get(run["suggested_middle"]),
                "below": danger_names.get(run["suggested_lower"]),
                "overall": danger_names.get(run["suggested_overall"]),
            },
            "notes": [v for v in (run.get("justifications") or {}).values() if v],
        }

    if fmt != "text":
        raise HTTPException(422, "fmt must be text, afp or backdrop")
    lines = [f"Avalanche forecast guidance — {run['center_name']}{' / ' + run['zone_name'] if run['zone_name'] else ''} — {run['valid_date']}"]
    if run["status"] == "finalized":
        lines.append(
            "Danger: "
            + " | ".join(f"{band_labels.get(b, b)}: {danger_names.get(run['final_' + b], '?')}" for b in engine.BANDS)
        )
    else:
        lines.append("Danger: not yet set (draft)")
    lines.append("")
    for p in problems:
        like = LIKELIHOOD_NAMES.get(p["likelihood"], "?")
        size = _fmt_size(p["size_min"]) if p["size_min"] == p["size_max"] else f"{_fmt_size(p['size_min'])}–{_fmt_size(p['size_max'])}"
        lines.append(f"Problem {p['rank']}: {AFP_PROBLEM_NAMES.get(p['type'])} — {like.capitalize()} — {size}")
        lines.append(f"  Where: {_locations_text(p['locations'], band_labels)}")
        if p.get("discussion"):
            lines.append(f"  {p['discussion'].strip()}")
    for key, text in (run.get("justifications") or {}).items():
        if text:
            lines.append(f"Pairing note ({key}): {text}")
    if run.get("bottom_line"):
        lines += ["", f"Bottom line: {run['bottom_line']}"]
    lines += ["", f"Tool config v{run['config_version']}, engine {run['engine_version']}. Guidance only — the forecaster sets the rating."]
    return PlainTextResponse("\n".join(lines))