Blame for src/stage0_contracts.py

"""Canonical Stage 0 contracts shared by every Pure system.

The module is intentionally self-contained so separately deployed repositories
can vendor the exact source until the package is published to a package index.
"""

from __future__ import annotations

import hashlib
import inspect
import json
import re
import uuid
from datetime import datetime, timezone
from typing import Any, Literal, Protocol, runtime_checkable

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator


CONTRACT_VERSION = "1.0.0"
CLOUD_EVENTS_SPEC_VERSION = "1.0"
ID_PATTERN = r"^[A-Za-z0-9][A-Za-z0-9._:-]{7,127}$"
SHA256_PATTERN = r"^[0-9a-fA-F]{64}$"
SPEAKER_PATTERN = re.compile(r"^(?:SPEAKER_[0-9]{2,}|P[12]|spk_[A-Za-z0-9_-]+)$")


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


def _new_id(prefix: str) -> str:
    return f"{prefix}_{uuid.uuid4().hex}"


def _canonical_json(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=str)


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


class ArtifactReference(StrictModel):
    uri: str = Field(min_length=1, max_length=2048)
    sha256: str = Field(pattern=SHA256_PATTERN)
    media_type: str = Field(min_length=1, max_length=128)
    size_bytes: int | None = Field(default=None, ge=0)
    encrypted: bool = False
    consent_scope: str | None = Field(default=None, max_length=128)


class ConsentRecord(StrictModel):
    consent_id: str = Field(min_length=8, max_length=128)
    policy_version: str = Field(min_length=1, max_length=128)
    privacy_notice_version: str = Field(min_length=1, max_length=128)
    granted: bool
    permitted_modalities: list[
        Literal["text", "speech", "image", "multimodal", "judgment", "telemetry"]
    ] = Field(min_length=1, max_length=6)
    purposes: list[str] = Field(min_length=1, max_length=16)
    retain_audio: bool = False
    derive_voice_features: bool = False
    granted_at: datetime = Field(default_factory=_utc_now)
    expires_at: datetime | None = None

    @field_validator("granted_at", "expires_at")
    @classmethod
    def timestamps_are_timezone_aware(cls, value: datetime | None) -> datetime | None:
        if value is not None and value.tzinfo is None:
            raise ValueError("consent timestamps must include a timezone")
        return value

    @model_validator(mode="after")
    def audio_permissions_are_explicit(self) -> "ConsentRecord":
        if (self.retain_audio or self.derive_voice_features) and "speech" not in self.permitted_modalities:
            raise ValueError("speech must be permitted before retaining audio or deriving voice features")
        if self.expires_at and self.expires_at <= self.granted_at:
            raise ValueError("expires_at must be later than granted_at")
        return self


class TaskConstraints(StrictModel):
    max_steps: int = Field(default=6, ge=1, le=100)
    max_latency_ms: int = Field(default=10_000, ge=1, le=86_400_000)
    max_cost_usd: float | None = Field(default=None, ge=0)
    allowed_tools: list[str] = Field(default_factory=list, max_length=64)
    human_gates: list[str] = Field(default_factory=list, max_length=32)


class TaskEnvelope(StrictModel):
    specversion: Literal["1.0"] = CLOUD_EVENTS_SPEC_VERSION
    schema_version: Literal["1.0.0"] = CONTRACT_VERSION
    task_id: str = Field(default_factory=lambda: _new_id("task"), pattern=ID_PATTERN)
    goal_id: str | None = Field(default=None, pattern=ID_PATTERN)
    parent_task_id: str | None = Field(default=None, pattern=ID_PATTERN)
    conversation_id: str | None = Field(default=None, max_length=128)
    correlation_id: str | None = Field(default=None, max_length=128)
    causation_id: str | None = Field(default=None, max_length=128)
    system_id: str = Field(min_length=2, max_length=64)
    tenant_id: str = Field(min_length=2, max_length=64)
    app_id: str = Field(min_length=2, max_length=64)
    capability: str = Field(min_length=3, max_length=128)
    idempotency_key: str = Field(pattern=ID_PATTERN)
    created_at: datetime = Field(default_factory=_utc_now)
    deadline: datetime | None = None
    payload: dict[str, Any] = Field(default_factory=dict)
    artifacts: list[ArtifactReference] = Field(default_factory=list, max_length=32)
    consent: ConsentRecord | None = None
    constraints: TaskConstraints = Field(default_factory=TaskConstraints)
    traceparent: str | None = Field(default=None, max_length=256)
    metadata: dict[str, Any] = Field(default_factory=dict)

    @field_validator("created_at", "deadline")
    @classmethod
    def task_timestamps_are_timezone_aware(cls, value: datetime | None) -> datetime | None:
        if value is not None and value.tzinfo is None:
            raise ValueError("task timestamps must include a timezone")
        return value

    @field_validator("payload", "metadata")
    @classmethod
    def json_objects_are_bounded(cls, value: dict[str, Any]) -> dict[str, Any]:
        encoded = _canonical_json(value).encode("utf-8")
        if len(encoded) > 128 * 1024:
            raise ValueError("JSON object must not exceed 128 KiB")
        return value

    @model_validator(mode="after")
    def normalize_task_lineage(self) -> "TaskEnvelope":
        if not self.correlation_id:
            self.correlation_id = self.task_id
        if self.deadline and self.deadline <= self.created_at:
            raise ValueError("deadline must be later than created_at")
        if self.artifacts and not self.consent:
            raise ValueError("artifact-bearing tasks require an explicit consent record")
        return self


