app/generation/domain/runtime.py

from __future__ import annotations

import math
import re
from typing import Any

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

from app.generation.domain.enums import (
    GenerationModality,
    WorkerCancellationStatus,
    WorkerErrorCategory,
    WorkerHealthStatus,
    WorkerJobStatus,
    WorkerReadinessStatus,
)

_EXTERNAL_ID = re.compile(r"^[A-Za-z0-9._:-]{1,255}$")
_SHA256 = re.compile(r"^[0-9a-f]{64}$")
_SAFE_FILENAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,254}$")
_BEARER = re.compile(r"(?i)\bbearer\s+[A-Za-z0-9._~+/=-]+")
_HTTP_URL = re.compile(r"(?i)\bhttps?://[^\s\"'<>]+")
_ASSIGNED_SECRET = re.compile(
    r"(?i)(access[_-]?token|refresh[_-]?token|id[_-]?token|client[_-]?secret|"
    r"authorization|api[_-]?key|password|secret|credential)"
    r"([\"']?\s*[:=]\s*[\"']?)([^\"'\s,&}]+)"
)
_SENSITIVE_METADATA_PARTS = frozenset(
    {
        "access_token",
        "refresh_token",
        "id_token",
        "token",
        "secret",
        "authorization",
        "cookie",
        "password",
        "api_key",
        "credential",
        "url",
        "uri",
    }
)
_MAX_METADATA_DEPTH = 8
_MAX_METADATA_ITEMS = 256
_MAX_METADATA_STRING_LENGTH = 8_192
_DROP = object()


def _safe_metadata(value: Any, *, depth: int = 0) -> Any:
    """Return bounded JSON-safe metadata with credential-like content removed."""

    if depth > _MAX_METADATA_DEPTH:
        return _DROP
    if value is None or isinstance(value, bool) or isinstance(value, int):
        return value
    if isinstance(value, float):
        return value if math.isfinite(value) else _DROP
    if isinstance(value, str):
        without_bearer = _BEARER.sub("Bearer [REDACTED]", value)
        without_urls = _HTTP_URL.sub("[REDACTED_URL]", without_bearer)
        redacted = _ASSIGNED_SECRET.sub(
            lambda match: f"{match.group(1)}{match.group(2)}[REDACTED]",
            without_urls,
        )
        return redacted[:_MAX_METADATA_STRING_LENGTH]
    if isinstance(value, dict):
        result: dict[str, object] = {}
        for key, item in list(value.items())[:_MAX_METADATA_ITEMS]:
            if not isinstance(key, str) or len(key) > 255:
                continue
            normalized_key = key.lower().replace("-", "_")
            if any(part in normalized_key for part in _SENSITIVE_METADATA_PARTS):
                continue
            cleaned = _safe_metadata(item, depth=depth + 1)
            if cleaned is not _DROP:
                result[key] = cleaned
        return result
    if isinstance(value, (list, tuple)):
        result: list[object] = []
        for item in list(value)[:_MAX_METADATA_ITEMS]:
            cleaned = _safe_metadata(item, depth=depth + 1)
            if cleaned is not _DROP:
                result.append(cleaned)
        return result
    return _DROP


def safe_worker_metadata(value: Any) -> Any:
    """Remove secret-like fields and non-JSON values before persistence.

    Worker metadata is useful for diagnostics, but it is never a credential
    store.  This helper is deliberately conservative and is applied both at
    worker-transport parsing and at persistence boundaries.
    """

    cleaned = _safe_metadata(value)
    return {} if cleaned is _DROP and isinstance(value, dict) else cleaned


def _safe_metadata_dict(value: Any) -> dict[str, object]:
    cleaned = safe_worker_metadata(value)
    if not isinstance(cleaned, dict):
        raise ValueError("worker metadata must be a JSON object")
    return cleaned


class WorkerModelInfo(BaseModel):
    """One worker-discovered model; never an operator-configured endpoint."""

    model_config = ConfigDict(extra="forbid")

    id: str = Field(min_length=1, max_length=255)
    name: str = Field(min_length=1, max_length=255)
    media_types: list[GenerationModality] = Field(default_factory=list)
    metadata: dict[str, object] = Field(default_factory=dict)

    @field_validator("metadata", mode="before")
    @classmethod
    def clean_metadata(cls, value: Any) -> dict[str, object]:
        return _safe_metadata_dict(value)


class WorkerInfo(BaseModel):
    """Verified non-secret identity returned by a configured worker."""

    model_config = ConfigDict(extra="forbid")

    id: str = Field(min_length=1, max_length=255)
    name: str = Field(min_length=1, max_length=255)
    media_types: list[GenerationModality] = Field(default_factory=list)
    models: list[WorkerModelInfo] = Field(default_factory=list)
    status: WorkerHealthStatus = WorkerHealthStatus.UNKNOWN
    metadata: dict[str, object] = Field(default_factory=dict)

    @field_validator("models")
    @classmethod
    def unique_models(cls, values: list[WorkerModelInfo]) -> list[WorkerModelInfo]:
        if len({value.id for value in values}) != len(values):
            raise ValueError("worker metadata contains duplicate model IDs")
        return values

    @field_validator("metadata", mode="before")
    @classmethod
    def clean_metadata(cls, value: Any) -> dict[str, object]:
        return _safe_metadata_dict(value)


class WorkerHealth(BaseModel):
    """Normalised liveness result. It never implies model availability."""

    model_config = ConfigDict(extra="forbid")

    status: WorkerHealthStatus
    metadata: dict[str, object] = Field(default_factory=dict)

    @field_validator("metadata", mode="before")
    @classmethod
    def clean_metadata(cls, value: Any) -> dict[str, object]:
        return _safe_metadata_dict(value)


