"""Bounded finding-specific Stage Q orchestration and batch execution."""

from __future__ import annotations

from collections import Counter
from dataclasses import replace
from statistics import mean, median
from typing import Callable, Mapping

from architecture.models import EvidenceReference, ExecutionStatus, FindingOutcome
from validator.capabilities import CapabilityFamily, CapabilityRole, TOOL_METADATA
from validator.models import (
    CapabilityObservation,
    CaptureScreenshotObservation,
    CheckPackageInstalledObservation,
    ForceStopPackageObservation,
    InspectResumedActivityObservation,
    LaunchActivityObservation,
    ToolCall,
    ToolExecutionStatus,
)
from validator.restricted_tools import RestrictedToolExecutor

from .capabilities import ValidationEnvironment, derive_capability_view
from .context import ValidationContextBuilder, ValidationContextError, canonical_json
from .evaluator import evaluate_validation
from .models import (
    CapabilityView,
    FinishValidationAction,
    ValidationAction,
    ValidationAssessmentAction,
    ValidationBatchResult,
    ValidationBudgetSnapshot,
    ValidationBudgets,
    ValidationContext,
    ValidationMetrics,
    ValidationObservation,
    ValidationResult,
    ValidationStep,
    ValidationTerminationReason,
    ValidationToolAction,
    ValidationPolicyDecision, ValidationPolicyStatus,
)
from .policy import ValidationPolicy, action_key
from .predicates import evaluate_assertions
from .providers import (
    ValidationProvider,
    ValidationProviderError,
    ValidationProviderMetadata,
)


def _provider_metadata(provider: ValidationProvider) -> dict[str, object]:
    if isinstance(provider, ValidationProviderMetadata):
        try:
            return dict(provider.metadata_snapshot())
        except Exception:
            return {
                "provider": provider.provider_type.casefold(),
                "implementation": type(provider).__name__,
                "telemetry_error": "METADATA_UNAVAILABLE",
                "input_tokens": 0,
                "output_tokens": 0,
                "total_latency_ms": 0,
                "estimated_cost_usd": None,
            }
    return {
        "provider": provider.provider_type.casefold(),
        "implementation": type(provider).__name__,
        "input_tokens": 0,
        "output_tokens": 0,
        "total_latency_ms": 0,
        "estimated_cost_usd": None,
    }


def _observation_from_result(observation_id: str, result) -> ValidationObservation:
    raw = result.observation
    refs: tuple[EvidenceReference, ...] = ()
    limitations: tuple[str, ...] = ()
    failure_reason = None
    success = result.execution_status is ToolExecutionStatus.SUCCESS
    summary = f"Stage O {result.tool_name} returned typed evidence."
    data: Mapping[str, object] = {}
    truncated = False
    if isinstance(raw, CapabilityObservation):
        refs = tuple(raw.evidence_refs)
        limitations = raw.limitations
        failure_reason = None if raw.failure_reason is None else raw.failure_reason.value
        success = raw.success and success
        summary = raw.summary
        data = raw.data
        truncated = raw.truncated
    elif isinstance(raw, CheckPackageInstalledObservation):
        data = {
            "return_code": raw.return_code,
            "installed": raw.installed,
            "package_path_present": raw.package_path_present,
        }
        success = success and raw.installed
    elif isinstance(raw, ForceStopPackageObservation):
        data = {"return_code": raw.return_code, "completed": raw.completed}
        success = success and raw.completed
    elif isinstance(raw, LaunchActivityObservation):
        data = {
            "return_code": raw.return_code,
            "status_ok": raw.status_ok,
            "timeout_detected": raw.timeout_detected,
            "stdout": raw.stdout,
            "stderr": raw.stderr,
        }
        success = success and raw.status_ok and not raw.timeout_detected
        if raw.timeout_detected:
            failure_reason = "TOOL_TIMEOUT"
    elif isinstance(raw, InspectResumedActivityObservation):
        data = {
            "return_code": raw.return_code,
            "resumed_activity": raw.resumed_activity,
            "target_matches": raw.target_matches,
        }
    elif isinstance(raw, CaptureScreenshotObservation):
        data = {
            "return_code": raw.return_code,
            "bytes_written": raw.bytes_written,
            "artifact_filename": raw.artifact_filename,
            "valid_png": raw.valid_png,
        }
        success = success and raw.valid_png
    else:
        summary = result.error or "Stage O returned no typed observation."
        success = False
        failure_reason = "TOOL_EXECUTION_FAILURE"
    if result.execution_status is ToolExecutionStatus.TIMEOUT:
        failure_reason = "TOOL_TIMEOUT"
    elif not success and failure_reason is None:
        failure_reason = "TOOL_EXECUTION_FAILURE"
    return ValidationObservation(
        observation_id=observation_id,
        source="STAGE_O_TOOL",
        originating_tool=str(result.tool_name),
        success=success,
        execution_status=result.execution_status.value,
        summary=summary,
        data=data,
        evidence_refs=refs,
        truncated=truncated,
        limitations=limitations,
        failure_reason=failure_reason,
    )