class TranscriptSegment(StrictModel):
    segment_id: str = Field(default_factory=lambda: _new_id("seg"), pattern=ID_PATTERN)
    speaker_id: str = Field(min_length=2, max_length=128)
    start_ms: int = Field(ge=0)
    end_ms: int = Field(ge=0)
    text: str = Field(min_length=1, max_length=20_000)
    language: str | None = Field(default=None, max_length=128)
    asr_confidence: float | None = Field(default=None, ge=0, le=1)
    diarization_confidence: float | None = Field(default=None, ge=0, le=1)
    overlap: bool = False
    speaker_source: Literal["diarization", "channel", "turn", "legacy"]

    @field_validator("speaker_id")
    @classmethod
    def speaker_is_stable_identifier(cls, value: str) -> str:
        if not SPEAKER_PATTERN.fullmatch(value):
            raise ValueError("speaker_id must be SPEAKER_00, P1/P2, or an spk_* identifier")
        return value

    @model_validator(mode="after")
    def segment_has_valid_interval(self) -> "TranscriptSegment":
        if self.end_ms < self.start_ms:
            raise ValueError("end_ms must be greater than or equal to start_ms")
        return self


class TranscriptResult(StrictModel):
    schema_version: Literal["1.0.0"] = CONTRACT_VERSION
    recording_id: str = Field(default_factory=lambda: _new_id("recording"), pattern=ID_PATTERN)
    full_text: str = Field(min_length=1, max_length=200_000)
    segments: list[TranscriptSegment] = Field(min_length=1, max_length=10_000)
    language: str | None = Field(default=None, max_length=128)
    asr_model: str = Field(default="legacy-unspecified", min_length=1, max_length=256)
    audio_sha256: str | None = Field(default=None, pattern=SHA256_PATTERN)


class InterpretationCandidate(StrictModel):
    candidate_id: str = Field(default_factory=lambda: _new_id("candidate"), pattern=ID_PATTERN)
    task_id: str = Field(pattern=ID_PATTERN)
    interpretation: str = Field(min_length=1, max_length=20_000)
    dialect: str | None = Field(default=None, max_length=256)
    pragmatics: str | None = Field(default=None, max_length=20_000)
    confidence: float = Field(ge=0, le=1)
    evidence_refs: list[str] = Field(default_factory=list, max_length=64)
    model_id: str = Field(min_length=1, max_length=256)
    model_revision: str = Field(min_length=1, max_length=256)
    prompt_hash: str = Field(pattern=SHA256_PATTERN)


class CriticVerdict(StrictModel):
    verdict_id: str = Field(default_factory=lambda: _new_id("verdict"), pattern=ID_PATTERN)
    task_id: str = Field(pattern=ID_PATTERN)
    candidate_id: str = Field(pattern=ID_PATTERN)
    decision: Literal["accept", "revise", "delegate", "abstain", "escalate"]
    reason_codes: list[str] = Field(min_length=1, max_length=32)
    issues: list[str] = Field(default_factory=list, max_length=64)
    confidence: float = Field(ge=0, le=1)
    evidence_refs: list[str] = Field(default_factory=list, max_length=64)


