"""Provider-neutral execution identity value objects and versioned readers."""
from __future__ import annotations

import hashlib
import json
from collections.abc import Mapping
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from .domain.provider import HostModelBinding, ServedModelAttestation
from .domain.worker_exec import WorkerWriteCapability
from .json_boundary import JsonBoundaryError, load_owned_object


V2_SCHEMA_VERSION = "2.0"
EXECUTION_IDENTITY_VERSION = 2
_V1_SCHEMA_VERSIONS = {1, "1", "1.0"}
_V2_IDENTITY_KEYS = frozenset({
    "executionIdentityVersion",
    "participantAssignments",
    "roleExecutions",
    "invocations",
    "attempts",
    "attemptEvidenceSeals",
    "mutationRecoveryDecisions",
    "roleExecutionIndex",
    "staticParticipantAssignments",
    "staticRoleExecutions",
    "dynamicRolePolicy",
})


class ExecutionManifestError(ValueError):
    """Raised when execution identity state is ambiguous or inconsistent."""


@dataclass(frozen=True)
class ParticipantAssignment:
    participant_ref: str
    provider: str
    model_ref: str | None
    model_id: str
    runner: str
    host_runtime: str
    host_model_value: str | None
    session_ref: str | None
    window_ref: str | None
    worker_write_capability: WorkerWriteCapability | None
    entry_mode: str
    status: str

    def to_payload(self) -> dict[str, Any]:
        return {
            "participantRef": self.participant_ref,
            "provider": self.provider,
            "modelRef": self.model_ref,
            "modelId": self.model_id,
            "runner": self.runner,
            "hostRuntime": self.host_runtime,
            "hostModelValue": self.host_model_value,
            "sessionRef": self.session_ref,
            "windowRef": self.window_ref,
            "workerWriteCapability": _capability_payload(
                self.worker_write_capability
            ),
            "entryMode": self.entry_mode,
            "status": self.status,
        }


@dataclass(frozen=True)
class RoleExecution:
    role_execution_ref: str
    participant_ref: str
    source_role_execution_ref: str | None
    role: str
    provider: str
    model_ref: str | None
    model_id: str
    ordinal: int
    execution_label: str
    model_spec_digest: str | None
    binding: HostModelBinding | None
    served_model_attestation: ServedModelAttestation
    status: str

    def to_payload(self) -> dict[str, Any]:
        return {
            "roleExecutionRef": self.role_execution_ref,
            "participantRef": self.participant_ref,
            "sourceRoleExecutionRef": self.source_role_execution_ref,
            "role": self.role,
            "provider": self.provider,
            "modelRef": self.model_ref,
            "modelId": self.model_id,
            "ordinal": self.ordinal,
            "executionLabel": self.execution_label,
            "modelSpecDigest": self.model_spec_digest,
            "binding": _binding_payload(self.binding),
            "servedModelAttestation": _attestation_payload(
                self.served_model_attestation
            ),
            "status": self.status,
        }


@dataclass(frozen=True)
class Invocation:
    invocation_ref: str
    participant_ref: str
    role_execution_ref: str
    source_invocation_ref: str | None
    recovery_ref: str | None
    duty_id: str
    dispatch_kind: str
    round: int
    input_digest: str
    write_policy: Mapping[str, Any]
    write_policy_digest: str
    write_enforcement: Mapping[str, Any]

    def to_payload(self) -> dict[str, Any]:
        return {
            "invocationRef": self.invocation_ref,
            "participantRef": self.participant_ref,
            "roleExecutionRef": self.role_execution_ref,
            "sourceInvocationRef": self.source_invocation_ref,
            "recoveryRef": self.recovery_ref,
            "dutyId": self.duty_id,
            "dispatchKind": self.dispatch_kind,
            "round": self.round,
            "inputDigest": self.input_digest,
            "writePolicy": dict(self.write_policy),
            "writePolicyDigest": self.write_policy_digest,
            "writeEnforcement": dict(self.write_enforcement),
        }


@dataclass(frozen=True)
class Attempt:
    invocation_ref: str
    attempt: int
    started_at: str
    finished_at: str | None
    status: str
    result_path: str | None
    error_path: str | None
    change_summary: Mapping[str, Any]
    evidence_seal_ref: str | None

    def to_payload(self) -> dict[str, Any]:
        return {
            "invocationRef": self.invocation_ref,
            "attempt": self.attempt,
            "startedAt": self.started_at,
            "finishedAt": self.finished_at,
            "status": self.status,
            "resultPath": self.result_path,
            "errorPath": self.error_path,
            "changeSummary": dict(self.change_summary),
            "evidenceSealRef": self.evidence_seal_ref,
        }


