app/generation/services/generation_service.py

from __future__ import annotations

import hashlib
import json
from datetime import datetime, timedelta, timezone
from typing import Literal

from app.core.config import Settings
from app.core.logger import get_logger
from app.generation.domain.enums import (
    GenerationJobStatus,
    WorkerCancellationStatus,
    WorkerJobStatus,
)
from app.generation.domain.errors import (
    GenerationCapabilityUnsupportedError,
    GenerationCancellationError,
    GenerationIdempotencyConflictError,
    GenerationInputAssetNotFoundError,
    GenerationOutputError,
    GenerationProviderUnavailableError,
    GenerationValidationError,
)
from app.generation.model_registry import GenerationModelRegistry, GenerationModelView
from app.generation.models import GenerationJob, GenerationRequest
from app.generation.domain.runtime import WorkerJob
from app.generation.providers.registry import GenerationProviderRegistry
from app.generation.repositories.generation import GenerationRepository
from app.generation.repositories.generation import GenerationDispatchRecord
from app.generation.schemas.requests import (
    GenerationJobView,
    GenerationProviderView,
    GenerationRequestCreate,
    GenerationRequestView,
)
from app.generation.services.output_ingestion import GenerationOutputIngestor
from app.security.assets import CanonicalAssetNotFoundError, CanonicalAssetService
from app.security.database import SecurityDatabase

logger = get_logger(__name__)


