"""Claude reviews of a draft forecast, informed by the current guidance and by
similar past forecasts. Claude advises; it never sets or dictates a rating.

Request layout (prompt-cache friendly — stable content first):
  system[0]  frozen reviewer instructions            (never changes)
  system[1]  guidance + center config for this run   (stable per config version) <- cache breakpoint
  messages   the draft, analogs, and the task         (varies every call)
"""

from __future__ import annotations

import json
from decimal import Decimal

import anthropic

from ..engine import ASPECTS
from ..mapping import AFP_PROBLEM_NAMES, LIKELIHOOD_NAMES
from ..sources import reference
from ..style import rules as style_rules

PROMPT_VERSION = "2026-09-18.1"
FALLBACK_BETA = "server-side-fallback-2026-07-01"
MAX_TOKENS = 32000

# USD per million tokens: input, output, cache write (5 min), cache read.
PRICES = {
    "claude-opus-5": (5.0, 25.0, 6.25, 0.50),
    "claude-fable-5-1": (10.0, 50.0, 12.5, 0.25),
    "claude-opus-4-8": (5.0, 25.0, 6.25, 0.50),
}

MODES = {
    "review_draft": "Review this draft for clarity and internal consistency: do the bottom line, danger ratings, avalanche problems and travel advice agree with each other and with the guidance? Flag jargon, passive voice, and missing what/where/what-to-do.",
    "draft_bottom_line": "Write a suggested Bottom Line for this draft (three sentences or fewer; what is it, where is it, what can readers do) from the structured problems and ratings, following the writing guidance. Also note anything in the draft that makes the bottom line hard to write.",
    "sanity_check": "Sanity-check the draft's danger ratings against its problems (likelihood, size, location), the CMAH hazard-chart suggestion, the on-the-line guidance, and the similar past forecasts. Where a rating looks high or low, explain why and what to double-check — never state which rating is correct.",
}

SYSTEM = """You are an experienced avalanche forecaster and editor reviewing a public avalanche forecast before it is published.

Ground rules:
- The forecaster decides. You advise. Never state what a danger rating must be, and never present your view as the correct rating.
- The current guidance supplied below is the standard, in this order of precedence: (1) the USDA FS Avalanche Forecast Guidance (FG), (2) the center's problem descriptions, (3) Statham et al. 2018, "A conceptual model of avalanche hazard" (CMAH), the research paper the FG builds on. Where they differ, the higher one wins; use CMAH only where the FG and center descriptions are silent, and cite it as "CMAH Table 4". Past forecasts are shown only as context: many were written under older conventions (for example "Normal Caution" or different problem choices) and may not follow current guidance — do not treat them as correct just because they were published.
- Writing style: apply the FG writing guidance (§3) first, then Writing for Busy Readers (Rogers & Lasky-Fink 2023), which is advisory. Any text you suggest must follow both: plain words, short sentences, no wordy phrases from the list below. Never swap out Forecast Guidance vocabulary (danger levels, likelihood, size, distribution, problem or terrain terms) for plainer words. Cite Busy Readers points as "Busy Readers".
- Be concise and practical. Cite the guidance section when it supports a point (e.g. "FG §3.3", "FG Table 1").
- Do not rewrite the whole forecast.
- Respond only with the JSON object described by the output schema."""

OUTPUT_SCHEMA = {
    "type": "object",
    "properties": {
        "summary": {"type": "string"},
        "issues": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "severity": {"type": "string", "enum": ["info", "suggestion", "important"]},
                    "section": {"type": "string", "enum": ["bottom_line", "danger", "problems", "travel_advice", "discussion", "general"]},
                    "text": {"type": "string"},
                    "guidance_ref": {"type": "string"},
                },
                "required": ["severity", "section", "text", "guidance_ref"],
                "additionalProperties": False,
            },
        },
        "suggested_bottom_line": {"type": "string"},
        "questions_for_forecaster": {"type": "array", "items": {"type": "string"}},
    },
    "required": ["summary", "issues", "suggested_bottom_line", "questions_for_forecaster"],
    "additionalProperties": False,
}


class AIUnavailable(Exception):
    pass


def enabled(api_key: str | None) -> bool:
    return bool(api_key)