@dataclass(frozen=True)
class ExecutionManifest:
    entry_mode: str
    participant_assignments: tuple[ParticipantAssignment, ...] = ()
    role_executions: tuple[RoleExecution, ...] = ()
    invocations: tuple[Invocation, ...] = ()
    attempts: tuple[Attempt, ...] = ()
    attempt_evidence_seals: tuple[Mapping[str, Any], ...] = ()
    mutation_recovery_decisions: tuple[Mapping[str, Any], ...] = ()
    legacy: bool = False
    legacy_payload: Mapping[str, Any] = field(default_factory=dict, compare=False)

    def to_payload(self) -> dict[str, Any]:
        return {
            "schemaVersion": V2_SCHEMA_VERSION,
            "executionIdentityVersion": EXECUTION_IDENTITY_VERSION,
            "entryMode": self.entry_mode,
            "participantAssignments": [
                row.to_payload() for row in self.participant_assignments
            ],
            "roleExecutions": [row.to_payload() for row in self.role_executions],
            "invocations": [row.to_payload() for row in self.invocations],
            "attempts": [row.to_payload() for row in self.attempts],
            "attemptEvidenceSeals": list(self.attempt_evidence_seals),
            "mutationRecoveryDecisions": list(self.mutation_recovery_decisions),
        }


@dataclass(frozen=True)
class IdentitySurface:
    surface: str
    legacy: bool
    payload: Mapping[str, Any]


def execution_label(role: str, provider: str, model_id: str, ordinal: int) -> str:
    return f"{role}-{provider}-{model_id}-{ordinal}"


def execution_identity_for_ref(
    manifest: ExecutionManifest, ref: str
) -> RoleExecution:
    """Resolve one canonical role execution for state, usage, errors and activity views."""
    match = next(
        (row for row in manifest.role_executions if row.role_execution_ref == ref),
        None,
    )
    if match is None:
        raise ExecutionManifestError(f"unknown roleExecutionRef: {ref}")
    return match


def stored_execution_label(row: Mapping[str, Any]) -> str:
    """Return the stored label; do not parse it into role/provider/model."""
    label = row.get("executionLabel")
    if not isinstance(label, str) or not label.strip():
        raise ExecutionManifestError("executionLabel is missing")
    return label


STORED_IDENTITY_KEYS = (
    "participantRef",
    "roleExecutionRef",
    "invocationRef",
    "attempt",
    "executionLabel",
)


def stored_identity(payload: Mapping[str, Any] | None) -> dict[str, Any]:
    """Copy stored identity fields. Never infer them from labels."""
    if not isinstance(payload, Mapping):
        return {}
    copied: dict[str, Any] = {}
    for key in STORED_IDENTITY_KEYS:
        if key not in payload:
            continue
        value = payload[key]
        if value in (None, ""):
            continue
        copied[key] = value
    if "executionLabel" in copied:
        copied["executionLabel"] = stored_execution_label(copied)
    return copied if "roleExecutionRef" in copied or "executionLabel" in copied else {}


def reject_partial_v2_payload(payload: Mapping[str, Any]) -> None:
    present = [key for key in _V2_IDENTITY_KEYS if key in payload]
    if not present:
        return
    required = (
        "executionIdentityVersion",
        "participantAssignments",
        "roleExecutions",
        "invocations",
        "attempts",
    )
    missing = [key for key in required if key not in payload]
    if missing:
        raise ExecutionManifestError(
            "partial v2 execution identity payload is rejected: "
            + ", ".join(missing)
        )


def model_spec_digest(
    model_spec: Mapping[str, Any] | object,
    binding: HostModelBinding,
    *,
    write_policy_digest: str | None = None,
    served_model_attestation: ServedModelAttestation | None = None,
) -> str:
    """Digest preparation facts; invocation policy and observations are excluded."""
    del write_policy_digest, served_model_attestation
    canonical = {
        "modelSpec": _canonical_model_spec(model_spec),
        "binding": _binding_payload(binding),
    }
    encoded = json.dumps(
        canonical, sort_keys=True, separators=(",", ":"), ensure_ascii=False
    ).encode("utf-8")
    return "sha256:" + hashlib.sha256(encoded).hexdigest()