class ValidationOrchestrator:
    def __init__(
        self,
        *,
        finding,
        environment: ValidationEnvironment,
        provider: ValidationProvider,
        supporting_evidence: tuple[ValidationObservation, ...] = (),
        budgets: ValidationBudgets = ValidationBudgets(),
        capability_view: CapabilityView | None = None,
        executor=None,
        context_builder: ValidationContextBuilder | None = None,
    ) -> None:
        self.activity_experiment = environment.activity_experiment
        self.finding = finding
        self.environment = environment
        self.provider = provider
        self.supporting_evidence = supporting_evidence
        self.budgets = budgets
        self.capability_view = capability_view or derive_capability_view(finding, environment)
        self.executor = executor or RestrictedToolExecutor(
            environment.execution_context,
            role=CapabilityRole.VALIDATION,
            max_budget_weight=budgets.max_weighted_tool_cost,
        )
        self.context_builder = context_builder or ValidationContextBuilder(
            budgets.max_context_bytes
        )
        self.policy = ValidationPolicy(budgets)
        self.observations: list[ValidationObservation] = []
        self.assessments: list[ValidationAssessmentAction] = []
        self.steps: list[ValidationStep] = []
        self.provider_turns = 0
        self.provider_failures = 0
        self.provider_retries = 0
        self.malformed_actions = 0
        self.policy_rejections = 0
        self.weighted_cost = 0
        self.mutations = 0
        self.ui_actions = 0
        self.read_only_calls = 0
        self.read_only_cost = 0
        self.mutating_calls = 0
        self.mutating_cost = 0
        self.timeouts = 0
        self.last_action_key: str | None = None
        self.identical_consecutive = 0
        self.finished = False
        self.final_context: ValidationContext | None = None
        self._assertion_text = canonical_json(finding.to_dict())
        if self.activity_experiment is not None:
            self._assertion_text += " " + canonical_json(self.activity_experiment.to_dict())

    def _budget(self) -> ValidationBudgetSnapshot:
        return ValidationBudgetSnapshot(
            self.provider_turns,
            self.weighted_cost,
            self.mutations,
            self.ui_actions,
        )

    def _context(self) -> ValidationContext:
        return self.context_builder.build(
            workspace=self.environment.workspace,
            finding=self.finding,
            supporting_evidence=self.supporting_evidence,
            observations=tuple(self.observations),
            assessments=tuple(self.assessments),
            capability_view=self.capability_view,
            budgets=self.budgets,
            provider_turns=self.provider_turns,
            weighted_cost=self.weighted_cost,
            mutations=self.mutations,
            ui_actions=self.ui_actions,
            **({"activity_experiment": self.activity_experiment} if self.activity_experiment is not None else {}),
        )

    def run(self) -> ValidationResult:
        close = getattr(self.executor, "close", None)
        try:
            result = self._run_loop()
        except Exception:
            if callable(close):
                close()
            raise
        if callable(close):
            try:
                close()
            except Exception:
                return self._result(
                    ValidationTerminationReason.TOOL_FAILURE,
                    "Original permission state could not be restored and verified.",
                )
        return result

    def _run_loop(self) -> ValidationResult:
        if self.capability_view.feasibility.value == "UNSUPPORTED_CAPABILITY":
            return self._result(
                ValidationTerminationReason.CAPABILITY_UNAVAILABLE,
                "Required validation capability is outside the Stage O route.",
            )
        if self.capability_view.feasibility.value == "DYNAMIC_UNAVAILABLE":
            return self._result(
                ValidationTerminationReason.DEVICE_UNAVAILABLE,
                "Dynamic validation requires a usable trusted device.",
            )
        if (
            self.capability_view.feasibility.value
            == "REQUIRED_CAPABILITY_UNAVAILABLE"
        ):
            return self._result(
                ValidationTerminationReason.CAPABILITY_UNAVAILABLE,
                "A required Stage O capability is unavailable in this environment.",
            )
        while self.provider_turns < self.budgets.max_provider_turns:
            try:
                context = self._context()
            except ValidationContextError as error:
                return self._result(ValidationTerminationReason.EVIDENCE_INSUFFICIENT, str(error))
            except Exception:
                return self._result(
                    ValidationTerminationReason.TOOL_FAILURE,
                    "Trusted validation context construction failed safely.",
                )
            self.final_context = context
            before = self._budget()
            self.provider_turns += 1
            try:
                action = self.provider.next_action(context)
            except ValidationProviderError as error:
                self.provider_failures += 1
                category = error.category.value
                ambiguous = category == "AMBIGUOUS_PROVIDER_OUTPUT"
                if category == "MALFORMED_OUTPUT":
                    self.malformed_actions += 1
                    if self.malformed_actions >= self.budgets.max_malformed_actions:
                        return self._provider_failure_step(before, ValidationTerminationReason.MALFORMED_ACTION_LIMIT, category)
                if self.provider_failures > self.budgets.max_provider_retries:
                    reason = (
                        ValidationTerminationReason.AMBIGUOUS_PROVIDER_OUTPUT
                        if ambiguous
                        else ValidationTerminationReason.PROVIDER_FAILURE
                    )
                    return self._provider_failure_step(before, reason, category)
                self.provider_retries += 1
                message = (
                    "Ambiguous provider response rejected without executing an action; retained for bounded retry."
                    if ambiguous
                    else "Provider failure retained for bounded retry."
                )
                self.steps.append(
                    self._step(None, None, before, message, failure=category)
                )
                continue
            except Exception as error:
                self.provider_failures += 1
                return self._provider_failure_step(before, ValidationTerminationReason.PROVIDER_FAILURE, type(error).__name__)

            if not isinstance(action, (ValidationToolAction, ValidationAssessmentAction, FinishValidationAction)):
                self.malformed_actions += 1
                self.steps.append(self._step(None, None, before, "Provider returned an unsupported action object.", failure="MALFORMED_ACTION"))
                if self.malformed_actions >= self.budgets.max_malformed_actions:
                    return self._result(ValidationTerminationReason.MALFORMED_ACTION_LIMIT, "Malformed-action limit was exhausted.")
                continue

            try:
                key = action_key(action) if isinstance(action, ValidationToolAction) else None
            except (TypeError, ValueError):
                self.malformed_actions += 1
                self.steps.append(self._step(None, None, before, "Action arguments are not serializable.", failure="MALFORMED_ACTION"))
                if self.malformed_actions >= self.budgets.max_malformed_actions:
                    return self._result(ValidationTerminationReason.MALFORMED_ACTION_LIMIT, "Malformed-action limit was exhausted.")
                continue
            if key is not None:
                if key == self.last_action_key:
                    self.identical_consecutive += 1
                else:
                    self.last_action_key = key
                    self.identical_consecutive = 0
            else:
                self.last_action_key = None
                self.identical_consecutive = 0
            all_observations = (*self.supporting_evidence, *self.observations)
            decision = self.policy.evaluate(
                action,
                execution_context=self.environment.execution_context,
                available_tools=self.capability_view.tool_names,
                observations=all_observations,
                finished=self.finished,
                weighted_cost=self.weighted_cost,
                mutations=self.mutations,
                ui_actions=self.ui_actions,
                identical_consecutive=self.identical_consecutive,
                allowed_assertion_text=self._assertion_text,
            )
            if decision.approved and self.activity_experiment is not None:
                from .activity_strategy import activity_action_error
                route_error = activity_action_error(self.activity_experiment, action)
                if route_error:
                    decision = ValidationPolicyDecision(ValidationPolicyStatus.REJECTED, route_error, ValidationTerminationReason.POLICY_REJECTED)
            if not decision.approved:
                self.policy_rejections += 1
                feedback = ValidationObservation(
                    observation_id=f"VAL-OBS-{len(self.observations) + 1:03d}",
                    source="POLICY",
                    originating_tool=None,
                    success=False,
                    execution_status="REJECTED",
                    summary=decision.message,
                    data={},
                    failure_reason="POLICY_REJECTED",
                )
                self.observations.append(feedback)
                self.steps.append(self._step(action, decision, before, decision.message, observation_id=feedback.observation_id, failure=(decision.termination_reason.value if decision.termination_reason else "POLICY_REJECTED")))
                if decision.termination_reason in {
                    ValidationTerminationReason.TOOL_BUDGET_EXHAUSTED,
                    ValidationTerminationReason.MUTATION_BUDGET_EXHAUSTED,
                    ValidationTerminationReason.UI_BUDGET_EXHAUSTED,
                    ValidationTerminationReason.REPEATED_ACTION_LIMIT,
                }:
                    return self._result(decision.termination_reason, decision.message)
                if self.policy_rejections >= self.budgets.max_policy_rejections:
                    return self._result(ValidationTerminationReason.POLICY_REJECTED, "Policy-rejection limit was exhausted.")
                continue

            if isinstance(action, FinishValidationAction):
                self.finished = True
                self.steps.append(self._step(action, decision, before, "Provider finished finding-specific validation."))
                return self._result(ValidationTerminationReason.FINISHED, action.summary)
            if isinstance(action, ValidationAssessmentAction):
                self.assessments.append(action)
                self.steps.append(self._step(action, decision, before, "Recorded non-authoritative provider assessment."))
                continue

            metadata = TOOL_METADATA[action.tool_name]
            self.weighted_cost += metadata.budget_weight
            if metadata.read_only:
                self.read_only_calls += 1
                self.read_only_cost += metadata.budget_weight
            else:
                self.mutating_calls += 1
                self.mutating_cost += metadata.budget_weight
            if metadata.mutates_device:
                self.mutations += 1
            if action.tool_name in {"click_ui_element", "input_text", "press_key"}:
                self.ui_actions += 1
            try:
                result = self.executor.execute(
                    ToolCall(
                        call_id=f"VAL-TOOL-{len([x for x in self.observations if x.originating_tool]) + 1:03d}",
                        tool_name=action.tool_name,  # type: ignore[arg-type]
                        arguments=action.arguments,
                    )
                )
                observation = _observation_from_result(
                    f"VAL-OBS-{len(self.observations) + 1:03d}", result
                )
            except Exception:
                observation = ValidationObservation(
                    observation_id=f"VAL-OBS-{len(self.observations) + 1:03d}",
                    source="STAGE_O_TOOL",
                    originating_tool=action.tool_name,
                    success=False,
                    execution_status="FAILURE",
                    summary="Stage O execution failed safely.",
                    data={},
                    failure_reason="TOOL_EXECUTION_FAILURE",
                )
                self.observations.append(observation)
                self.steps.append(self._step(action, decision, before, "Stage O execution failed safely.", observation_id=observation.observation_id, failure=observation.failure_reason))
                return self._result(
                    ValidationTerminationReason.TOOL_FAILURE,
                    "A restricted capability failed safely.",
                )
            if self.activity_experiment is not None:
                observation = replace(observation, data={**observation.data,
                    "executed_arguments": dict(action.arguments),
                    "activity_experiment_sha256": self.activity_experiment.digest})
            self.observations.append(observation)
            self._assertion_text += " " + canonical_json(action.to_dict()["arguments"])
            if result.execution_status is ToolExecutionStatus.TIMEOUT:
                self.timeouts += 1
            self.steps.append(self._step(action, decision, before, f"Executed Stage O tool {action.tool_name}.", observation_id=observation.observation_id, failure=observation.failure_reason))
            if result.execution_status is ToolExecutionStatus.TIMEOUT:
                return self._result(ValidationTerminationReason.TOOL_TIMEOUT, "A restricted capability timed out.")

        return self._result(ValidationTerminationReason.TURN_BUDGET_EXHAUSTED, "Provider-turn budget was exhausted.")

    def _provider_failure_step(self, before, reason, failure):
        self.steps.append(self._step(None, None, before, "Validation provider failed safely.", failure=failure))
        return self._result(reason, "Validation provider could not return one trusted action.")

    def _step(self, action: ValidationAction | None, decision, before, message: str, *, observation_id: str | None = None, failure: str | None = None) -> ValidationStep:
        return ValidationStep(
            self.provider_turns,
            None if action is None else action.to_dict(),
            decision,
            observation_id,
            message,
            before,
            self._budget(),
            failure,
        )

    def _result(self, termination: ValidationTerminationReason, detail: str) -> ValidationResult:
        metadata = _provider_metadata(self.provider)
        if self.activity_experiment is not None:
            metadata.update({"activity_experiment": self.activity_experiment.to_dict(), "activity_experiment_sha256": self.activity_experiment.digest, "oracle_version": self.activity_experiment.version})
        assessment = self.assessments[-1] if self.assessments else None
        all_observations = (*self.supporting_evidence, *self.observations)
        predicates = evaluate_assertions(
            finding=self.finding,
            assessment=assessment,
            feasibility=self.capability_view.feasibility,
            observations=all_observations,
            activity_experiment=self.activity_experiment,
        )
        outcome = evaluate_validation(
            finding=self.finding,
            feasibility=self.capability_view.feasibility,
            termination_reason=termination,
            observations=all_observations,
            assessment=assessment,
            predicates=predicates,
        )
        metrics = ValidationMetrics(
            provider_turns=self.provider_turns,
            provider_retries=self.provider_retries,
            provider_failures=self.provider_failures,
            tool_calls=self.read_only_calls + self.mutating_calls,
            weighted_tool_cost=self.weighted_cost,
            read_only_tool_calls=self.read_only_calls,
            read_only_tool_cost=self.read_only_cost,
            mutating_tool_calls=self.mutating_calls,
            mutating_tool_cost=self.mutating_cost,
            ui_operations=self.ui_actions,
            policy_rejections=self.policy_rejections,
            malformed_actions=self.malformed_actions,
            timeouts=self.timeouts,
            input_tokens=int(metadata.get("input_tokens", 0) or 0),
            output_tokens=int(metadata.get("output_tokens", 0) or 0),
            latency_ms=int(metadata.get("total_latency_ms", 0) or 0),
            api_cost_usd=(None if metadata.get("estimated_cost_usd") is None else str(metadata["estimated_cost_usd"])),
        )
        return ValidationResult(
            finding=self.finding,
            canonical_finding_id=self.finding.speculative_finding.finding_id,
            outcome=outcome,
            feasibility=self.capability_view.feasibility,
            termination_reason=termination,
            termination_detail=detail,
            observations=all_observations,
            steps=tuple(self.steps),
            assessment=assessment,
            predicates=predicates,
            capability_view=self.capability_view,
            budgets=self.budgets,
            metrics=metrics,
            provider_metadata=metadata,
            limitations=tuple(self.environment.workspace.metadata.limitations),
            final_context=self.final_context,
        )