class GenerationService:
    """Transport-neutral generation facade.

    This service owns request validation, tenant isolation, durable job
    creation, and idempotency.  It intentionally never opens a WAN or FLUX
    connection; adding a worker adapter later is the only route to accepting
    generation work.
    """

    def __init__(
        self,
        *,
        settings: Settings,
        database: SecurityDatabase,
        assets: CanonicalAssetService,
        repository: GenerationRepository,
        providers: GenerationProviderRegistry,
        models: GenerationModelRegistry,
        output_ingestor: GenerationOutputIngestor,
    ) -> None:
        self.settings = settings
        self.database = database
        self.assets = assets
        self.repository = repository
        self.providers = providers
        self.models = models
        self.output_ingestor = output_ingestor
        self.ready = False

    async def initialize(self) -> None:
        # SecurityDatabase initialization is owned by the application
        # lifespan.  It creates the local SQLite metadata or verifies the
        # production PostgreSQL migration before this service becomes ready.
        self.ready = self.settings.generation_enabled and await self.database.schema_ready()
        if self.ready:
            logger.info("generation foundation initialized")
        elif self.settings.generation_enabled:
            logger.warning(
                "generation foundation schema unavailable; apply the generation migration"
            )
        else:
            logger.info("generation foundation disabled")

    def ensure_ready(self) -> None:
        if not self.settings.generation_enabled:
            raise GenerationProviderUnavailableError("Generation is disabled.")
        if not self.ready:
            raise GenerationProviderUnavailableError(
                "Generation storage is unavailable. Apply the generation migration."
            )

    def list_providers(self) -> list[GenerationProviderView]:
        self.ensure_ready()
        return [self._provider_view(adapter.provider) for adapter in self.providers.list()]

    def get_provider(self, provider: str) -> GenerationProviderView:
        self.ensure_ready()
        return self._provider_view(provider)

    def _provider_view(self, provider: str) -> GenerationProviderView:
        """Return public availability only after model verification.

        An adapter's ``available`` flag intentionally means its server-owned
        configuration can be contacted; it is used by the dispatcher to run
        health/readiness discovery.  It must not by itself be advertised to
        clients as a usable provider.  At least one registered model must have
        passed the exact health, identity, modality, and readiness checks.
        """

        adapter = self.providers.get(provider)
        verified_model_available = any(
            model.available for model in self.models.list(provider_id=adapter.provider)
        )
        return GenerationProviderView(
            capabilities=adapter.capabilities,
            available=adapter.available and verified_model_available,
        )

    def list_models(self, *, provider: str | None = None) -> list[GenerationModelView]:
        self.ensure_ready()
        return self.models.list(provider_id=provider.strip().lower() if provider else None)

    def get_model(self, provider: str, model_id: str) -> GenerationModelView:
        self.ensure_ready()
        return self.models.get(provider.strip().lower(), model_id)

    async def refresh_provider_runtime(self, provider: str) -> list[GenerationModelView]:
        """Verify configured worker metadata and readiness before advertising models.

        No startup caller invokes this while the registry is empty.  Future
        provider integrations may call it after their trusted configuration is
        loaded; a health success alone never makes a model available.
        """

        self.ensure_ready()
        adapter = self.providers.get(provider)
        try:
            health = await adapter.health()
            if health.status.value != "healthy":
                self.models.mark_unavailable(adapter.provider)
                return self.models.list(provider_id=adapter.provider)
            info = await adapter.info()
            readiness = await adapter.ready()
        except Exception:
            # Do not expose a worker exception or leave a stale model marked
            # available. The future provider worker owns detailed diagnostics.
            self.models.mark_unavailable(adapter.provider)
            raise
        return self.models.verify_readiness(
            provider_id=adapter.provider,
            worker_info=info,
            readiness=readiness,
            provider_configured=adapter.available,
        )

    async def create(
        self,
        *,
        workspace_id: str,
        user_id: str,
        payload: GenerationRequestCreate,
        idempotency_key: str,
        project_id: str | None = None,
        brand_kit_version_id: str | None = None,
        product_surface: Literal["generation", "ai_studio"] = "generation",
    ) -> GenerationRequestView:
        self.ensure_ready()
        key = idempotency_key.strip()
        if not key:
            raise GenerationValidationError("Idempotency-Key must not be blank.")

        base_spec = self._request_spec(payload)
        requested_fingerprint = self._fingerprint(
            provider=payload.provider,
            model_id=payload.model_id,
            modality=payload.modality.value,
            input_asset_id=payload.input_asset_id,
            project_id=project_id,
            brand_kit_version_id=brand_kit_version_id,
            product_surface=product_surface,
            spec=base_spec,
        )
        existing = await self.repository.get_by_idempotency(workspace_id, key, user_id=user_id)
        if existing is not None:
            request, job = existing
            if request.request_fingerprint != requested_fingerprint:
                raise GenerationIdempotencyConflictError(
                    "Idempotency-Key is already associated with a different generation request."
                )
            return GenerationRequestView.from_records(request, job)

        adapter = self.providers.get(payload.provider)
        model = adapter.model_capability(payload.model_id)
        registered_model = self.models.get(adapter.provider, payload.model_id)
        if not registered_model.available:
            raise GenerationProviderUnavailableError(
                f"Generation model '{payload.model_id}' is not ready to accept work."
            )
        if registered_model.model.modality != model.modality:
            raise GenerationCapabilityUnsupportedError(
                f"Generation model '{payload.model_id}' has inconsistent capability metadata."
            )
        if model.modality != payload.modality:
            raise GenerationCapabilityUnsupportedError(
                f"Model '{payload.model_id}' does not support {payload.modality.value} generation."
            )
        if payload.input_asset_id and not model.input_asset_supported:
            raise GenerationCapabilityUnsupportedError(
                f"Model '{payload.model_id}' does not support canonical input assets."
            )
        if not adapter.available:
            raise GenerationProviderUnavailableError(
                f"Generation provider '{adapter.provider}' is not ready to accept work."
            )

        input_asset = None
        if payload.input_asset_id:
            try:
                input_asset = await self.assets.get_owned_by_id(
                    workspace_id=workspace_id,
                    user_id=user_id,
                    asset_id=payload.input_asset_id,
                )
            except CanonicalAssetNotFoundError as exc:
                raise GenerationInputAssetNotFoundError(
                    "Input asset was not found in this workspace."
                ) from exc

        normalized = await adapter.validate_request(payload)
        if input_asset is not None:
            await adapter.validate_input_asset(payload, input_asset)
        spec = self._validated_spec(payload, normalized)
        fingerprint = self._fingerprint(
            provider=adapter.provider,
            model_id=payload.model_id,
            modality=payload.modality.value,
            input_asset_id=payload.input_asset_id,
            project_id=project_id,
            brand_kit_version_id=brand_kit_version_id,
            product_surface=product_surface,
            spec=spec,
        )
        if fingerprint != requested_fingerprint:
            raise GenerationCapabilityUnsupportedError(
                "Generation adapter returned a request outside the foundation contract."
            )

        request = GenerationRequest(
            workspace_id=workspace_id,
            created_by_user_id=user_id,
            provider=adapter.provider,
            model_id=payload.model_id,
            modality=payload.modality.value,
            input_asset_id=payload.input_asset_id,
            project_id=project_id,
            product_surface=product_surface,
            spec_json=spec,
            request_fingerprint=fingerprint,
            idempotency_key=key,
            status=GenerationJobStatus.QUEUED.value,
        )
        job = GenerationJob(
            generation_request_id="",
            workspace_id=workspace_id,
            provider=adapter.provider,
            status=GenerationJobStatus.QUEUED.value,
            brand_kit_version_id=brand_kit_version_id,
            max_attempts=self.settings.generation_job_retry_limit + 1,
        )
        created_request, created_job = await self.repository.create(
            request=request, job=job, user_id=user_id
        )
        if created_request.request_fingerprint != fingerprint:
            raise GenerationIdempotencyConflictError(
                "Idempotency-Key is already associated with a different generation request."
            )
        logger.info(
            "generation request queued",
            extra={
                "generation_request_id": created_request.id,
                "generation_job_id": created_job.id,
                "workspace_id": workspace_id,
                "provider": adapter.provider,
                "model_id": payload.model_id,
            },
        )
        return GenerationRequestView.from_records(created_request, created_job)

    async def list_requests(
        self,
        workspace_id: str,
        user_id: str,
        *,
        offset: int = 0,
        limit: int = 100,
        product_surface: str | None = None,
    ) -> list[GenerationRequestView]:
        self.ensure_ready()
        records = await self.repository.list_requests(
            workspace_id,
            user_id=user_id,
            offset=offset,
            limit=limit,
            product_surface=product_surface,
        )
        return [GenerationRequestView.from_records(request, job) for request, job in records]

    async def get_request(
        self, workspace_id: str, user_id: str, request_id: str
    ) -> GenerationRequestView:
        self.ensure_ready()
        request, job = await self.repository.get_request(workspace_id, request_id, user_id=user_id)
        return GenerationRequestView.from_records(request, job)

    async def get_job(self, workspace_id: str, user_id: str, job_id: str) -> GenerationJobView:
        self.ensure_ready()
        return GenerationJobView.from_record(
            await self.repository.get_job(workspace_id, job_id, user_id=user_id)
        )

    async def cancel(self, workspace_id: str, user_id: str, job_id: str) -> GenerationJobView:
        self.ensure_ready()
        current = await self.repository.get_job(workspace_id, job_id, user_id=user_id)
        status = GenerationJobStatus(current.status)
        active = {
            GenerationJobStatus.QUEUED,
            GenerationJobStatus.SUBMITTING,
            GenerationJobStatus.RUNNING,
            GenerationJobStatus.CANCEL_REQUESTED,
        }
        if status in active and current.external_job_id:
            adapter = self.providers.get(current.provider)
            try:
                result = await adapter.cancel(external_job_id=current.external_job_id)
            except Exception as exc:
                try:
                    normalized = adapter.normalize_error(exc)
                except Exception:
                    # An adapter's error normalizer is diagnostic-only. A
                    # programming error there must not expose worker context
                    # or turn cancellation into a false success.
                    normalized = None
                if isinstance(normalized, GenerationCancellationError):
                    raise normalized
                # Never surface provider exception text: worker clients can
                # include request context and third-party adapters may not.
                raise GenerationCancellationError() from None
            if result.status is WorkerCancellationStatus.UNSUPPORTED:
                raise GenerationCapabilityUnsupportedError(
                    f"{adapter.provider} does not support remote generation cancellation."
                )
            if result.status is WorkerCancellationStatus.FAILED:
                raise GenerationCancellationError()
            if result.status is WorkerCancellationStatus.CANCELLED:
                job = await self.repository.transition_job(
                    workspace_id,
                    job_id,
                    GenerationJobStatus.CANCELLED,
                    user_id=user_id,
                )
            else:
                # The worker accepted a request but did not claim that its GPU
                # inference stopped.  Preserve this distinction locally.
                if status is GenerationJobStatus.CANCEL_REQUESTED:
                    job = current
                else:
                    job = await self.repository.transition_job(
                        workspace_id,
                        job_id,
                        GenerationJobStatus.CANCEL_REQUESTED,
                        user_id=user_id,
                    )
        else:
            # Queued jobs are safe to cancel locally. Active jobs without an
            # external identity remain cancellation_requested for the worker
            # to observe before/after it establishes a remote job.
            job = await self.repository.cancel(workspace_id, job_id, user_id=user_id)
        logger.info(
            "generation cancellation requested",
            extra={
                "generation_job_id": job.id,
                "workspace_id": workspace_id,
                "status": job.status,
            },
        )
        return GenerationJobView.from_record(job)

    async def bind_provider_job(
        self,
        *,
        workspace_id: str,
        user_id: str,
        job_id: str,
        worker_job_id: str,
        provider_metadata: dict[str, object] | None = None,
    ) -> GenerationJobView:
        """Bind a worker-issued opaque job ID through the tenant boundary.

        This is an internal worker integration hook, not an HTTP endpoint.
        Future dispatchers must use it after their state-machine transition,
        rather than updating ``GenerationJob.external_job_id`` directly.
        """

        self.ensure_ready()
        current = await self.repository.get_job(workspace_id, job_id, user_id=user_id)
        adapter = self.providers.get(current.provider)
        job = await self.repository.bind_external_job(
            workspace_id,
            job_id,
            user_id=user_id,
            provider=adapter.provider,
            external_job_id=worker_job_id,
            provider_metadata=provider_metadata,
        )
        return GenerationJobView.from_record(job)

    async def ingest_completed_provider_output(
        self,
        *,
        workspace_id: str,
        user_id: str,
        job_id: str,
        worker_job: WorkerJob | None = None,
    ) -> GenerationJobView:
        """Ingest a completed worker output into the canonical asset pipeline.

        The public API cannot call this method or provide an output URL/path.
        A future trusted dispatcher calls it only after status reconciliation.
        The job's bound provider and external ID are the sole source of worker
        selection, preserving workspace ownership throughout the handoff.
        """

        self.ensure_ready()
        current = await self.repository.get_job(workspace_id, job_id, user_id=user_id)
        if current.output_asset_id:
            return GenerationJobView.from_record(current)
        if not current.external_job_id:
            raise GenerationOutputError("Generation job has no bound worker output.")
        if GenerationJobStatus(current.status) not in {
            GenerationJobStatus.QUEUED,
            GenerationJobStatus.SUBMITTING,
            GenerationJobStatus.RUNNING,
            GenerationJobStatus.CANCEL_REQUESTED,
        }:
            raise GenerationOutputError("Generation job is not ready to ingest a worker output.")

        adapter = self.providers.get(current.provider)
        worker_job = worker_job or await adapter.get_job(external_job_id=current.external_job_id)
        if worker_job.external_job_id != current.external_job_id:
            raise GenerationOutputError("Generation worker returned an unexpected job identity.")
        if worker_job.status is not WorkerJobStatus.COMPLETED or worker_job.output is None:
            raise GenerationOutputError("Generation worker output is not ready.")

        async with adapter.stream_output(worker_job.output) as chunks:
            asset = await self.output_ingestor.ingest(
                workspace_id=workspace_id,
                user_id=user_id,
                generation_request_id=current.generation_request_id,
                generation_job_id=current.id,
                provider_id=adapter.provider,
                project_id=(
                    await self.repository.get_request(
                        workspace_id,
                        current.generation_request_id,
                        user_id=user_id,
                    )
                )[0].project_id,
                output=worker_job.output,
                chunks=chunks,
            )
        job = await self.repository.complete_with_output_asset(
            workspace_id,
            job_id,
            user_id=user_id,
            external_job_id=current.external_job_id,
            output_asset_id=asset.id,
            provider_metadata={
                "provider_output_id": worker_job.output.provider_output_id,
                "output_type": worker_job.output.output_type.value,
                "mime_type": worker_job.output.mime_type,
                "worker_metadata": worker_job.metadata,
            },
        )
        logger.info(
            "generation worker output ingested",
            extra={
                "generation_job_id": job.id,
                "workspace_id": workspace_id,
                "provider": adapter.provider,
                "output_asset_id": asset.id,
            },
        )
        return GenerationJobView.from_record(job)

    # The following methods form the trusted worker boundary.  They are not
    # mounted as REST/MCP/SDK endpoints, so external callers cannot claim
    # arbitrary jobs, supply an external job identity, or alter state.

    async def claim_dispatchable_jobs(self, *, limit: int) -> list[GenerationDispatchRecord]:
        self.ensure_ready()
        return await self.repository.claim_dispatchable(limit=limit)

    async def list_reconcilable_jobs(self, *, limit: int) -> list[GenerationDispatchRecord]:
        self.ensure_ready()
        return await self.repository.list_reconcilable(
            limit=limit,
            lease_seconds=self.settings.generation_job_stale_after_seconds,
        )

    async def set_worker_reconciliation_due(
        self,
        *,
        workspace_id: str,
        user_id: str,
        job_id: str,
        due_at: datetime | None,
    ) -> None:
        self.ensure_ready()
        await self.repository.set_reconciliation_due(
            workspace_id,
            job_id,
            user_id=user_id,
            due_at=due_at,
        )

    async def fail_stale_unbound_submissions(self) -> int:
        self.ensure_ready()
        stale_before = datetime.now(timezone.utc) - timedelta(
            seconds=self.settings.generation_job_stale_after_seconds
        )
        return await self.repository.fail_stale_unbound_submissions(stale_before=stale_before)

    async def transition_worker_job(
        self,
        *,
        workspace_id: str,
        user_id: str,
        job_id: str,
        status: GenerationJobStatus,
        error_code: str | None = None,
        error_message: str | None = None,
        next_attempt_at: datetime | None = None,
    ) -> GenerationJobView:
        self.ensure_ready()
        job = await self.repository.transition_job(
            workspace_id,
            job_id,
            status,
            error_code=error_code,
            error_message=error_message,
            next_attempt_at=next_attempt_at,
            user_id=user_id,
        )
        return GenerationJobView.from_record(job)

    def submission_retry_at(
        self,
        *,
        attempt_number: int,
        category,
        http_status: int | None,
    ) -> datetime | None:
        """Classify a known-not-accepted submission with the shared policy."""

        from app.generation.domain.retry import GenerationRetryPolicy

        decision = GenerationRetryPolicy(
            max_retries=self.settings.ai_worker_max_retries,
            backoff_seconds=self.settings.ai_worker_retry_backoff_seconds,
        ).decide(
            category=category,
            http_status=http_status,
            retry_number=max(0, attempt_number - 1),
            idempotent=True,
        )
        if not decision.retryable:
            return None
        return datetime.now(timezone.utc) + timedelta(seconds=decision.delay_seconds)

    async def close(self) -> None:
        await self.providers.close()

    @staticmethod
    def _validated_spec(
        payload: GenerationRequestCreate, normalized: dict[str, object]
    ) -> dict[str, object]:
        """Guard adapter output before it becomes durable API-visible input."""

        if not isinstance(normalized, dict):
            raise GenerationCapabilityUnsupportedError(
                "Generation adapter returned an invalid validated request."
            )
        # A provider cannot rewrite the base intent or smuggle an opaque
        # payload into persistence. Provider controls are accepted only when
        # a closed typed schema has been reviewed at this boundary.
        allowed = {"prompt", "wan", "flux"}
        extra = set(normalized) - allowed
        if extra:
            raise GenerationCapabilityUnsupportedError(
                "Generation adapter returned unsupported request fields."
            )
        prompt = normalized.get("prompt")
        if not isinstance(prompt, str) or prompt != payload.prompt:
            raise GenerationCapabilityUnsupportedError(
                "Generation adapter returned an invalid prompt."
            )
        expected = GenerationService._request_spec(payload)
        if normalized != expected:
            raise GenerationCapabilityUnsupportedError(
                "Generation adapter returned unsupported request controls."
            )
        return expected

    @staticmethod
    def _request_spec(payload: GenerationRequestCreate) -> dict[str, object]:
        """Return the complete typed public generation intent for persistence."""

        spec: dict[str, object] = {"prompt": payload.prompt}
        if payload.wan is not None:
            # ``exclude_unset`` retains the worker's audited defaults when a
            # caller does not send a parameter, instead of silently inventing
            # a MediaRouter-side override.
            spec["wan"] = payload.wan.model_dump(exclude_unset=True)
        if payload.flux is not None:
            # As with WAN, preserve the audited worker defaults for omitted
            # controls.  The closed Pydantic schema rejects every unsupported
            # FLUX field before the request can be persisted or dispatched.
            spec["flux"] = payload.flux.model_dump(exclude_unset=True)
        return spec

    @staticmethod
    def _fingerprint(
        *,
        provider: str,
        model_id: str,
        modality: str,
        input_asset_id: str | None,
        project_id: str | None,
        brand_kit_version_id: str | None,
        product_surface: str,
        spec: dict[str, object],
    ) -> str:
        canonical = json.dumps(
            {
                "provider": provider,
                "model_id": model_id,
                "modality": modality,
                "input_asset_id": input_asset_id,
                "project_id": project_id,
                "brand_kit_version_id": brand_kit_version_id,
                "product_surface": product_surface,
                "spec": spec,
            },
            sort_keys=True,
            separators=(",", ":"),
            ensure_ascii=False,
        )
        return hashlib.sha256(canonical.encode("utf-8")).hexdigest()