def read_identity_surface(path: Path, *, surface: str) -> IdentitySurface:
    payload = _read_json_object(path)
    return identity_surface_from_payload(payload, surface=surface)


def identity_surface_from_payload(
    payload: Mapping[str, Any],
    *,
    surface: str,
) -> IdentitySurface:
    version = payload.get("schemaVersion")
    has_v2_identity = bool(_V2_IDENTITY_KEYS.intersection(payload))
    if version == V2_SCHEMA_VERSION and surface == "run-manifest":
        reject_partial_v2_payload(payload)
    if version in _V1_SCHEMA_VERSIONS:
        if has_v2_identity:
            raise ExecutionManifestError("mixed execution identity version")
        return IdentitySurface(surface, True, payload)
    if version != V2_SCHEMA_VERSION:
        raise ExecutionManifestError(f"unsupported execution identity schema: {version!r}")
    if payload.get("executionIdentityVersion") != EXECUTION_IDENTITY_VERSION:
        raise ExecutionManifestError("mixed execution identity version")
    required = _surface_required_keys(surface)
    if any(key not in payload for key in required):
        raise ExecutionManifestError("mixed execution identity version")
    if surface == "task-manifest":
        validate_task_role_execution_index(payload)
    return IdentitySurface(surface, False, payload)


def validate_task_role_execution_index(payload: Mapping[str, Any]) -> None:
    rows = payload.get("roleExecutionIndex")
    if not isinstance(rows, list):
        raise ExecutionManifestError("roleExecutionIndex must be an array")
    required_keys = {"runRef", "roleExecutionRef"}
    seen: set[tuple[str, str]] = set()
    for index, row in enumerate(rows):
        if not isinstance(row, Mapping):
            raise ExecutionManifestError(
                f"roleExecutionIndex[{index}] must be an object"
            )
        if set(row) != required_keys:
            raise ExecutionManifestError(
                f"roleExecutionIndex[{index}] must contain exactly "
                "runRef and roleExecutionRef"
            )
        run_ref = _required_index_text(row, "runRef", index)
        role_ref = _required_index_text(row, "roleExecutionRef", index)
        pair = (run_ref, role_ref)
        if pair in seen:
            raise ExecutionManifestError(
                f"roleExecutionIndex contains a duplicate pair: {run_ref}, {role_ref}"
            )
        seen.add(pair)


def synthesize_v1_execution_view(payload: Mapping[str, Any]) -> ExecutionManifest:
    assignments = _legacy_assignments(payload)
    participants: list[ParticipantAssignment] = []
    roles: list[RoleExecution] = []
    role_ordinals: dict[str, int] = {}
    for index, assignment in enumerate(assignments, start=1):
        participant, role = _legacy_identity_rows(assignment, index, role_ordinals)
        participants.append(participant)
        roles.append(role)
    return ExecutionManifest(
        entry_mode="legacy",
        participant_assignments=tuple(participants),
        role_executions=tuple(roles),
        legacy=True,
        legacy_payload=dict(payload),
    )


def _legacy_assignments(payload: Mapping[str, Any]) -> list[Mapping[str, Any]]:
    assignments: list[Mapping[str, Any]] = []
    lead = payload.get("leadAssignment")
    if isinstance(lead, Mapping):
        assignments.append({"role": "leader", **lead})
    workers = payload.get("workerAssignments")
    if isinstance(workers, list):
        assignments.extend(row for row in workers if isinstance(row, Mapping))
    return assignments


def _legacy_identity_rows(
    assignment: Mapping[str, Any],
    index: int,
    role_ordinals: dict[str, int],
) -> tuple[ParticipantAssignment, RoleExecution]:
    role = _canonical_legacy_role(str(assignment.get("role") or "analyser"))
    role_ordinals[role] = role_ordinals.get(role, 0) + 1
    ordinal = role_ordinals[role]
    provider = str(assignment.get("provider") or "unknown")
    model_id = str(
        assignment.get("modelExecutionValue") or assignment.get("model") or "unknown"
    )
    participant_ref = f"legacy-participant-{index:03d}"
    participant = ParticipantAssignment(
        participant_ref=participant_ref,
        provider=provider,
        model_ref=None,
        model_id=model_id,
        runner=str(assignment.get("runner") or "legacy"),
        host_runtime=str(assignment.get("hostRuntime") or "legacy"),
        host_model_value=_optional_text(assignment.get("hostModelValue")),
        session_ref=_optional_text(assignment.get("sessionId")),
        window_ref=_optional_text(assignment.get("windowId")),
        worker_write_capability=None,
        entry_mode="legacy",
        status="legacy",
    )
    role_row = RoleExecution(
        role_execution_ref=f"legacy-role-exec-{role}-{ordinal:03d}",
        participant_ref=participant_ref,
        source_role_execution_ref=None,
        role=role,
        provider=provider,
        model_ref=None,
        model_id=model_id,
        ordinal=ordinal,
        execution_label=execution_label(role, provider, model_id, ordinal),
        model_spec_digest=None,
        binding=None,
        served_model_attestation=ServedModelAttestation.unknown(),
        status="legacy",
    )
    return participant, role_row