class ReviewDecision(StrictModel):
    review_id: str = Field(default_factory=lambda: _new_id("review"), pattern=ID_PATTERN)
    task_id: str = Field(pattern=ID_PATTERN)
    reviewer_id: str = Field(pattern=ID_PATTERN)
    decision: Literal["approve", "reject", "request_changes", "abstain"]
    reason_codes: list[str] = Field(min_length=1, max_length=32)
    reviewer_signature: str | None = Field(default=None, max_length=4096)
    decided_at: datetime = Field(default_factory=_utc_now)

    @model_validator(mode="after")
    def terminal_decisions_are_signed(self) -> "ReviewDecision":
        if self.decision in {"approve", "reject"} and not self.reviewer_signature:
            raise ValueError("approve/reject decisions require a reviewer signature")
        return self


class TrainingProposal(StrictModel):
    proposal_id: str = Field(default_factory=lambda: _new_id("training"), pattern=ID_PATTERN)
    system_id: str = Field(min_length=2, max_length=64)
    tenant_id: str = Field(min_length=2, max_length=64)
    task_type: Literal["asr", "semantic_sft", "preference", "embedding"]
    dataset_snapshot_uri: str = Field(min_length=1, max_length=2048)
    dataset_sha256: str = Field(pattern=SHA256_PATTERN)
    base_model_id: str = Field(min_length=1, max_length=256)
    base_model_revision: str = Field(min_length=1, max_length=256)
    method: Literal["lora", "qlora", "full_finetune", "dpo", "orpo"]
    evaluation_suite: str = Field(min_length=1, max_length=256)
    requested_by: str = Field(pattern=ID_PATTERN)
    promotion_requires_human: Literal[True] = True


class AgentManifest(StrictModel):
    agent_id: str = Field(pattern=ID_PATTERN)
    version: str = Field(min_length=1, max_length=64)
    role: str = Field(min_length=1, max_length=256)
    goals: list[str] = Field(min_length=1, max_length=16)
    capabilities: list[str] = Field(min_length=1, max_length=64)
    input_schemas: list[str] = Field(min_length=1, max_length=32)
    output_schemas: list[str] = Field(min_length=1, max_length=32)
    allowed_tools: list[str] = Field(default_factory=list, max_length=64)
    model_revisions: list[str] = Field(default_factory=list, max_length=64)
    max_concurrency: int = Field(default=1, ge=1, le=10_000)


class AgentResult(StrictModel):
    schema_version: Literal["1.0.0"] = CONTRACT_VERSION
    task_id: str = Field(pattern=ID_PATTERN)
    agent_id: str = Field(pattern=ID_PATTERN)
    agent_version: str = Field(min_length=1, max_length=64)
    status: Literal["accepted", "completed", "failed", "escalated", "rejected"]
    action: str = Field(min_length=1, max_length=128)
    reason_codes: list[str] = Field(min_length=1, max_length=32)
    output: dict[str, Any] = Field(default_factory=dict)
    confidence: float | None = Field(default=None, ge=0, le=1)
    evidence_refs: list[str] = Field(default_factory=list, max_length=64)
    model_id: str | None = Field(default=None, max_length=256)
    model_revision: str | None = Field(default=None, max_length=256)
    prompt_hash: str | None = Field(default=None, pattern=SHA256_PATTERN)
    latency_ms: int | None = Field(default=None, ge=0)
    cost_usd: float | None = Field(default=None, ge=0)


class ProviderRequest(StrictModel):
    prompt: str = Field(min_length=1, max_length=100_000)
    model_id: str | None = Field(default=None, max_length=256)
    temperature: float = Field(default=0.2, ge=0, le=2)
    max_tokens: int = Field(default=512, ge=1, le=32_768)


class ProviderResponse(StrictModel):
    text: str = Field(min_length=1, max_length=200_000)
    model_id: str = Field(min_length=1, max_length=256)
    model_revision: str = Field(min_length=1, max_length=256)
    provider: str = Field(min_length=1, max_length=128)
    latency_ms: int | None = Field(default=None, ge=0)


@runtime_checkable
class ModelProvider(Protocol):
    async def generate(self, request: ProviderRequest) -> ProviderResponse: ...


async def invoke_provider(provider: Any, request: ProviderRequest) -> ProviderResponse:
    """Adapt Stage 0 providers plus legacy sync/async generate_* managers."""

    method = getattr(provider, "generate", None)
    argument: Any = request
    if method is None:
        method = getattr(provider, "generate_smart", None) or getattr(provider, "generate_fast", None)
        argument = request.prompt
    if method is None:
        raise TypeError("provider must expose generate(), generate_smart(), or generate_fast()")

    value = method(argument)
    if inspect.isawaitable(value):
        value = await value
    if isinstance(value, ProviderResponse):
        return value

    text = getattr(value, "text", value)
    if not isinstance(text, str) or not text.strip():
        raise TypeError("provider response must contain non-empty text")
    model_id = request.model_id or str(
        getattr(provider, "model_id", None)
        or getattr(provider, "model_flash", None)
        or getattr(provider, "model_name", None)
        or "legacy-unspecified"
    )
    return ProviderResponse(
        text=text.strip(),
        model_id=model_id,
        model_revision=str(getattr(provider, "model_revision", "unversioned")),
        provider=provider.__class__.__name__,
    )