def guidance_context(config: dict, center_name: str) -> str:
    """Stable per config version: the guidance the review applies."""
    lines = [f"# Guidance for {center_name}", "", "## Avalanche problem types (FG Table 1; center descriptions)"]
    for p in config["problemTypes"]:
        lines.append(f"### {p['name']}")
        lines.append(p["desc"])
        for label, key in (("Initiation", "initiation"), ("Termination/transition", "termination"), ("Considerations", "considerations")):
            if p.get(key):
                lines.append(f"{label}: " + " ".join(p[key]))
        if p.get("sac"):
            lines.append(f"Center notes: {p['sac']}")
    lines += ["", "## Danger levels (NAPADS, FG §2.1)"]
    for d in config["dangerLevels"]:
        if d["level"]:
            lines.append(f"- {d['level']} {d['name']}: {d['advice']} Likelihood: {d.get('likelihood', '')} Size/distribution: {d.get('sizeDist', '')}")
    lines += ["", "## On-the-line and special situations (FG §2.2)"]
    lines += [f"- {t['title']}: {t['text']}" for t in config["onTheLine"]]
    lines += ["", "## Writing guidance (FG §3)"]
    for section, items in config["writing"].items():
        lines.append(f"{section}: " + " ".join(items))
    lines += ["", "## Background: CMAH Table 4 (Statham et al. 2018) — lower precedence than the FG and center descriptions",
              "Relative size is on the R scale (relative to the path), not the D scale."]
    names = {str(p["id"]): p["name"] for p in config["problemTypes"]}
    for pid, note in sorted(reference()["cmah_problems"].items(), key=lambda kv: int(kv[0])):
        if pid in names:
            lines.append(f"- {names[pid]}: persistence: {note['persistence']} Typical relative size {note['relative_size']}. "
                         f"Typical risk mitigation: {note['mitigation']}")
    style = style_rules()
    lines += ["", "## Writing for Busy Readers (Rogers & Lasky-Fink 2023) — after the FG writing guidance"]
    lines += [f"- {p['title']}: " + "; ".join(p["items"]) for p in style["principles"]]
    common = sorted((w for w in style["wordy"] if w["archive_count"]), key=lambda w: (-w["archive_count"], w["phrase"]))
    lines.append("Wordy phrases to avoid (→ plainer alternatives), most common in past forecasts first: "
                 + "; ".join(f"{w['phrase']} → {' / '.join(w['options'])}" for w in common))
    lines += ["", "## CMAH hazard-chart danger table used by the tool (a suggestion, not a rule)",
              "Rows likelihood 1-5 (Unlikely..Certain); columns D1, D2, D3, D4-5; values 0-5 (No Rating..Extreme):",
              json.dumps(config["dangerTable"], sort_keys=True),
              f"Band labels: {json.dumps(config.get('bandLabels', {}), sort_keys=True)}"]
    return "\n".join(lines)


def _loc_summary(locations: list[str]) -> str:
    by_band: dict[str, list[str]] = {}
    for loc in locations:
        aspect, band = loc.split()
        by_band.setdefault(band, []).append(aspect)
    order = {a: i for i, a in enumerate(ASPECTS)}
    return "; ".join(f"{b}: {', '.join(sorted(v, key=order.get))}" for b, v in by_band.items())


def draft_payload(run: dict, config: dict) -> dict:
    names = {d["level"]: d["name"] for d in config["dangerLevels"]}
    problems = []
    for i, p in enumerate(run["inputs"]["problems"]):
        c = run["computed"]["problems"][i]
        problems.append({
            "rank": i + 1,
            "type": AFP_PROBLEM_NAMES.get(p["type"]),
            "likelihood": LIKELIHOOD_NAMES.get(c["likelihood"]),
            "size_range": [p["size"]["min"], p["size"]["max"]],
            "where": _loc_summary(p["locations"]),
        })
    final = None
    if run.get("final_upper") is not None:
        final = {b: names.get(run[f"final_{b}"]) for b in ("upper", "middle", "lower")}
    return {
        "valid_date": str(run["valid_date"]),
        "zone": run.get("zone_name"),
        "observed_conditions": run["inputs"].get("conditions", []),
        "problems": problems,
        "forecaster_danger": final or "not yet set",
        "tool_suggestion": {b: names.get(run["computed"]["suggestions"][b]) for b in ("upper", "middle", "lower")},
        "pairing_justifications": run.get("justifications") or {},
        "guidance_checks": [c["text"] for c in run["computed"].get("finalChecks", run["computed"].get("checks", []))],
        "bottom_line": run.get("bottom_line") or "",
        "notes": run.get("notes") or "",
    }