def _canonical_legacy_role(role: str) -> str:
    return {"lead": "leader", "executor": "implementer"}.get(role, role)


def _canonical_model_spec(model_spec: Mapping[str, Any] | object) -> dict[str, Any]:
    if isinstance(model_spec, Mapping):
        return {key: _canonical_value(value) for key, value in model_spec.items()}
    fields = {
        "modelRef": str(getattr(model_spec, "model_ref")),
        "providerId": getattr(model_spec, "provider_id"),
        "modelId": getattr(getattr(model_spec, "model_ref"), "model_id"),
        "displayName": getattr(model_spec, "display_name"),
        "executionValue": getattr(model_spec, "execution_value"),
        "aliases": getattr(model_spec, "aliases"),
        "versionKind": getattr(model_spec, "version_kind"),
        "channelFamily": getattr(model_spec, "channel_family"),
        "supportedRoles": getattr(model_spec, "supported_roles"),
        "selectable": getattr(model_spec, "selectable"),
        "pricing": getattr(model_spec, "pricing"),
    }
    return {key: _canonical_value(value) for key, value in fields.items()}


def _canonical_value(value: Any) -> Any:
    if isinstance(value, (set, frozenset, tuple)):
        return sorted(value) if all(isinstance(item, str) for item in value) else list(value)
    return value


def _binding_payload(binding: HostModelBinding | None) -> dict[str, Any] | None:
    if binding is None:
        return None
    return {
        "runner": binding.runner,
        "catalogExecutionValue": binding.catalog_execution_value,
        "resolvedExecutionValue": binding.resolved_execution_value,
        "hostModelValue": binding.host_model_value,
        "bindingFidelity": binding.binding_fidelity,
        "workerWriteCapability": _capability_payload(
            binding.worker_write_capability
        ),
    }


def _attestation_payload(attestation: ServedModelAttestation) -> dict[str, Any]:
    return {
        "observedModel": attestation.observed_model,
        "normalizedModelRef": attestation.normalized_model_ref,
        "level": attestation.level,
        "source": attestation.source,
    }


def _capability_payload(
    capability: WorkerWriteCapability | None,
) -> dict[str, str] | None:
    if capability is None:
        return None
    return {"maxBoundaryPrecision": capability.max_boundary_precision}


def _read_json_object(path: Path) -> dict[str, Any]:
    try:
        payload = load_owned_object(path, artifact="execution identity")
    except JsonBoundaryError as exc:
        raise ExecutionManifestError(f"cannot read execution identity: {path}") from exc
    if not isinstance(payload, dict):
        raise ExecutionManifestError("execution identity surface must be an object")
    return payload


def _surface_required_keys(surface: str) -> tuple[str, ...]:
    required = {
        "task-manifest": ("roleExecutionIndex",),
        "run-manifest": (
            "entryMode",
            "participantAssignments",
            "roleExecutions",
            "invocations",
            "attempts",
            "attemptEvidenceSeals",
            "mutationRecoveryDecisions",
        ),
        "input-snapshot": (
            "staticParticipantAssignments",
            "staticRoleExecutions",
            "dynamicRolePolicy",
        ),
    }
    try:
        return required[surface]
    except KeyError as exc:
        raise ExecutionManifestError(f"unknown execution identity surface: {surface}") from exc


def _required_index_text(
    row: Mapping[str, Any],
    key: str,
    index: int,
) -> str:
    value = row.get(key)
    if not isinstance(value, str) or not value.strip():
        raise ExecutionManifestError(
            f"roleExecutionIndex[{index}].{key} must be a non-empty string"
        )
    return value


def _optional_text(value: object) -> str | None:
    text = str(value or "").strip()
    return text or None