from dataclasses import replace

import pytest
from conftest import GOOD_SCENE

from api import budget, routing
from api.llm import ScriptedLLM
from api.logstore import LogStore
from api.simulate import generate_scene


def fenced(xml):
    return f"```xml\n{xml}\n```"


def test_code_complexity_picks_the_tier(settings):
    simple = routing.plan_code("Publish hello on /chatter every second", "humble", "python", settings)
    medium = routing.plan_code(
        "A service client that reads a parameter and publishes odometry", "jazzy", "python", settings
    )
    complex_ = routing.plan_code(
        "A Nav2 behavior tree node that calls MoveIt for pick and place", "jazzy", "cpp", settings
    )
    assert (simple.tier, medium.tier, complex_.tier) == ("low", "mid", "high")
    assert simple.model == "deepseek/deepseek-v4-flash" and complex_.model == "z-ai/glm-5.2"
    assert simple.category == "coding" and "navigation" in complex_.reason


def test_scene_plan_and_simulator_choice(settings):
    p = routing.plan_scene("A box on the floor", "unitree_g1", settings)
    assert p.tier == "mid" and p.simulator == "mujoco"  # scenes never use the cheapest tier
    assert routing.plan_scene("A door that opens on a hinge", None, settings).tier == "high"
    assert routing.plan_scene("How often does it fall on random slopes?", None, settings).simulator == "mujoco_warp"
    assert routing.plan_scene("A photoreal kitchen for camera images", None, settings).simulator == "isaac_sim"
    assert not routing.SIMULATORS["gazebo"]["available"]


def test_repairs_escalate_to_a_stronger_model(store, settings):
    bad = GOOD_SCENE.replace('mass="3"', 'mass="3" frictoin="1"')
    llm = ScriptedLLM([fenced(bad), fenced(GOOD_SCENE)])
    plan = routing.plan_scene("G1 next to a crate", "unitree_g1", settings)
    result = generate_scene("G1 next to a crate", "unitree_g1", llm, store, settings, plan=plan)
    assert result.valid and llm.models == ["z-ai/glm-5.2", "moonshotai/kimi-k2.6"]  # scene models
    assert [a.model for a in result.attempts] == llm.models and result.model == "moonshotai/kimi-k2.6"


def test_estimates_grow_with_tier_and_backend(settings):
    oc = replace(settings, codegen_backend="opencode")
    low = routing.plan_code("Publish hello", "humble", "python", oc).estimate_inr
    high = routing.plan_code("Nav2 planner with MoveIt and SLAM", "humble", "python", oc).estimate_inr
    assert 0 < low[0] <= low[1] and high[0] > low[0]
    direct = routing.plan_code("Publish hello", "humble", "python", replace(settings, codegen_backend="direct"))
    assert direct.estimate_inr[0] < low[0]  # OpenCode's agent loop reads more context


# --------------------------------------------------------------------------- rupee budget


def spend(logs, settings, tester, usd, cost=None, model="deepseek/deepseek-v4-pro", category="coding"):
    spec = routing.model_for_tier("mid")
    usage = {"prompt_tokens": 1000, "completion_tokens": 100}
    if cost is not None:
        usage["cost"] = cost
    elif usd:
        usage["cost"] = usd
    budget.record_call(logs, settings, tester=tester, generation_id="g", category=category, model=spec, usage=usage)


def test_budget_counts_billable_calls_only(tmp_path, settings):
    logs = LogStore(tmp_path / "l.db")
    spend(logs, settings, "alice", 0.5)  # OpenRouter reported $0.50
    spend(logs, settings, "alice", 0, cost=0)  # OpenRouter-exempt: reported no charge
    spend(logs, settings, "bob", 0.25, category="simulation_ai")
    s = budget.summary(logs, settings, "alice")
    assert s["used"] == pytest.approx(0.5 * settings.usd_inr) and s["exempt_calls"] == 1
    assert s["by_category"]["coding"] == s["used"] and s["remaining"] == pytest.approx(125 - s["used"])
    assert s["pool"]["used"] == pytest.approx(0.75 * settings.usd_inr) and s["pool"]["total"] == 500
    free = routing.ModelSpec("some/model:free", "low", "x", 0, 0)
    assert budget.cost_of_call(free, {"prompt_tokens": 10}).exempt_reason.startswith("free model")
    listed = budget.cost_of_call(routing.model_for_tier("low"), {"prompt_tokens": 1_000_000, "completion_tokens": 0})
    assert listed.estimated and listed.cost_usd == pytest.approx(0.089)