async def assert_provider_conforms(provider: Any) -> ProviderResponse:
    response = await invoke_provider(
        provider,
        ProviderRequest(prompt="Return the single word: ready", model_id="conformance-probe"),
    )
    if not response.text.strip():
        raise AssertionError("provider returned empty text")
    return response


class LegacyGradioCall(StrictModel):
    api_name: str = Field(min_length=2, max_length=128)
    args: list[Any] = Field(default_factory=list, max_length=64)
    task_id: str | None = Field(default=None, pattern=ID_PATTERN)
    correlation_id: str | None = Field(default=None, max_length=128)
    idempotency_key: str | None = Field(default=None, pattern=ID_PATTERN)
    conversation_id: str | None = Field(default=None, max_length=128)
    metadata: dict[str, Any] = Field(default_factory=dict)


class TranscriptNormalizeRequest(StrictModel):
    segments: str | list[dict[str, Any]]
    language: str | None = Field(default=None, max_length=128)


class TaskValidationReceipt(StrictModel):
    valid: Literal[True] = True
    contract_version: Literal["1.0.0"] = CONTRACT_VERSION
    task: TaskEnvelope


class SystemContractManifest(StrictModel):
    contract_version: Literal["1.0.0"] = CONTRACT_VERSION
    system_id: str
    tenant_id: str
    app_id: str
    current_stage: Literal["stage0-contracts"] = "stage0-contracts"
    autonomous_agents_deployed: Literal[False] = False
    capabilities: list[str]
    legacy_api_names: list[str]
    route_prefix: str


def legacy_gradio_to_task(
    call: LegacyGradioCall,
    *,
    system_id: str,
    tenant_id: str,
    app_id: str,
) -> TaskEnvelope:
    normalized_name = call.api_name.strip()
    if not normalized_name.startswith("/"):
        normalized_name = f"/{normalized_name}"
    fingerprint = hashlib.sha256(
        _canonical_json(
            {"system_id": system_id, "api_name": normalized_name, "args": call.args}
        ).encode("utf-8")
    ).hexdigest()
    task_id = call.task_id or f"task_{fingerprint[:32]}"
    idempotency_key = call.idempotency_key or f"idem_{fingerprint[:32]}"
    capability_name = re.sub(r"[^a-z0-9]+", ".", normalized_name.lower()).strip(".")
    return TaskEnvelope(
        task_id=task_id,
        conversation_id=call.conversation_id,
        correlation_id=call.correlation_id or task_id,
        system_id=system_id,
        tenant_id=tenant_id,
        app_id=app_id,
        capability=f"legacy.gradio.{capability_name}",
        idempotency_key=idempotency_key,
        payload={"api_name": normalized_name, "args": call.args},
        metadata={"compatibility_mode": True, **call.metadata},
    )


def _legacy_speaker_id(value: Any, index: int) -> tuple[str, Literal["diarization", "channel", "turn", "legacy"]]:
    text = str(value or "").strip()
    if SPEAKER_PATTERN.fullmatch(text):
        source: Literal["diarization", "channel", "turn", "legacy"] = (
            "diarization" if text.startswith(("SPEAKER_", "spk_")) else "turn"
        )
        return text, source
    match = re.fullmatch(r"speaker\s*(\d+)", text, flags=re.IGNORECASE)
    if match:
        return f"SPEAKER_{max(int(match.group(1)) - 1, 0):02d}", "legacy"
    return f"SPEAKER_{index:02d}", "legacy"


def _milliseconds(item: dict[str, Any], key_ms: str, key_seconds: str, default: int) -> int:
    if item.get(key_ms) is not None:
        return max(int(float(item[key_ms])), 0)
    if item.get(key_seconds) is not None:
        return max(int(float(item[key_seconds]) * 1000), 0)
    return default


