from __future__ import annotations

from datetime import datetime, timezone
from enum import Enum
from hashlib import sha256
import json
from typing import Any, Literal
from uuid import uuid4

from pydantic import BaseModel, ConfigDict, Field


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


def _canonical_payload(value: BaseModel | dict[str, Any]) -> str:
    if isinstance(value, BaseModel):
        payload = value.model_dump(mode="json", exclude={"contract_id", "created_at"})
    else:
        payload = value
    return json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True)


class ContractKind(str, Enum):
    object = "Object"
    shared_signal = "SharedSignal"
    plan_delta = "PlanDelta"
    dataset_schema = "DatasetSchema"
    access_policy = "AccessPolicy"
    audit_log = "AuditLog"
    reproducibility_manifest = "ReproducibilityManifest"


class BaseContract(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    contract_id: str = Field(default_factory=lambda: uuid4().hex)
    kind: str
    name: str
    version: str = "1.0.0"
    created_at: datetime = Field(default_factory=_utcnow)
    parents: list[str] = Field(default_factory=list)

    def digest(self) -> str:
        return sha256(_canonical_payload(self).encode("utf-8")).hexdigest()


class SchemaField(BaseModel):
    model_config = ConfigDict(frozen=True, extra="forbid")

    name: str
    dtype: str
    required: bool = True
    description: str | None = None


class DatasetSchemaContract(BaseContract):
    kind: Literal["DatasetSchema"] = ContractKind.dataset_schema.value
    modality: str
    discipline: str
    fields: list[SchemaField]
    provenance: dict[str, Any] = Field(default_factory=dict)


class AccessPolicyContract(BaseContract):
    kind: Literal["AccessPolicy"] = ContractKind.access_policy.value
    allowed_actions: list[str]
    privacy_budget_epsilon: float | None = None
    requires_attestation: bool = True
    data_residency: list[str] = Field(default_factory=list)


class ObjectContract(BaseContract):
    kind: Literal["Object"] = ContractKind.object.value
    media_type: str
    content_hash: str
    byte_size: int
    source_uri: str | None = None


class SharedSignalContract(BaseContract):
    kind: Literal["SharedSignal"] = ContractKind.shared_signal.value
    signal_name: str
    source_object_id: str
    payload_hash: str
    privacy_budget_spent: float = 0.0


class PlanDeltaContract(BaseContract):
    kind: Literal["PlanDelta"] = ContractKind.plan_delta.value
    objective: str
    requested_contracts: list[str]
    steps: list[str]
    deterministic_key: str


class AuditLogContract(BaseContract):
    kind: Literal["AuditLog"] = ContractKind.audit_log.value
    actor: str
    action: str
    target_contract_id: str
    details: dict[str, Any] = Field(default_factory=dict)


class ReproducibilityManifestContract(BaseContract):
    kind: Literal["ReproducibilityManifest"] = ContractKind.reproducibility_manifest.value
    execution_id: str
    code_hash: str
    data_hashes: list[str]
    environment: dict[str, Any] = Field(default_factory=dict)