import copy
import hashlib
import json
from pathlib import Path
import re

import pytest

from hessian_recipes.agent_rl_evaluation import compare, run, verify_report
from hessian_recipes.agent_rl_tasks import generate_tasks, parse_response
from hessian_recipes.evaluation import checkpoint_sha256, sha256_file
from hessian_recipes.trajectories import read_trajectories


class FixturePolicy:
    template_metadata = {
        "chat_template_sha256": hashlib.sha256(b"explicit-byte-fixture").hexdigest()
    }
    tool_schemas = []
    runtime = {"mode": "fixture"}
    packages = {}

    def __init__(self, mode="success"):
        self.mode = mode

    def encode(self, messages):
        text = messages[0]["content"] + "\n<assistant>"
        for message in messages[1:]:
            if message["role"] == "assistant":
                text += message["raw_response"]
            else:
                text += "\n<tool>" + message["content"] + "\n<assistant>"
        return list(text.encode())

    def __call__(self, messages, *, seed, max_new_tokens):
        if self.mode == "infrastructure":
            raise TimeoutError("fixture infrastructure unavailable")
        last = messages[-1]
        if last["role"] == "user":
            record = re.search(r"beleg-[0-9a-f]+", last["content"]).group()
            if self.mode == "skip":
                raw = "<answer>6</answer>"
            elif self.mode == "unknown":
                raw = '<tool_call>{"name":"shell","arguments":{"command":"touch forbidden"}}</tool_call>'
            elif self.mode == "same_turn":
                raw = (
                    "<tool_call>"
                    + json.dumps({"name": "lookup", "arguments": {"record_id": record}})
                    + "</tool_call>"
                )
                raw += (
                    '<tool_call>{"name":"calculator","arguments":{"expression":"2*3"}}</tool_call>'
                )
            else:
                raw = (
                    "<tool_call>"
                    + json.dumps({"name": "lookup", "arguments": {"record_id": record}})
                    + "</tool_call>"
                )
        else:
            observation = json.loads(last["content"])
            if "quantity" in observation:
                expression = f"{observation['quantity']}*{observation['unit_price']}"
                raw = (
                    "<tool_call>"
                    + json.dumps({"name": "calculator", "arguments": {"expression": expression}})
                    + "</tool_call>"
                )
            else:
                raw = "<answer>" + observation.get("result", "1") + "</answer>"
        tokens = list(raw.encode())[:max_new_tokens]
        text = bytes(tokens).decode()
        assistant = {"role": "assistant", "content": text, "raw_response": text}
        try:
            actions = parse_response(text)
            if all("tool" in action for action in actions):
                assistant["content"] = ""
                assistant["tool_calls"] = [
                    {
                        "type": "function",
                        "function": {
                            "name": action["tool"],
                            "arguments": {k: v for k, v in action.items() if k != "tool"},
                        },
                    }
                    for action in actions
                ]
        except (ValueError, TypeError):
            if self.mode == "unknown" and last["role"] == "user" and len(tokens) < max_new_tokens:
                assistant["tool_calls"] = [
                    {
                        "type": "function",
                        "function": {"name": "shell", "arguments": {"command": "touch forbidden"}},
                    }
                ]
                assistant["content"] = ""
        return {
            "assistant": assistant,
            "prompt_token_ids": self.encode(messages),
            "response_token_ids": tokens,
            "finish_reason": "length" if len(tokens) == max_new_tokens else "stop",
        }


def setup(tmp_path, *, budget=768, iterations=3, count=3):
    model = tmp_path / "checkpoint"
    model.mkdir()
    (model / "weights.bin").write_bytes(b"fixture only; no actual policy")
    tokenizer = tmp_path / "tokenizer"
    tokenizer.mkdir()
    (tokenizer / "fixture.txt").write_text("UTF-8 bytes")
    tasks = generate_tasks("test", count)
    for task in tasks:
        task["task"]["private_facts"] = {"quantity": 2, "unit_price": 3}
    path = tmp_path / "tasks.jsonl"
    path.write_text("".join(json.dumps(row) + "\n" for row in tasks))
    spec = {
        "schema_version": 1,
        "model": {"path": "checkpoint", "sha256": checkpoint_sha256(model)},
        "tokenizer": {"path": "tokenizer", "sha256": checkpoint_sha256(tokenizer)},
        "tasks": {"path": "tasks.jsonl", "sha256": sha256_file(path)},
        "sampling": {"temperature": 0.6, "top_p": 0.95, "top_k": 20, "min_p": 0, "seed": 42},
        "max_completion_length": budget,
        "max_tool_calling_iterations": iterations,
        "chat_template_kwargs": {"enable_thinking": True},
        "trials": 1,
        "device": "cpu",
        "dtype": "float32",
    }
    spec_path = tmp_path / "spec.json"
    spec_path.write_text(json.dumps(spec))
    return spec_path