def test_budget_check_refuses_and_asks(tmp_path, settings):
    logs = LogStore(tmp_path / "l.db")
    budget.check(logs, settings, "alice", (1.0, 3.0), confirmed=False)  # cheap: no question
    with pytest.raises(budget.BudgetError) as e:
        budget.check(logs, settings, "alice", (4.0, 12.0), confirmed=False)
    assert e.value.status == 409 and e.value.detail["needs_confirmation"]
    budget.check(logs, settings, "alice", (4.0, 12.0), confirmed=True)
    spend(logs, settings, "alice", 124.5 / settings.usd_inr)
    with pytest.raises(budget.BudgetError) as e:
        budget.check(logs, settings, "alice", (1.0, 2.0), confirmed=True)
    assert e.value.status == 402 and "Your AI budget" in str(e.value)
    for name in ("bob", "carol", "dave"):
        spend(logs, settings, name, 125 / settings.usd_inr)
    with pytest.raises(budget.BudgetError) as e:
        budget.check(logs, settings, "erin", (1.0, 2.0), confirmed=True)
    assert "team's AI budget" in str(e.value)


# --------------------------------------------------------------------------- through the API


def test_generate_reports_routing_and_cost(run_env):
    client, keys_, backend, services, gen_id = run_env
    h = {"X-API-Key": keys_["alice"]}
    services.llm_factory().responses[:] = [fenced(GOOD_SCENE)]
    r = client.post("/simulate/generate", json={"description": "G1 near a crate and a ramp"}, headers=h).json()
    assert r["routing"]["tier"] == "mid" and r["routing"]["simulator"] == "mujoco" and not r["routing"]["escalated"]
    assert r["cost"]["inr"] >= 0 and r["cost"]["budget"]["allocation"] == 125
    usage = client.get("/me/usage", headers=h).json()
    assert usage["ai"]["calls"] >= 2 and usage["ai"]["pool"]["total"] == 500 and not usage["admin"]
    g = services.logs.get(r["id"])
    assert g["request"]["routing"]["model"] == "z-ai/glm-5.2"


def test_expensive_task_needs_confirmation(run_env):
    client, keys_, backend, services, gen_id = run_env
    services.settings = replace(services.settings, confirm_above_inr=0.0001)
    h = {"X-API-Key": keys_["alice"]}
    r = client.post("/simulate/generate", json={"description": "G1 near a crate and a ramp"}, headers=h)
    assert r.status_code == 409 and r.json()["detail"]["needs_confirmation"]
    services.llm_factory().responses[:] = [fenced(GOOD_SCENE)]
    body = {"description": "G1 near a crate and a ramp", "confirm_cost": True}
    assert client.post("/simulate/generate", json=body, headers=h).status_code == 200


def test_feedback_is_stored_with_context(run_env):
    client, keys_, backend, services, gen_id = run_env
    h = {"X-API-Key": keys_["alice"]}
    body = {"target_type": "generation", "target_id": gen_id, "rating": "not_helpful", "note": "box is floating"}
    assert client.post("/feedback", json=body, headers=h).json() == {"ok": True}
    (fb,) = services.logs.feedback_rows()
    assert fb["rating"] == "not_helpful" and fb["note"] == "box is floating" and fb["context"]["kind"] == "simulate"
    assert services.logs.get(gen_id)["tester_rating"] == "broken"
    other = {"X-API-Key": keys_["bob"]}
    assert client.post("/feedback", json=body, headers=other).status_code == 404


def test_admin_sees_the_team(run_env):
    client, keys_, backend, services, gen_id = run_env
    services.settings = replace(services.settings, admin_testers=("alice",))
    team = client.get("/me/usage", headers={"X-API-Key": keys_["alice"]}).json()["team"]
    assert {t["tester"] for t in team} >= {"alice", "bob"} and all(t["allocation"] == 125 for t in team)