def _equivalent(left, right) -> bool:
    if left.category is not right.category or left.vulnerability_type != right.vulnerability_type:
        return False
    lsf, rsf = left.speculative_finding, right.speculative_finding
    locations = [
        bool(lsf.entry_point and rsf.entry_point and lsf.entry_point == rsf.entry_point),
        bool(lsf.affected_component and rsf.affected_component and lsf.affected_component == rsf.affected_component),
    ]
    left_refs = {(item.artifact, item.locator) for item in lsf.evidence_refs}
    right_refs = {(item.artifact, item.locator) for item in rsf.evidence_refs}
    return any(locations) and bool(left_refs & right_refs)


def validate_findings(
    *,
    findings: tuple,
    environment: ValidationEnvironment | None = None,
    environment_factory: Callable[[object], ValidationEnvironment] | None = None,
    provider_factory: Callable[[object], ValidationProvider],
    supporting_evidence_factory: Callable[[object], tuple[ValidationObservation, ...]] = lambda _: (),
    budgets: ValidationBudgets = ValidationBudgets(),
    executor_factory: Callable[[object], object] | None = None,
) -> ValidationBatchResult:
    if (environment is None) == (environment_factory is None):
        raise ValueError(
            "Provide exactly one validation environment or environment factory."
        )
    environments = {
        finding.speculative_finding.finding_id: environment if environment_factory is None else environment_factory(finding)
        for finding in findings
    }
    canonical: list = []
    aliases: dict[str, str] = {}
    for finding in findings:
        matched = next((item for item in canonical if _equivalent(item, finding)), None)
        if environments[finding.speculative_finding.finding_id].activity_experiment is not None or (matched is not None and environments[matched.speculative_finding.finding_id].activity_experiment is not None):
            matched = None
        if matched is None:
            canonical.append(finding)
        else:
            aliases[finding.speculative_finding.finding_id] = matched.speculative_finding.finding_id
    canonical_results: dict[str, ValidationResult] = {}
    for finding in canonical:
        executor = None if executor_factory is None else executor_factory(finding)
        finding_environment = environments[finding.speculative_finding.finding_id]
        assert finding_environment is not None
        result = ValidationOrchestrator(
            finding=finding,
            environment=finding_environment,
            provider=provider_factory(finding),
            supporting_evidence=supporting_evidence_factory(finding),
            budgets=budgets,
            executor=executor,
        ).run()
        canonical_results[finding.speculative_finding.finding_id] = result
    results: list[ValidationResult] = []
    for finding in findings:
        finding_id = finding.speculative_finding.finding_id
        canonical_id = aliases.get(finding_id, finding_id)
        source = canonical_results[canonical_id]
        if canonical_id == finding_id:
            results.append(source)
            continue
        outcome = FindingOutcome(
            finding_id,
            source.outcome.finding_status,
            source.outcome.execution_status,
            source.outcome.failure_reason,
            source.outcome.failure_stage,
            f"Reused canonical validation evidence from {canonical_id}.",
        )
        results.append(
            replace(
                source,
                finding=finding,
                outcome=outcome,
                alias_of=canonical_id,
                metrics=ValidationMetrics(),
                final_context=None,
            )
        )
    values = list(canonical_results.values())
    counts = Counter(item.outcome.finding_status.value for item in results)
    tool_calls = [item.metrics.tool_calls for item in values] or [0]
    costs = [item.metrics.weighted_tool_cost for item in values] or [0]
    tokens = [item.metrics.input_tokens + item.metrics.output_tokens for item in values] or [0]
    total = len(results)
    aggregate: dict[str, object] = {
        "total_findings": total,
        **{f"{status.casefold()}_count": counts[status] for status in ("VALIDATED", "NOT_VALIDATED", "INCONCLUSIVE", "UNSUPPORTED")},
        **{f"{status.casefold()}_rate": (counts[status] / total if total else 0.0) for status in ("VALIDATED", "NOT_VALIDATED", "INCONCLUSIVE", "UNSUPPORTED")},
        "mean_tool_calls": mean(tool_calls),
        "median_tool_calls": median(tool_calls),
        "mean_weighted_cost": mean(costs),
        "median_weighted_cost": median(costs),
        "mean_tokens": mean(tokens),
        "median_tokens": median(tokens),
    }
    return ValidationBatchResult(tuple(results), aggregate, len(canonical), len(aliases))