class WorkerReadiness(BaseModel):
    """Normalised inference readiness result."""

    model_config = ConfigDict(extra="forbid")

    status: WorkerReadinessStatus
    model_loaded: bool = False
    model_ids: list[str] = Field(default_factory=list)
    metadata: dict[str, object] = Field(default_factory=dict)

    @field_validator("model_ids")
    @classmethod
    def valid_model_ids(cls, values: list[str]) -> list[str]:
        if len(values) != len(set(values)):
            raise ValueError("worker readiness model_ids must be unique")
        for value in values:
            if not value or len(value) > 255 or any(character.isspace() for character in value):
                raise ValueError("worker readiness model_ids contain an invalid value")
        return values

    @field_validator("metadata", mode="before")
    @classmethod
    def clean_metadata(cls, value: Any) -> dict[str, object]:
        return _safe_metadata_dict(value)


class WorkerOutput(BaseModel):
    """A worker-issued descriptor, never a filesystem path or arbitrary URL."""

    model_config = ConfigDict(extra="forbid")

    output_type: GenerationModality
    mime_type: str = Field(
        min_length=3,
        max_length=255,
        pattern=r"^[a-z0-9!#$&^_.+-]+/[a-z0-9!#$&^_.+-]+$",
    )
    provider_output_id: str = Field(min_length=1, max_length=255)
    # A worker may expose a download endpoint only under its configured origin.
    # RemoteWorkerClient rejects absolute URLs, traversal, queries, and fragments.
    download_path: str = Field(min_length=2, max_length=2048)
    filename: str | None = Field(default=None, max_length=255)
    sha256: str | None = Field(default=None, max_length=64)
    byte_size: int | None = Field(default=None, ge=0)
    metadata: dict[str, object] = Field(default_factory=dict)

    @field_validator("provider_output_id")
    @classmethod
    def valid_provider_output_id(cls, value: str) -> str:
        if _EXTERNAL_ID.fullmatch(value) is None:
            raise ValueError("provider_output_id contains unsupported characters")
        return value

    @field_validator("download_path")
    @classmethod
    def valid_download_path(cls, value: str) -> str:
        if (
            not value.startswith("/")
            or "//" in value
            or "\\" in value
            or "%" in value
            or "?" in value
            or "#" in value
            or any(part in {"", ".", ".."} for part in value.split("/")[1:])
        ):
            raise ValueError("download_path must be a safe absolute worker-relative path")
        return value

    @field_validator("filename")
    @classmethod
    def valid_filename(cls, value: str | None) -> str | None:
        if value is not None and _SAFE_FILENAME.fullmatch(value) is None:
            raise ValueError("filename must not contain a path")
        return value

    @field_validator("sha256")
    @classmethod
    def valid_sha256(cls, value: str | None) -> str | None:
        if value is not None and _SHA256.fullmatch(value) is None:
            raise ValueError("sha256 must be a lowercase SHA-256 hex digest")
        return value

    @field_validator("metadata", mode="before")
    @classmethod
    def clean_metadata(cls, value: Any) -> dict[str, object]:
        return _safe_metadata_dict(value)

    @model_validator(mode="after")
    def media_type_matches_output_type(self) -> "WorkerOutput":
        if not self.mime_type.startswith(f"{self.output_type.value}/"):
            raise ValueError("output MIME type does not match its generation modality")
        return self


class WorkerJob(BaseModel):
    """Provider-neutral remote job representation."""

    model_config = ConfigDict(extra="forbid")

    external_job_id: str = Field(min_length=1, max_length=255)
    status: WorkerJobStatus
    output: WorkerOutput | None = None
    error_category: WorkerErrorCategory | None = None
    error_code: str | None = Field(default=None, max_length=100)
    error_message: str | None = Field(default=None, max_length=500)
    metadata: dict[str, object] = Field(default_factory=dict)

    @field_validator("external_job_id")
    @classmethod
    def valid_external_job_id(cls, value: str) -> str:
        if _EXTERNAL_ID.fullmatch(value) is None:
            raise ValueError("external_job_id contains unsupported characters")
        return value

    @field_validator("error_code")
    @classmethod
    def valid_error_code(cls, value: str | None) -> str | None:
        if value is not None and _EXTERNAL_ID.fullmatch(value) is None:
            raise ValueError("worker error_code contains unsupported characters")
        return value

    @field_validator("error_message", mode="before")
    @classmethod
    def clean_error_message(cls, value: Any) -> str | None:
        if value is None:
            return None
        if not isinstance(value, str):
            raise ValueError("worker error_message must be a string")
        cleaned = safe_worker_metadata(value)
        if not isinstance(cleaned, str):  # Defensive: strings are JSON-safe.
            raise ValueError("worker error_message is invalid")
        return cleaned

    @model_validator(mode="after")
    def completed_job_requires_output(self) -> "WorkerJob":
        if self.status is WorkerJobStatus.COMPLETED and self.output is None:
            raise ValueError("completed worker jobs must include an output descriptor")
        return self

    @field_validator("metadata", mode="before")
    @classmethod
    def clean_metadata(cls, value: Any) -> dict[str, object]:
        return _safe_metadata_dict(value)


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

    status: WorkerCancellationStatus
    metadata: dict[str, object] = Field(default_factory=dict)

    @field_validator("metadata", mode="before")
    @classmethod
    def clean_metadata(cls, value: Any) -> dict[str, object]:
        return _safe_metadata_dict(value)