from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass, field
from enum import Enum
from typing import Any


class State(str, Enum):
    RECEIVED = "RECEIVED"
    PARSING = "PARSING"
    EXTRACTING = "EXTRACTING"
    VALIDATING = "VALIDATING"
    NEEDS_INFORMATION = "NEEDS_INFORMATION"
    POLICY_CHECK = "POLICY_CHECK"
    READY_FOR_REVIEW = "READY_FOR_REVIEW"
    APPROVAL_PENDING = "APPROVAL_PENDING"
    APPROVED = "APPROVED"
    COMPLETED = "COMPLETED"
    RETRY_PENDING = "RETRY_PENDING"
    MANUAL_REVIEW = "MANUAL_REVIEW"
    PERMANENT_FAILURE = "PERMANENT_FAILURE"
    CANCELLED = "CANCELLED"


TERMINAL_STATES = {
    State.NEEDS_INFORMATION,
    State.COMPLETED,
    State.MANUAL_REVIEW,
    State.PERMANENT_FAILURE,
    State.CANCELLED,
}


class ErrorClass(str, Enum):
    RETRIABLE = "RETRIABLE"
    INPUT_VALIDATION = "INPUT_VALIDATION"
    POLICY = "POLICY"
    AUTHORIZATION = "AUTHORIZATION"
    PERMANENT = "PERMANENT"
    EXTERNAL_STATE_UNKNOWN = "EXTERNAL_STATE_UNKNOWN"


class RuntimeErrorBase(Exception):
    category = ErrorClass.PERMANENT


class RetriableError(RuntimeErrorBase):
    category = ErrorClass.RETRIABLE


class ValidationError(RuntimeErrorBase):
    category = ErrorClass.INPUT_VALIDATION


class PolicyError(RuntimeErrorBase):
    category = ErrorClass.POLICY


class AuthorizationError(RuntimeErrorBase):
    category = ErrorClass.AUTHORIZATION


class ExternalStateUnknown(RuntimeErrorBase):
    category = ErrorClass.EXTERNAL_STATE_UNKNOWN


class CheckpointIncompatible(RuntimeErrorBase):
    category = ErrorClass.PERMANENT


def canonical_hash(value: Any) -> str:
    raw = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
    return hashlib.sha256(raw.encode("utf-8")).hexdigest()


@dataclass(frozen=True)
class EvidenceValue:
    value: Any
    document_id: str
    location: str
    review_required: bool = False


@dataclass
class DocumentResult:
    fields: dict[str, EvidenceValue] = field(default_factory=dict)
    problems: list[dict[str, Any]] = field(default_factory=list)
    untrusted_instructions: int = 0

    @property
    def outcome(self) -> str:
        kinds = {item["kind"] for item in self.problems}
        if "UNREADABLE" in kinds:
            return "MANUAL_REVIEW"
        if "MISSING" in kinds:
            return "NEEDS_INFORMATION"
        if kinds:
            return "MANUAL_REVIEW"
        return "VALID"


@dataclass(frozen=True)
class AgentAction:
    tool: str
    arguments: dict[str, Any]
    request_id: str
    idempotency_key: str | None = None

    @classmethod
    def from_payload(cls, payload: dict[str, Any], allowed_tools: set[str]) -> "AgentAction":
        if type(payload) is not dict:
            raise ValidationError("action payload must be an object")
        expected = {"tool", "arguments", "request_id", "idempotency_key"}
        if set(payload) - expected:
            raise ValidationError("unknown action field")
        if not {"tool", "arguments", "request_id"}.issubset(payload):
            raise ValidationError("missing action field")
        if type(payload["tool"]) is not str or payload["tool"] not in allowed_tools:
            raise ValidationError("tool is not allowlisted")
        if type(payload["arguments"]) is not dict:
            raise ValidationError("arguments must be an object")
        request_id = payload["request_id"]
        if type(request_id) is not str or not (1 <= len(request_id) <= 128):
            raise ValidationError("request_id must be a non-empty bounded string")
        idempotency_key = payload.get("idempotency_key")
        if idempotency_key is not None and (type(idempotency_key) is not str or not (1 <= len(idempotency_key) <= 128)):
            raise ValidationError("idempotency_key must be absent or a non-empty bounded string")
        if payload["tool"] in {"save_review_draft", "request_approval", "apply_mock_update"} and idempotency_key is None:
            raise ValidationError("write action requires idempotency_key")
        return cls(
            tool=payload["tool"],
            arguments=payload["arguments"],
            request_id=payload["request_id"],
            idempotency_key=payload.get("idempotency_key"),
        )


