from __future__ import annotations

import json
import hashlib
import random
from dataclasses import dataclass, field
from typing import Any, Callable

from .models import AuthorizationError, Budget, ExternalStateUnknown, RetriableError, ValidationError, canonical_hash
from .policy import PolicyGateway, WRITE_TOOLS
from .storage import Store


@dataclass
class FaultInjector:
    seed: int = 0
    plans: dict[str, str] = field(default_factory=dict)
    hits: dict[str, int] = field(default_factory=dict)
    trigger_attempts: dict[str, int] = field(default_factory=dict)
    trace: list[dict[str, Any]] = field(default_factory=list)

    def __post_init__(self) -> None:
        for tool in self.plans:
            self.trigger_attempts.setdefault(tool, 1)

    @classmethod
    def seeded(cls, seed: int, tool: str = "apply_mock_update") -> "FaultInjector":
        rng = random.Random(seed)
        mode = rng.choice(["transient_before_commit", "timeout_after_commit", "none"])
        trigger = 1 + rng.randrange(2)
        plans = {} if mode == "none" else {tool: mode}
        return cls(seed, plans, {}, {tool: trigger})

    def plan_digest(self) -> str:
        payload = {"plans": self.plans, "trigger_attempts": self.trigger_attempts}
        return hashlib.sha256(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()

    def before(self, tool: str) -> None:
        self.hits[tool] = self.hits.get(tool, 0) + 1
        event = {"tool": tool, "attempt": self.hits[tool], "phase": "before", "plan": self.plans.get(tool, "none")}
        self.trace.append(event)
        if self.plans.get(tool) == "transient_before_commit" and self.hits[tool] == self.trigger_attempts.get(tool, 1):
            event["raised"] = "RetriableError"
            raise RetriableError("injected transient error")

    def after(self, tool: str) -> None:
        event = {"tool": tool, "attempt": self.hits.get(tool, 0), "phase": "after", "plan": self.plans.get(tool, "none")}
        self.trace.append(event)
        if self.plans.get(tool) == "timeout_after_commit" and self.hits.get(tool) == self.trigger_attempts.get(tool, 1):
            event["raised"] = "ExternalStateUnknown"
            raise ExternalStateUnknown("injected response loss")


class ToolExecutor:
    SCHEMAS: dict[str, dict[str, type]] = {
        "get_case": {"case_id": str},
        "get_policy": {"case_id": str},
        "get_result": {"case_id": str},
        "save_review_draft": {"case_id": str, "draft_version": int},
        "request_approval": {"case_id": str, "approval_id": str, "requested_tool": str, "requested_args": dict},
        "apply_mock_update": {"case_id": str, "target": str, "value": str, "expected_version": int},
    }

    def __init__(self, store: Store, policy: PolicyGateway, run_id: str, faults: FaultInjector | None = None, budget: Budget | None = None, now: Callable[[], float] | None = None, idempotency_enabled: bool = True, state_policy_enabled: bool = True) -> None:
        self.store = store
        self.policy = policy
        self.run_id = run_id
        self.faults = faults or FaultInjector()
        self.budget = budget
        self.now = now or (lambda: 0.0)
        self.idempotency_enabled = idempotency_enabled
        self.state_policy_enabled = state_policy_enabled

    def _validate(self, target_case: str, tool: str, arguments: dict[str, Any]) -> None:
        schema = self.SCHEMAS.get(tool)
        if schema is None or set(arguments) != set(schema):
            raise ValidationError("tool argument fields do not match strict schema")
        for name, expected in schema.items():
            if type(arguments[name]) is not expected:
                raise ValidationError(f"tool argument type invalid:{name}")
            if expected is str and not (1 <= len(arguments[name]) <= 128):
                raise ValidationError(f"tool argument length invalid:{name}")
        if arguments["case_id"] != target_case:
            raise ValidationError("case_id argument mismatch")
        if tool == "request_approval":
            if arguments["requested_tool"] != "apply_mock_update":
                raise ValidationError("only mock update can be requested")
            nested = arguments["requested_args"]
            expected_nested = self.SCHEMAS["apply_mock_update"]
            if set(nested) != set(expected_nested) or any(type(nested[key]) is not expected_nested[key] for key in expected_nested):
                raise ValidationError("requested arguments fail strict schema")

    def execute(self, actor_case: str, target_case: str, tool: str, arguments: dict[str, Any], request_id: str, idempotency_key: str | None = None, approval_id: str | None = None) -> dict[str, Any]:
        if self.budget:
            try:
                self.budget.preflight_tool_call()
            except Exception as error:
                self.store.safety_block(self.run_id, target_case, "TOOL_PREFLIGHT", {"tool": tool, "error_class": type(error).__name__}, self.now())
                raise
        self.store.log(self.run_id, target_case, "STARTED", {"received": True}, request_id, tool, self.now())
        try:
            self._validate(target_case, tool, arguments)
            self.store.log(self.run_id, target_case, "VALIDATED", {"strict_schema": True}, request_id, tool, self.now())
            if tool in WRITE_TOOLS and (not request_id or not idempotency_key):
                raise ValidationError("write requires request and idempotency identifiers")
            if actor_case != target_case:
                raise AuthorizationError("cross-case access blocked")
            if idempotency_key and self.idempotency_enabled:
                existing = self.store.effect(idempotency_key, self.run_id, target_case)
                if existing:
                    if existing["args_hash"] != canonical_hash(arguments):
                        raise ValidationError("idempotency key reused with different arguments")
                    self.store.log(self.run_id, target_case, "RESULT_REUSED", {"idempotent": True}, request_id, tool, self.now())
                    return json.loads(existing["result"])
            elif idempotency_key:
                self.store.log(self.run_id, target_case, "IDEMPOTENCY_DISABLED", {"result_reuse": False}, request_id, tool, self.now())
            self.policy.authorize(actor_case, target_case, tool, arguments, approval_id, self.run_id, enforce_state=self.state_policy_enabled)
            if not self.state_policy_enabled:
                self.store.log(self.run_id, target_case, "STATE_POLICY_DISABLED", {"ablation_only": True}, request_id, tool, self.now())
            self.faults.before(tool)
            result = self._dispatch(target_case, tool, arguments)
            if tool in WRITE_TOOLS:
                assert idempotency_key is not None
                if tool == "apply_mock_update":
                    if not approval_id:
                        raise ValidationError("approval identifier required")
                    self.policy.validate_approval(approval_id, target_case, self.run_id, tool, arguments)
                    self.store.apply_business_update(target_case, self.run_id, arguments["target"], arguments["value"], arguments["expected_version"], tool, idempotency_key, canonical_hash(arguments), result, approval_id)
                else:
                    self.store.add_effect(target_case, tool, idempotency_key, canonical_hash(arguments), result, self.run_id)
                self.store.log(self.run_id, target_case, "COMMITTED", {"effect_recorded": True}, request_id, tool, self.now())
            self.faults.after(tool)
            self.store.log(self.run_id, target_case, "RESPONSE_RECEIVED", {"result_type": result["type"]}, request_id, tool, self.now())
            self.store.log(self.run_id, target_case, "COMPLETED", {"ok": True}, request_id, tool, self.now())
            return result
        except Exception as error:
            self.store.log(self.run_id, target_case, "FAILED", {"error_class": type(error).__name__}, request_id, tool, self.now())
            raise

    def _dispatch(self, case_id: str, tool: str, arguments: dict[str, Any]) -> dict[str, Any]:
        if tool == "get_case":
            return {"type": "case", "case_id": case_id, "run_id": self.run_id, "state": self.store.state(case_id, self.run_id)}
        if tool == "get_policy":
            return {"type": "policy", "version": self.policy.version}
        if tool == "get_result":
            return {"type": "result", "case_id": case_id, "run_id": self.run_id, "state": self.store.state(case_id, self.run_id)}
        if tool == "save_review_draft":
            return {"type": "draft", "case_id": case_id, "saved": True}
        if tool == "request_approval":
            self.store.request_approval({"id": arguments["approval_id"], "run_id": self.run_id, "case_id": case_id, "tool": arguments["requested_tool"], "args_hash": canonical_hash(arguments["requested_args"]), "requested_at": self.now(), "policy_version": self.policy.version})
            return {"type": "approval_request", "case_id": case_id, "approval_id": arguments["approval_id"], "status": "PENDING"}
        if tool == "apply_mock_update":
            return {"type": "mock_update", "case_id": case_id, "recorded": True, "target": arguments["target"], "value": arguments["value"]}
        raise ValidationError("unknown tool")
