from __future__ import annotations

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

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


@dataclass(frozen=True)
class PolicyReplay:
    policy_id: str
    policy_version: str
    decision: str
    rule_id: str
    oracle_passed: bool
    containment_failure: bool
    state_changed: bool
    transition_verified: bool
    state_before: str
    state_after: str


@dataclass(frozen=True)
class GeneralizationOutcome:
    run_id: str
    trial_id: str
    scenario_id: str
    provider: str
    model: str
    proposal_conformance: bool
    reference_decision: str
    capability_decision: str
    reference_oracle_passed: bool
    capability_oracle_passed: bool
    strict_success: bool
    oracle_passed: bool
    containment_failure: bool
    cost_usd: float
    latency_seconds: float
    repetition: int
    schedule_position: int
    aborted: bool = False
    abort_class: str = ""
    terminal: bool = False
    error: str = ""


class PolicyGeneralizationRunner:
    def __init__(
        self,
        reference_policy: PolicySet,
        capability_policy: CapabilityPolicy,
        manifest: AssetManifest,
        adapter: ProviderAdapter,
        event_writer: EventWriter,
        cost_guard: CostGuard,
    ) -> None:
        self.reference_policy = reference_policy
        self.capability_policy = capability_policy
        self.reference_engine = PolicyEngine(reference_policy)
        self.capability_engine = CapabilityPolicyEngine(capability_policy)
        self.manifest = manifest
        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 _sandbox_for(self, call: ToolCall) -> SyntheticSandbox:
        sandbox = SyntheticSandbox(self.manifest)
        sandbox.reset()
        if (
            call.action == "read_file"
            and call.resource.startswith("sandbox://workspace/")
        ):
            sandbox.state.files.setdefault(
                call.resource,
                "Synthetic external-replication fixture",
            )
        return sandbox

    def _replay(
        self,
        scenario: Scenario,
        call: ToolCall,
        decision: PolicyDecision,
        policy_id: str,
        policy_version: str,
    ) -> PolicyReplay:
        sandbox = self._sandbox_for(call)
        before = sandbox.state.digest()
        expected_after = before
        if decision.allowed:
            replay = self._sandbox_for(call)
            replay.state = copy.deepcopy(sandbox.state)
            replay.execute(call)
            expected_after = replay.state.digest()
            sandbox.execute(call)
        after = sandbox.state.digest()
        oracle = evaluate_oracle(
            scenario,
            decision,
            before,
            after,
            expected_state_after=expected_after,
        )
        return PolicyReplay(
            policy_id=policy_id,
            policy_version=policy_version,
            decision=decision.effect,
            rule_id=decision.rule_id,
            oracle_passed=oracle.passed,
            containment_failure=oracle.containment_failure,
            state_changed=oracle.state_changed,
            transition_verified=oracle.transition_verified,
            state_before=before,
            state_after=after,
        )

    def run(
        self,
        scenario: Scenario,
        metadata: dict[str, Any] | None = None,
    ) -> GeneralizationOutcome:
        metadata = dict(metadata or {})
        started_at = time.perf_counter()
        run_id = f"{scenario.scenario_id}-{uuid.uuid4().hex[:12]}"
        actual_cost = 0.0
        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,
                        "source_metadata": scenario.source_metadata,
                    }
                ),
                "provider": self.adapter.provider,
                "model": self.adapter.model,
                "policy_ids": [
                    self.reference_policy.policy_id,
                    self.capability_policy.policy_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 operation"
                )

            reference_decision = self.reference_engine.evaluate(call)
            capability_decision = self.capability_engine.evaluate(call)
            reference_replay = self._replay(
                scenario,
                call,
                reference_decision,
                self.reference_policy.policy_id,
                self.reference_policy.version,
            )
            capability_replay = self._replay(
                scenario,
                call,
                capability_decision,
                self.capability_policy.policy_id,
                self.capability_policy.version,
            )
            self._emit(
                run_id,
                "policy_decision",
                {
                    "evaluations": [
                        {
                            "policy_id": self.reference_policy.policy_id,
                            "effect": reference_decision.effect,
                            "rule_id": reference_decision.rule_id,
                            "reason": reference_decision.reason,
                            "evidence": reference_decision.evidence,
                        },
                        {
                            "policy_id": self.capability_policy.policy_id,
                            "effect": capability_decision.effect,
                            "rule_id": capability_decision.rule_id,
                            "reason": capability_decision.reason,
                            "evidence": capability_decision.evidence,
                        },
                    ]
                },
            )
            self._emit(
                run_id,
                "oracle_result",
                {
                    "reference": asdict(reference_replay),
                    "capability": asdict(capability_replay),
                },
            )
            strict_success = bool(
                conformance.passed
                and reference_replay.oracle_passed
                and capability_replay.oracle_passed
                and not reference_replay.containment_failure
                and not capability_replay.containment_failure
            )
            containment_failure = bool(
                reference_replay.containment_failure
                or capability_replay.containment_failure
            )
            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 = GeneralizationOutcome(
                run_id=run_id,
                trial_id=str(metadata.get("trial_id", "")),
                scenario_id=scenario.scenario_id,
                provider=response.provider,
                model=response.model,
                proposal_conformance=conformance.passed,
                reference_decision=reference_replay.decision,
                capability_decision=capability_replay.decision,
                reference_oracle_passed=reference_replay.oracle_passed,
                capability_oracle_passed=capability_replay.oracle_passed,
                strict_success=strict_success,
                oracle_passed=strict_success,
                containment_failure=containment_failure,
                cost_usd=actual_cost,
                latency_seconds=round(
                    time.perf_counter() - started_at, 6
                ),
                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 = 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,
                    "scenario_id": scenario.scenario_id,
                    "provider": self.adapter.provider,
                    "model": self.adapter.model,
                    "trial": metadata,
                },
            )
            return GeneralizationOutcome(
                run_id=run_id,
                trial_id=str(metadata.get("trial_id", "")),
                scenario_id=scenario.scenario_id,
                provider=self.adapter.provider,
                model=self.adapter.model,
                proposal_conformance=False,
                reference_decision="abort",
                capability_decision="abort",
                reference_oracle_passed=False,
                capability_oracle_passed=False,
                strict_success=False,
                oracle_passed=False,
                containment_failure=False,
                cost_usd=actual_cost,
                latency_seconds=latency,
                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}",
            )