def normalize_transcript_segments(
    value: str | list[dict[str, Any]], *, language: str | None = None
) -> list[TranscriptSegment]:
    raw_segments = [{"text": value}] if isinstance(value, str) else value
    if not raw_segments:
        raise ValueError("at least one transcript segment is required")
    normalized: list[TranscriptSegment] = []
    for index, item in enumerate(raw_segments):
        text = str(item.get("text", "")).strip()
        if not text:
            continue
        speaker_id, speaker_source = _legacy_speaker_id(
            item.get("speaker_id", item.get("speaker")), index
        )
        start_ms = _milliseconds(item, "start_ms", "start", 0)
        end_ms = _milliseconds(item, "end_ms", "end", start_ms)
        normalized.append(
            TranscriptSegment(
                segment_id=str(item.get("segment_id") or _new_id("seg")),
                speaker_id=speaker_id,
                start_ms=start_ms,
                end_ms=max(end_ms, start_ms),
                text=text,
                language=str(item.get("language") or language or "") or None,
                asr_confidence=item.get("asr_confidence"),
                diarization_confidence=item.get("diarization_confidence"),
                overlap=bool(item.get("overlap", False)),
                speaker_source=speaker_source,
            )
        )
    if not normalized:
        raise ValueError("at least one non-empty transcript segment is required")
    return normalized


def normalize_transcript_result(
    value: str | list[dict[str, Any]],
    *,
    language: str | None = None,
    asr_model: str = "legacy-unspecified",
    recording_id: str | None = None,
) -> TranscriptResult:
    segments = normalize_transcript_segments(value, language=language)
    return TranscriptResult(
        recording_id=recording_id or _new_id("recording"),
        full_text=" ".join(segment.text for segment in segments),
        segments=segments,
        language=language,
        asr_model=asr_model,
    )


SCHEMA_MODELS = (
    AgentManifest,
    AgentResult,
    ArtifactReference,
    ConsentRecord,
    CriticVerdict,
    InterpretationCandidate,
    LegacyGradioCall,
    ProviderRequest,
    ProviderResponse,
    ReviewDecision,
    TaskEnvelope,
    TrainingProposal,
    TranscriptResult,
    TranscriptSegment,
)


def create_stage0_router(
    *,
    system_id: str,
    tenant_id: str,
    app_id: str,
    capabilities: list[str],
    legacy_api_names: list[str],
    route_prefix: str = "/api/agent/v1",
) -> APIRouter:
    """Create side-effect-free Stage 0 validation and compatibility routes."""

    if not route_prefix.startswith("/"):
        raise ValueError("route_prefix must start with /")
    router = APIRouter(prefix=route_prefix.rstrip("/"), tags=["stage0-contracts"])
    declared_legacy_names = {
        name if name.startswith("/") else f"/{name}" for name in legacy_api_names
    }
    manifest = SystemContractManifest(
        system_id=system_id,
        tenant_id=tenant_id,
        app_id=app_id,
        capabilities=capabilities,
        legacy_api_names=sorted(declared_legacy_names),
        route_prefix=route_prefix.rstrip("/"),
    )

    @router.get("/contracts/manifest", response_model=SystemContractManifest)
    async def get_contract_manifest() -> SystemContractManifest:
        return manifest

    @router.get("/contracts/schemas")
    async def get_contract_schemas() -> dict[str, Any]:
        return {
            "contract_version": CONTRACT_VERSION,
            "schemas": {model.__name__: model.model_json_schema() for model in SCHEMA_MODELS},
        }

    @router.post("/tasks/validate", response_model=TaskValidationReceipt)
    async def validate_task(task: TaskEnvelope) -> TaskValidationReceipt:
        if task.system_id != system_id or task.tenant_id != tenant_id:
            raise HTTPException(status_code=409, detail="task system_id/tenant_id does not match this gateway")
        return TaskValidationReceipt(task=task)

    @router.post("/compat/gradio", response_model=TaskEnvelope)
    async def normalize_legacy_gradio_call(call: LegacyGradioCall) -> TaskEnvelope:
        requested_name = call.api_name if call.api_name.startswith("/") else f"/{call.api_name}"
        if declared_legacy_names and requested_name not in declared_legacy_names:
            raise HTTPException(status_code=422, detail="legacy api_name is not declared by this system")
        call = call.model_copy(update={"api_name": requested_name})
        return legacy_gradio_to_task(
            call, system_id=system_id, tenant_id=tenant_id, app_id=app_id
        )

    @router.post("/transcripts/normalize", response_model=TranscriptResult)
    async def normalize_transcript(request: TranscriptNormalizeRequest) -> TranscriptResult:
        try:
            return normalize_transcript_result(request.segments, language=request.language)
        except ValueError as exc:
            raise HTTPException(status_code=422, detail=str(exc)) from exc

    return router