@dataclass
class Budget:
    max_steps: int = 30
    max_calls: int = 12
    max_agent_calls: int = 8
    max_tool_calls: int = 12
    max_cost_units: int = 20
    max_repeat: int = 2
    steps: int = 0
    calls: int = 0
    agent_calls: int = 0
    tool_calls: int = 0
    cost_units: int = 0
    seen_actions: dict[str, int] = field(default_factory=dict)

    def preflight_step(self) -> None:
        if self.steps + 1 > self.max_steps:
            raise PolicyError("step limit exceeded before transition")
        self.steps += 1

    def preflight_agent_call(self, cost: int = 1) -> None:
        if self.agent_calls + 1 > self.max_agent_calls or self.cost_units + cost > self.max_cost_units:
            raise PolicyError("agent or cost limit exceeded before call")
        self.agent_calls += 1
        self.cost_units += cost

    def record_agent_signature(self, action_signature: str) -> None:
        self.seen_actions[action_signature] = self.seen_actions.get(action_signature, 0) + 1
        if self.seen_actions[action_signature] > self.max_repeat:
            raise PolicyError("repeat limit exceeded")

    def agent_call(self, action_signature: str, cost: int = 1) -> None:
        self.preflight_agent_call(cost)
        self.record_agent_signature(action_signature)

    def preflight_tool_call(self, cost: int = 1) -> None:
        if self.tool_calls + 1 > self.max_tool_calls or self.calls + 1 > self.max_calls or self.cost_units + cost > self.max_cost_units:
            raise PolicyError("tool or cost limit exceeded before call")
        self.tool_calls += 1
        self.calls += 1
        self.cost_units += cost

    def tool_call(self, cost: int = 1) -> None:
        self.preflight_tool_call(cost)

    def to_dict(self) -> dict[str, Any]:
        return {
            "max_steps": self.max_steps, "max_calls": self.max_calls, "max_agent_calls": self.max_agent_calls,
            "max_tool_calls": self.max_tool_calls, "max_cost_units": self.max_cost_units, "max_repeat": self.max_repeat,
            "steps": self.steps, "calls": self.calls, "agent_calls": self.agent_calls, "tool_calls": self.tool_calls,
            "cost_units": self.cost_units, "seen_actions": dict(self.seen_actions),
        }

    @classmethod
    def from_dict(cls, value: dict[str, Any]) -> "Budget":
        fields = {name: value[name] for name in ("max_steps", "max_calls", "max_agent_calls", "max_tool_calls", "max_cost_units", "max_repeat", "steps", "calls", "agent_calls", "tool_calls", "cost_units", "seen_actions") if name in value}
        return cls(**fields)

    def step(self, action_signature: str | None = None, cost: int = 0) -> None:
        self.preflight_step()
        if self.cost_units + cost > self.max_cost_units:
            self.steps -= 1
            raise PolicyError("cost limit exceeded before step")
        self.cost_units += cost
        if action_signature:
            self.calls += 1
            self.seen_actions[action_signature] = self.seen_actions.get(action_signature, 0) + 1
            if self.seen_actions[action_signature] > self.max_repeat:
                raise PolicyError("repeat limit exceeded")
        if self.calls > self.max_calls:
            raise PolicyError("resource limit exceeded")