def test_actual_multiturn_observations_drive_actions_and_export_captured_masks(tmp_path):
    spec = setup(tmp_path)
    report = run(spec, tmp_path / "run", generator=FixturePolicy())
    assert report["status"] == "complete"
    assert report["metrics"]["success_rate"] == report["metrics"]["workflow_rate"] == 1
    assert report["metrics"]["tool_calls"] == 6
    assert report["metrics"]["tool_errors"] == 0
    assert report["metrics"]["malformed_final_outputs"] == 0
    assert report["source_kind"] == "fixture"
    assert report["language_audit"]["status"] == "unassessed"
    verify_report(report)
    trajectories = read_trajectories(tmp_path / "run/trajectories.jsonl")
    assert len(trajectories) == 3
    assert all(len(episode["calls"]) == 3 for episode in trajectories)
    assert all(
        call["response_loss_mask"] == [1] * len(call["response_token_ids"])
        for episode in trajectories
        for call in episode["calls"]
    )
    assert trajectories[0]["calls"][1]["context"][-1]["role"] == "tool"
    with pytest.raises(ValueError, match="fresh"):
        run(spec, tmp_path / "run", generator=FixturePolicy())


@pytest.mark.parametrize("mode,errors", [("same_turn", 0), ("skip", 0), ("unknown", 1)])
def test_correct_final_without_legal_workflow_never_earns_reward(tmp_path, mode, errors):
    spec = setup(tmp_path, count=1)
    report = run(spec, tmp_path / "run", generator=FixturePolicy(mode))
    assert report["metrics"]["success_rate"] == 0
    assert report["metrics"]["mean_reward"] == 0
    assert report["metrics"]["tool_errors"] == errors
    verify_report(report)
    assert not Path("forbidden").exists()


def test_retained_stream_and_tool_iteration_caps_are_visible(tmp_path):
    spec = setup(tmp_path, budget=64, count=1)
    report = run(spec, tmp_path / "cap", generator=FixturePolicy())
    assert report["metrics"]["token_cap_episodes"] == 1
    assert report["episodes"][0]["completion_stream_tokens"] <= 64
    verify_report(report)
    value = json.loads(spec.read_text())
    value.update(max_completion_length=768, max_tool_calling_iterations=1)
    spec.write_text(json.dumps(value))
    report = run(spec, tmp_path / "turn", generator=FixturePolicy())
    assert report["metrics"]["turn_cap_episodes"] == 1
    assert report["metrics"]["success_rate"] == 0
    verify_report(report)


def test_comparison_pairing_and_failed_coverage(tmp_path):
    spec = setup(tmp_path)
    run(spec, tmp_path / "baseline", generator=FixturePolicy("skip"))
    run(spec, tmp_path / "candidate", generator=FixturePolicy())
    gate = compare(tmp_path / "baseline/report.json", tmp_path / "candidate/report.json")
    assert gate["passed"] and gate["release_eligible"] is False
    assert gate["paired"]["wins"] == 3 and gate["paired"]["losses"] == 0
    assert gate["paired"]["descriptive_bootstrap_95"] == [1, 1]
    report = run(spec, tmp_path / "failed", generator=FixturePolicy("infrastructure"))
    assert report["status"] == "failed" and report["metrics"]["success_rate"] is None
    assert not compare(tmp_path / "baseline/report.json", tmp_path / "failed/report.json")["passed"]


def test_replay_rejects_changed_observation_outcome_and_hidden_prompt_facts(tmp_path):
    spec = setup(tmp_path, count=1)
    report = run(spec, tmp_path / "run", generator=FixturePolicy())
    bad = copy.deepcopy(report)
    bad["episodes"][0]["outcome"]["reward"] = 0
    with pytest.raises(ValueError, match="re-scored"):
        verify_report(bad)
    bad = copy.deepcopy(report)
    bad["episodes"][0]["calls"][0]["observations"][0]["content"] = "fabricated result"
    with pytest.raises(ValueError, match="modified"):
        verify_report(bad)
    bad = copy.deepcopy(report)
    bad["episodes"][0]["calls"][0]["context"][0]["content"] += " hidden quantity=2"
    with pytest.raises(ValueError, match="modified"):
        verify_report(bad)


def test_comparison_rejects_changed_budget_even_with_same_scores(tmp_path):
    spec = setup(tmp_path, count=1)
    run(spec, tmp_path / "baseline", generator=FixturePolicy())
    value = json.loads(spec.read_text())
    value["max_completion_length"] = 700
    spec.write_text(json.dumps(value))
    run(spec, tmp_path / "candidate", generator=FixturePolicy())
    gate = compare(tmp_path / "baseline/report.json", tmp_path / "candidate/report.json")
    assert not gate["passed"]
    assert "incompatible max_completion_length" in gate["reasons"]