from __future__ import annotations

import copy
import uuid
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from .conformance import evaluate_proposal_conformance
from .cost import CostGuard
from .events import EventWriter, RunEvent, sha256_json
from .models import ToolCall
from .oracles import OracleResult, evaluate_oracle
from .policy import PolicyEngine, PolicySet
from .providers.base import ProviderAdapter
from .sandbox import SyntheticSandbox
from .scenarios import Scenario
from .safety import SafetyViolation, contains_credential_like_value


@dataclass(frozen=True)
class RunOutcome:
    run_id: str
    scenario_id: str
    provider: str
    model: str
    decision: str
    rule_id: str
    oracle_passed: bool
    containment_failure: bool
    state_changed: bool
    cost_usd: float
    latency_seconds: float = 0.0
    trial_id: str = ""
    repetition: int = 0
    schedule_position: int = 0
    aborted: bool = False
    abort_class: str = ""
    terminal: bool = False
    error: str = ""


class ExperimentRunner:
    def __init__(
        self,
        policy: PolicySet,
        sandbox: SyntheticSandbox,
        adapter: ProviderAdapter,
        event_writer: EventWriter,
        cost_guard: CostGuard,
    ) -> None:
        self.policy = policy
        self.engine = PolicyEngine(policy)
        self.sandbox = sandbox
        self.adapter = adapter
        self.events = event_writer
        self.cost_guard = cost_guard

    def _emit(self, run_id: str, event_type: str, payload: dict[str, Any]) -> None:
        self.events.append(RunEvent.create(run_id, event_type, payload))

    def run(
        self, scenario: Scenario, metadata: dict[str, Any] | None = None
    ) -> RunOutcome:
        metadata = dict(metadata or {})
        started_at = time.perf_counter()
        run_id = f"{scenario.scenario_id}-{uuid.uuid4().hex[:12]}"
        actual_cost = 0.0
        reset_id = self.sandbox.reset()
        self._emit(
            run_id,
            "run_started",
            {
                "scenario_id": scenario.scenario_id,
                "scenario_hash": sha256_json(
                    {
                        "task": scenario.task,
                        "untrusted_content": scenario.untrusted_content,
                        "expected_decision": scenario.expected_decision,
                    }
                ),
                "provider": self.adapter.provider,
                "model": self.adapter.model,
                "policy_id": self.policy.policy_id,
                "policy_version": self.policy.version,
                "mutation_id": self.policy.mutation_id,
                "reset_id": reset_id,
                "trial": metadata,
            },
        )
        try:
            self.cost_guard.preflight(
                self.adapter.provider,
                input_tokens=1800 if self.adapter.provider != "mock" else 0,
                output_tokens=500 if self.adapter.provider != "mock" else 0,
            )
            self._emit(
                run_id,
                "model_request",
                {
                    "task_hash": sha256_json(scenario.task),
                    "provider": self.adapter.provider,
                    "model": self.adapter.model,
                },
            )
            response = self.adapter.propose(scenario)
            actual_cost = self.cost_guard.record(
                self.adapter.provider, response.input_tokens, response.output_tokens
            )
            self._emit(
                run_id,
                "model_response",
                {
                    "provider": response.provider,
                    "model": response.model,
                    "proposal_count": len(response.proposals),
                    "input_tokens": response.input_tokens,
                    "output_tokens": response.output_tokens,
                    "stop_reason": response.raw_stop_reason,
                },
            )
            if len(response.proposals) != 1:
                raise SafetyViolation("Exactly one synthetic proposal is required")
            proposal = response.proposals[0]
            if proposal.name != "propose_action":
                raise SafetyViolation("Unexpected provider tool name")
            if contains_credential_like_value(proposal.arguments):
                raise SafetyViolation("Credential-like value detected in proposal")
            call = ToolCall(
                subject=str(proposal.arguments["subject"]),
                action=str(proposal.arguments["action"]),
                resource=str(proposal.arguments["resource"]),
                arguments=dict(proposal.arguments.get("arguments", {})),
                context=dict(proposal.arguments.get("context", {})),
                call_id=proposal.call_id,
            )
            self._emit(run_id, "tool_proposed", {"call": call.as_dict()})
            conformance = evaluate_proposal_conformance(scenario, call)
            self._emit(run_id, "proposal_conformance", conformance.as_dict())
            if not conformance.passed:
                raise SafetyViolation(
                    "Proposal does not exactly represent the frozen synthetic operation"
                )
            call.context.setdefault("provenance", scenario.provenance)
            before = self.sandbox.state.digest()
            decision = self.engine.evaluate(call)
            self._emit(
                run_id,
                "policy_decision",
                {
                    "effect": decision.effect,
                    "rule_id": decision.rule_id,
                    "reason": decision.reason,
                    "evidence": decision.evidence,
                },
            )
            tool_result: dict[str, Any] = {"executed": False}
            expected_state_after = before
            if decision.allowed:
                replay = SyntheticSandbox(self.sandbox.manifest)
                replay.state = copy.deepcopy(self.sandbox.state)
                replay.execute(call)
                expected_state_after = replay.state.digest()
                result = self.sandbox.execute(call)
                tool_result = {"executed": True, "result": result}
            self._emit(run_id, "tool_result", tool_result)
            after = self.sandbox.state.digest()
            oracle = evaluate_oracle(
                scenario,
                decision,
                before,
                after,
                expected_state_after=expected_state_after,
            )
            self._emit(run_id, "oracle_result", asdict(oracle))
            self._emit(
                run_id,
                "cost_update",
                {
                    "actual_cost_usd": actual_cost,
                    "provider_spend_usd": self.cost_guard.provider_spend.get(
                        self.adapter.provider, 0.0
                    ),
                    "total_spend_usd": self.cost_guard.total_spend,
                },
            )
            outcome = RunOutcome(
                run_id=run_id,
                scenario_id=scenario.scenario_id,
                provider=response.provider,
                model=response.model,
                decision=decision.effect,
                rule_id=decision.rule_id,
                oracle_passed=oracle.passed,
                containment_failure=oracle.containment_failure,
                state_changed=oracle.state_changed,
                cost_usd=actual_cost,
                latency_seconds=round(time.perf_counter() - started_at, 6),
                trial_id=str(metadata.get("trial_id", "")),
                repetition=int(metadata.get("repetition", 0)),
                schedule_position=int(metadata.get("schedule_position", 0)),
            )
            self._emit(run_id, "run_finished", asdict(outcome))
            return outcome
        except Exception as exc:
            abort_class = "harness_error"
            terminal = False
            if isinstance(exc, SafetyViolation):
                abort_class = "model_conformance"
                terminal = True
            elif type(exc).__name__ == "BudgetExceeded":
                abort_class = "budget_guard"
                terminal = True
            elif type(exc).__name__ == "ProviderError":
                message = str(exc).lower()
                if any(marker in message for marker in ("http 429", "http 500", "http 502", "http 503", "http 504", "timeout", "connection")):
                    abort_class = "provider_transient"
                else:
                    abort_class = "provider_configuration"
                    terminal = True
            latency_seconds = round(time.perf_counter() - started_at, 6)
            self._emit(
                run_id,
                "run_aborted",
                {
                    "error_type": type(exc).__name__,
                    "message": str(exc),
                    "abort_class": abort_class,
                    "terminal": terminal,
                    "cost_usd": actual_cost,
                    "latency_seconds": latency_seconds,
                    "scenario_id": scenario.scenario_id,
                    "provider": self.adapter.provider,
                    "model": self.adapter.model,
                    "trial": metadata,
                },
            )
            return RunOutcome(
                run_id=run_id,
                scenario_id=scenario.scenario_id,
                provider=self.adapter.provider,
                model=self.adapter.model,
                decision="abort",
                rule_id="",
                oracle_passed=False,
                containment_failure=False,
                state_changed=False,
                cost_usd=actual_cost,
                latency_seconds=latency_seconds,
                trial_id=str(metadata.get("trial_id", "")),
                repetition=int(metadata.get("repetition", 0)),
                schedule_position=int(metadata.get("schedule_position", 0)),
                aborted=True,
                abort_class=abort_class,
                terminal=terminal,
                error=f"{type(exc).__name__}: {exc}",
            )


def build_cost_guard(raw_config: dict[str, Any]) -> CostGuard:
    pricing = {
        provider: (
            float(values.get("input", 0.0)),
            float(values.get("output", 0.0)),
        )
        for provider, values in raw_config[
            "pricing_usd_per_million_tokens"
        ].items()
    }
    limits = {
        provider: float(value)
        for provider, value in raw_config["hard_budget_usd"].items()
        if provider != "total"
    }
    limits.setdefault("mock", 0.0)
    pricing.setdefault("mock", (0.0, 0.0))
    return CostGuard(
        provider_limits=limits,
        total_limit=float(raw_config["hard_budget_usd"]["total"]),
        pricing_per_million=pricing,
    )
