from __future__ import annotations

import hashlib
import json
import re
from datetime import UTC, datetime
from typing import Any, Literal, Protocol

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

from .enums import (
    DeploymentChannel,
    FailureCode,
    FailureOwner,
    HealthStatus,
    MigrationStatus,
    ResidencyPolicy,
    UpgradeClass,
    VerificationStatus,
    WorkloadPriority,
)

CAPABILITY_KEY = re.compile(r"^[a-z][a-z0-9_]*(?:\.[a-z][a-z0-9_]*)+$")
SHA256 = re.compile(r"^[a-f0-9]{64}$")
COMMIT_SHA = re.compile(r"^[a-f0-9]{40,64}$")


class StrictModel(BaseModel):
    model_config = ConfigDict(extra="forbid")


class SchemaDocument(StrictModel):
    schema_: dict[str, Any] = Field(alias="schema")


class ModalityContract(StrictModel):
    input: list[str] = Field(min_length=1)
    output: list[str] = Field(min_length=1)


class VectorContract(StrictModel):
    dimensionality: int | None = Field(default=None, gt=0)
    normalized: bool | None = None
    cross_deployment_compatible: bool = False


class SLOContract(StrictModel):
    latency_p95_ms: int | None = Field(default=None, gt=0)
    availability_percent: float | None = Field(default=None, ge=0, le=100)
    max_concurrency: int | None = Field(default=None, gt=0)


class PrivacyContract(StrictModel):
    classification: Literal["public", "internal", "confidential", "restricted"] = "internal"
    allow_persistence: bool = False
    allow_logging_payloads: bool = False
    allow_network_egress: bool = False


class ResourceRequirement(StrictModel):
    accelerator_required: bool = False
    minimum_vram_mb: int | None = Field(default=None, ge=0)
    cpu_fallback_allowed: bool = False


CapabilityCategory = Literal[
    "TEXT", "RAG", "DOCUMENT", "VISION", "AUDIO", "GENERATION", "ASSISTANTS"
]
CapabilityStability = Literal["stable", "experimental", "blocked", "planned"]
EvaluationType = Literal[
    "retrieval", "ocr", "visual-retrieval", "asr", "tts", "generation", "llm"
]
ResourceClass = Literal["LIGHT", "MEDIUM", "HEAVY", "EXCLUSIVE_GPU"]


class PayloadLimits(StrictModel):
    max_bytes: int = Field(ge=1, le=67_108_864)
    max_batch_count: int = Field(default=1, ge=1, le=64)
    max_width: int | None = Field(default=None, ge=1, le=16_384)
    max_height: int | None = Field(default=None, ge=1, le=16_384)
    max_pages: int | None = Field(default=None, ge=1, le=128)
    max_duration_seconds: float | None = Field(default=None, gt=0, le=3600)


class CapabilityEstateMetadata(StrictModel):
    category: CapabilityCategory
    purpose: str = Field(min_length=1, max_length=1000)
    stability: CapabilityStability
    resource_class: ResourceClass
    evaluation_type: EvaluationType
    consumers: list[str] = Field(default_factory=list)
    payload_limits: PayloadLimits


class FallbackContract(StrictModel):
    allowed: bool = False
    mode: Literal["none", "compatible_deployment", "degrade", "hard_fail"] = "hard_fail"
    capability: str | None = None

    @model_validator(mode="after")
    def validate_fallback(self) -> FallbackContract:
        if not self.allowed and self.mode not in {"none", "hard_fail"}:
            raise ValueError("disabled fallback must use none or hard_fail")
        return self


class CapabilityContractManifest(StrictModel):
    capability: str
    version: int = Field(ge=1)
    description: str = Field(min_length=1)
    input_schema: dict[str, Any]
    output_schema: dict[str, Any]
    modalities: ModalityContract
    languages: list[str] = Field(default_factory=list)
    streaming: bool = False
    structured_output: bool = False
    vector: VectorContract | None = None
    slo: SLOContract = Field(default_factory=SLOContract)
    quality_metrics: list[str] = Field(min_length=1)
    upgrade_class: UpgradeClass
    fallback: FallbackContract = Field(default_factory=FallbackContract)
    privacy: PrivacyContract = Field(default_factory=PrivacyContract)
    resources: ResourceRequirement = Field(default_factory=ResourceRequirement)
    production_priority: WorkloadPriority
    default_residency: ResidencyPolicy
    estate: CapabilityEstateMetadata

    @field_validator("capability")
    @classmethod
    def validate_key(cls, value: str) -> str:
        if not CAPABILITY_KEY.fullmatch(value):
            raise ValueError("capability must be a dotted, lowercase logical key")
        return value

    @model_validator(mode="after")
    def vector_safety(self) -> CapabilityContractManifest:
        vector_output = self.vector is not None or "embedding" in self.capability
        if vector_output and not self.vector:
            raise ValueError("vector-producing capabilities require a vector contract")
        if (
            vector_output
            and self.vector is not None
            and not self.vector.cross_deployment_compatible
            and self.upgrade_class is not UpgradeClass.REQUIRES_REINDEX
        ):
            raise ValueError("incompatible vector spaces require requires_reindex")
        return self