# --------------------------------------------------------------------------- free debugger


class FlakyFree:
    """Stands in for ZenLLM: answers from a script, or fails like a rate-limited free model."""

    model = "opencode/deepseek-v4-flash-free"

    def __init__(self, responses, fail=False):
        self.responses, self.fail, self.calls = list(responses), fail, 0

    def complete(self, messages, temperature=0.2, model=None):
        from api.llm import Completion, LLMError

        self.calls += 1
        if self.fail:
            raise LLMError("free model: 429 rate limited")
        return Completion(self.responses.pop(0), model, 0, {"input_tokens": 900, "output_tokens": 400, "cost": 0.0})


def test_first_repair_goes_to_the_free_debugger(store, settings):
    from api.zen import WithFreeModels

    s = replace(settings, free_debug=True)
    bad = GOOD_SCENE.replace('mass="3"', 'mass="3" frictoin="1"')
    paid, free = ScriptedLLM([fenced(bad)]), FlakyFree([fenced(GOOD_SCENE)])
    plan = routing.plan_scene("G1 next to a crate", "unitree_g1", s)
    assert plan.debugger == "opencode/deepseek-v4-flash-free"
    result = generate_scene("G1 next to a crate", "unitree_g1", WithFreeModels(paid, free), store, s, plan=plan)
    assert result.valid and paid.models == ["z-ai/glm-5.2"] and free.calls == 1
    assert [a.model for a in result.attempts] == ["z-ai/glm-5.2", "opencode/deepseek-v4-flash-free"]


def test_free_debugger_failure_falls_back_to_a_paid_tier(store, settings):
    from api.zen import WithFreeModels

    s = replace(settings, free_debug=True)
    bad = GOOD_SCENE.replace('mass="3"', 'mass="3" frictoin="1"')
    paid, free = ScriptedLLM([fenced(bad), fenced(GOOD_SCENE)]), FlakyFree([], fail=True)
    plan = routing.plan_scene("G1 next to a crate", "unitree_g1", s)
    result = generate_scene("G1 next to a crate", "unitree_g1", WithFreeModels(paid, free), store, s, plan=plan)
    assert result.valid and free.calls == 1  # tried, failed, and the same attempt went to the paid model
    assert paid.models == ["z-ai/glm-5.2", "moonshotai/kimi-k2.6"] and len(result.attempts) == 2


def test_attempt_schedule_and_worst_case(settings):
    s = replace(settings, free_debug=True)
    plan = routing.plan_code("Publish hello on /chatter", "humble", "python", s)
    ids = [[m.id for m in routing.candidates_for_attempt(plan, i)] for i in range(3)]
    assert ids == [
        ["deepseek/deepseek-v4-flash"],
        ["opencode/deepseek-v4-flash-free", "deepseek/deepseek-v4-pro"],
        ["z-ai/glm-5.2"],
    ]
    # the estimate assumes the free call fails, so its paid fallback is in the worst case
    tin, tout = routing.TOKENS[("code", "direct")]
    paid = [routing.model_for_attempt(plan, i) for i in range(3)]
    worst = sum(routing.call_cost_usd(m, tin * (1 + i // 2), tout) for i, m in enumerate(paid)) * s.usd_inr
    assert plan.estimate_inr[1] == pytest.approx(worst, abs=0.01)


def test_free_calls_are_recorded_but_never_billed(tmp_path, settings):
    free = routing.ModelSpec("opencode/deepseek-v4-flash-free", "free", "OpenCode Zen", 0, 0, free=True)
    logs = LogStore(tmp_path / "l.sqlite3")
    budget.record_call(logs, settings, tester="a", generation_id="g", category="debugging", model=free,
                       usage={"input_tokens": 5000, "output_tokens": 900})  # fmt: skip
    (row,) = logs.usage_rows()
    assert row["billable"] == 0 and row["provider"] == "opencode-zen" and row["exempt_reason"] == "opencode_zen_free"
    assert budget.summary(logs, settings, "a")["used"] == 0