def analog_payload(analogs: list[dict], config: dict) -> list[dict]:
    names = {d["level"]: d["name"] for d in config["dangerLevels"]}
    return [
        {
            "date": str(a["valid_date"]),
            "era": a["guidance_era"],
            "danger": {b: names.get(a[f"danger_{b}"]) for b in ("upper", "middle", "lower")},
            "problems": [
                {"type": AFP_PROBLEM_NAMES.get(p["type"]), "likelihood": LIKELIHOOD_NAMES.get(p["likelihood"]),
                 "size_max": p["size_max"]}
                for p in a["problems"]
            ],
            "bottom_line": (a.get("bottom_line") or "")[:400],
        }
        for a in analogs
    ]


def build_request(run: dict, config: dict, center_name: str, analogs: list[dict], mode: str, model: str) -> dict:
    if mode not in MODES:
        raise ValueError(f"unknown mode {mode}")
    user = (
        f"Task: {MODES[mode]}\n\n"
        f"<draft_forecast>\n{json.dumps(draft_payload(run, config), indent=1, default=str)}\n</draft_forecast>\n\n"
        f"<similar_past_forecasts note=\"context only; legacy-era forecasts may not follow current guidance\">\n"
        f"{json.dumps(analog_payload(analogs, config), indent=1, default=str)}\n</similar_past_forecasts>"
    )
    params = {
        "model": model,
        "max_tokens": MAX_TOKENS,
        "betas": [FALLBACK_BETA],
        "fallbacks": "default",
        "system": [
            {"type": "text", "text": SYSTEM},
            {"type": "text", "text": guidance_context(config, center_name), "cache_control": {"type": "ephemeral"}},
        ],
        "messages": [{"role": "user", "content": user}],
        "output_config": {"format": {"type": "json_schema", "schema": OUTPUT_SCHEMA}},
    }
    # Claude Fable 5.1 thinks by default and rejects explicit thinking config other than adaptive;
    # Opus 5 also defaults to adaptive — set it explicitly there.
    if model == "claude-opus-5":
        params["thinking"] = {"type": "adaptive"}
    return params


def cost_usd(model: str, usage) -> Decimal:
    p_in, p_out, p_cw, p_cr = PRICES.get(model, PRICES["claude-opus-5"])
    total = (
        (getattr(usage, "input_tokens", 0) or 0) * p_in
        + (getattr(usage, "output_tokens", 0) or 0) * p_out
        + (getattr(usage, "cache_creation_input_tokens", 0) or 0) * p_cw
        + (getattr(usage, "cache_read_input_tokens", 0) or 0) * p_cr
    ) / 1_000_000
    return Decimal(str(round(total, 5)))


def run_review(params: dict, api_key: str | None) -> dict:
    """Call Claude (streamed server-side) and return a result dict for storage."""
    if not api_key:
        raise AIUnavailable("AI review is not configured: a site admin can add the API key on the Admin page.")
    client = anthropic.Anthropic(api_key=api_key, timeout=600.0, max_retries=2)
    with client.beta.messages.stream(**params) as stream:
        message = stream.get_final_message()

    served_by = message.model
    fallback_model = served_by if served_by != params["model"] else None
    usage = message.usage
    result = {
        "stop_reason": message.stop_reason,
        "fallback_model": fallback_model,
        "input_tokens": usage.input_tokens,
        "output_tokens": usage.output_tokens,
        "cache_read_tokens": getattr(usage, "cache_read_input_tokens", None),
        "cache_write_tokens": getattr(usage, "cache_creation_input_tokens", None),
        "cost_usd": cost_usd(served_by, usage),
        "output": None,
        "error": None,
    }
    if message.stop_reason == "refusal":
        result["error"] = "The model declined this request."
        return result
    text = next((b.text for b in message.content if b.type == "text"), "")
    try:
        result["output"] = json.loads(text)
    except json.JSONDecodeError:
        result["error"] = f"Unreadable response (stop_reason={message.stop_reason})."
        result["output"] = {"raw": text[:4000]}
    return result