class ProjectDefinition(StrictModel):
    id: str = Field(pattern=r"^[a-z][a-z0-9-]*$")
    name: str = Field(min_length=1)
    description: str = Field(min_length=1)


class ProjectBindingManifest(StrictModel):
    contract_version: int = Field(ge=1)
    channel: DeploymentChannel
    priority: WorkloadPriority
    optional: bool = False
    fallback: FallbackContract = Field(default_factory=FallbackContract)
    migration_support: Literal["none", "reindex", "schema"] = "none"
    slo: SLOContract = Field(default_factory=SLOContract)
    benchmark_suites: list[str] = Field(default_factory=list)


class ProjectManifest(StrictModel):
    project: ProjectDefinition
    bindings: dict[str, ProjectBindingManifest]
    notes: list[str] = Field(default_factory=list)

    @field_validator("bindings")
    @classmethod
    def validate_bindings(
        cls, value: dict[str, ProjectBindingManifest]
    ) -> dict[str, ProjectBindingManifest]:
        for key in value:
            if not CAPABILITY_KEY.fullmatch(key):
                raise ValueError(f"invalid capability binding key: {key}")
        return value


class CandidateManifest(StrictModel):
    id: str
    display_name: str
    source: str
    intended_capabilities: list[str] = Field(min_length=1)
    proposed_role: str
    preferred_runtime: str | None = None
    verification_status: VerificationStatus = VerificationStatus.UNVERIFIED
    deployment_status: Literal["not_deployed"] = "not_deployed"


class CandidateRegistryManifest(StrictModel):
    candidates: list[CandidateManifest]


class BenchmarkSuiteDefinition(StrictModel):
    id: str
    capability: str
    version: int = Field(ge=1)
    purpose: str
    dataset_revision: str
    runnable: bool = False


class BenchmarkSuiteManifest(StrictModel):
    suite: BenchmarkSuiteDefinition
    metrics: dict[str, list[str]]
    performance: dict[str, list[str]]
    release_policy: dict[str, Any]

    @model_validator(mode="after")
    def require_pinned_dataset_for_runnable_suite(self) -> BenchmarkSuiteManifest:
        if self.suite.runnable and self.suite.dataset_revision.startswith("pending"):
            raise ValueError("runnable benchmark suites require a pinned dataset revision")
        return self


class SecurityDefaults(StrictModel):
    trust_remote_code: Literal[False] = False
    prefer_safetensors: Literal[True] = True
    require_exact_revision_for_approval: Literal[True] = True
    require_artifact_digest: Literal[True] = True
    inference_network_egress: Literal[False] = False
    expose_runtime_workers: Literal[False] = False


class LifecycleDefaults(StrictModel):
    automatic_production_promotion: Literal[False] = False
    require_local_benchmark_for_stable: Literal[True] = True
    retain_previous_stable_as_rollback_target: Literal[True] = True
    destructive_cleanup_requires_explicit_approval: Literal[True] = True


class SchedulerDefaults(StrictModel):
    priority_order: list[WorkloadPriority]
    reserve_vram_mb: int = Field(ge=0)
    default_warm_ttl_seconds: int = Field(ge=0)
    benchmark_may_disrupt_production: Literal[False] = False


class UpgradeDefaults(StrictModel):
    default_behavioral_change_class: Literal[UpgradeClass.BEHAVIORAL]
    vector_producer_default_change_class: Literal[UpgradeClass.REQUIRES_REINDEX]
    schema_change_requires_contract_version_bump: Literal[True] = True


class PolicyDefaults(StrictModel):
    security: SecurityDefaults
    lifecycle: LifecycleDefaults
    scheduler: SchedulerDefaults
    upgrades: UpgradeDefaults


class ArtifactProvenance(StrictModel):
    upstream_repository: str
    resolved_commit_sha: str
    filename: str
    artifact_type: str
    sha256: str
    source_artifact_sha256: str | None = None
    derivation_tool: str | None = None
    derivation_tool_version: str | None = None
    derivation_arguments: dict[str, Any] = Field(default_factory=dict)
    imported_at: datetime = Field(default_factory=lambda: datetime.now(UTC))

    @field_validator("resolved_commit_sha")
    @classmethod
    def validate_commit(cls, value: str) -> str:
        value = value.lower()
        if not COMMIT_SHA.fullmatch(value):
            raise ValueError("resolved commit must be a 40-64 character hexadecimal digest")
        return value

    @field_validator("sha256", "source_artifact_sha256")
    @classmethod
    def validate_digest(cls, value: str | None) -> str | None:
        if value is not None and not SHA256.fullmatch(value.lower()):
            raise ValueError("artifact digest must be lowercase SHA-256")
        return value.lower() if value else None

    @model_validator(mode="after")
    def validate_lineage(self) -> ArtifactProvenance:
        derived = self.source_artifact_sha256 is not None
        if derived and (not self.derivation_tool or not self.derivation_tool_version):
            raise ValueError("derived artifacts require derivation tool and version")
        return self


class RuntimeProfileContract(StrictModel):
    runtime_type: Literal["vllm", "transformers", "diffusers", "llama_cpp", "custom"]
    runtime_version: str
    runtime_image_digest: str | None = Field(default=None, pattern=r"^sha256:[a-f0-9]{64}$")
    artifact_sha256: str = Field(pattern=r"^[a-f0-9]{64}$")
    quantization: str | None = None
    context_length: int | None = Field(default=None, gt=0)
    max_concurrency: int = Field(default=1, gt=0)
    launch_arguments: dict[str, Any] = Field(default_factory=dict)
    environment_constraints: dict[str, Any] = Field(default_factory=dict)
    trust_remote_code: bool = False
    network_egress: bool = False

    @property
    def fingerprint(self) -> str:
        payload = self.model_dump(mode="json", exclude_none=True)
        return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()


class ResourceEnvelopeContract(StrictModel):
    deployment_id: str
    accelerator_kind: str
    context_length: int | None = Field(default=None, gt=0)
    concurrency: int = Field(ge=1)
    batch_size: int = Field(ge=1)
    idle_vram_mb: int = Field(ge=0)
    peak_vram_mb: int = Field(ge=0)
    safety_margin_mb: int = Field(default=1024, ge=0)

    @model_validator(mode="after")
    def peak_not_below_idle(self) -> ResourceEnvelopeContract:
        if self.peak_vram_mb < self.idle_vram_mb:
            raise ValueError("peak VRAM cannot be below idle VRAM")
        return self


class LeaseRequest(StrictModel):
    deployment_id: str
    priority: WorkloadPriority
    residency: ResidencyPolicy
    envelope: ResourceEnvelopeContract
    warm_ttl_seconds: int = Field(default=900, ge=0)
    exclusive: bool = False
    allow_cpu_fallback: bool = False


class BenchmarkEnvironmentFingerprint(StrictModel):
    runtime_type: str
    runtime_version: str
    runtime_image_digest: str | None = None
    launch_arguments: dict[str, Any]
    cuda_version: str | None
    driver_version: str | None
    accelerator_name: str | None
    accelerator_uuid: str | None
    context_length: int | None
    concurrency: int = Field(ge=1)
    seed: int | None
    operating_system: str | None
    python_version: str | None = None

    @property
    def digest(self) -> str:
        raw = json.dumps(self.model_dump(mode="json"), sort_keys=True, separators=(",", ":"))
        return hashlib.sha256(raw.encode()).hexdigest()

    def comparable_with(self, other: BenchmarkEnvironmentFingerprint) -> bool:
        relevant = (
            "runtime_type",
            "runtime_version",
            "runtime_image_digest",
            "launch_arguments",
            "cuda_version",
            "driver_version",
            "accelerator_name",
            "context_length",
            "concurrency",
        )
        return all(getattr(self, field) == getattr(other, field) for field in relevant)


class BenchmarkRunContract(StrictModel):
    deployment_id: str
    model_revision_sha: str = Field(pattern=r"^[a-f0-9]{40,64}$")
    artifact_sha256: str = Field(pattern=r"^[a-f0-9]{64}$")
    suite_key: str
    suite_revision: str
    dataset_revision: str
    environment: BenchmarkEnvironmentFingerprint
    started_at: datetime
    completed_at: datetime | None = None


class MigrationContract(StrictModel):
    project_id: str
    capability: str
    source_deployment_id: str
    target_deployment_id: str
    upgrade_class: UpgradeClass
    current_index_ref: str
    shadow_index_ref: str
    status: MigrationStatus = MigrationStatus.PLANNED
    rollback_retain_until: datetime

    @model_validator(mode="after")
    def require_reindex(self) -> MigrationContract:
        if (
            "embedding" in self.capability
            and self.upgrade_class is not UpgradeClass.REQUIRES_REINDEX
        ):
            raise ValueError("embedding migrations must be classified requires_reindex")
        if self.current_index_ref == self.shadow_index_ref:
            raise ValueError("shadow index must be distinct from current production index")
        return self


class LayerHealth(StrictModel):
    process: HealthStatus
    runtime: HealthStatus
    model: HealthStatus
    capability: HealthStatus
    project: HealthStatus
    reasons: list[str] = Field(default_factory=list)


class FailurePolicy(StrictModel):
    code: FailureCode
    owner: FailureOwner
    retryable: bool
    fallback_allowed: bool
    hard_failure: bool


class RuntimeAdapter(Protocol):
    """Stable control-plane boundary; concrete adapters arrive in M5."""

    runtime_type: str

    def validate_profile(self, profile: RuntimeProfileContract) -> None: ...
    def probe(self, profile: RuntimeProfileContract) -> LayerHealth: ...
    def load(self, profile: RuntimeProfileContract, lease_id: str) -> str: ...
    def unload(self, deployment_id: str) -> None: ...