app.py

from __future__ import annotations

import logging
import os
import sys

from contextlib import asynccontextmanager
from pathlib import Path
from typing import Literal, Optional, Tuple

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)

_HSSD_OBJECT_TYPES = (
    "FURNITURE",
    "MANIPULAND",
    "WALL_MOUNTED",
    "CEILING_MOUNTED",
)


class LegacyTextSearchRequest(BaseModel):
    query: str = Field(min_length=1)
    language: str = Field(default="english")
    cross_modal: bool = Field(default=False)
    top_k: int = Field(default=10, ge=1, le=100)
    high_quality_only: bool = Field(default=False)


class LegacySearchResult(BaseModel):
    asset_id: str
    similarity: float
    caption_en: str
    caption_cn: Optional[str] = None
    objaverse_id: Optional[str] = None
    model_url: Optional[str] = None


class LegacyTextSearchResponse(BaseModel):
    results: list[LegacySearchResult]
    query: str
    algorithm: str
    language: str
    cross_modal: bool
    high_quality_only: bool


class HssdSearchRequest(BaseModel):
    query: str = Field(min_length=1)
    object_type: Literal["FURNITURE", "MANIPULAND", "WALL_MOUNTED", "CEILING_MOUNTED"]
    desired_dimensions_m: Optional[Tuple[float, float, float]] = None
    top_k: int = Field(default=5, ge=1, le=20)


class HssdCandidate(BaseModel):
    artifact_id: str
    hssd_id: str
    name: str
    category: str
    similarity_score: float
    bbox_score: float
    size_m: tuple[float, float, float]
    mesh_format: Literal["glb"]
    download_url: str
    expires_at: str


class HssdSearchResponse(BaseModel):
    query: str
    candidates: list[HssdCandidate]


class AmbientCGSearchRequest(BaseModel):
    query: str = Field(min_length=1)
    top_k: int = Field(default=5, ge=1, le=20)


class AmbientCGTextureUrls(BaseModel):
    color_url: str
    normal_url: str
    roughness_url: str


class AmbientCGCandidate(BaseModel):
    artifact_id: str
    material_id: str
    category: str
    tags: list[str]
    similarity_score: float
    package_download_url: str
    textures: AmbientCGTextureUrls
    expires_at: str


class AmbientCGSearchResponse(BaseModel):
    query: str
    candidates: list[AmbientCGCandidate]


def _parse_env_bool(name: str, default: bool) -> bool:
    raw = os.getenv(name)
    if raw is None:
        return default
    return raw.strip().lower() in {"1", "true", "yes", "on"}


def _repo_root() -> Path:
    return Path(__file__).resolve().parent


def _ensure_local_pythonpath() -> Path:
    repo_root = _repo_root()
    src_path = repo_root / "src"
    if not src_path.exists():
        raise FileNotFoundError(
            f"SceneSmith retrieval source path does not exist: {src_path}"
        )
    for candidate in (repo_root, src_path):
        candidate_str = str(candidate)
        if candidate_str not in sys.path:
            sys.path.insert(0, candidate_str)
    return repo_root


def _candidate_value(candidate: object, field_name: str, default=None):
    if isinstance(candidate, dict):
        return candidate.get(field_name, default)
    return getattr(candidate, field_name, default)


def _runtime_ready(runtime: Optional[object]) -> bool:
    if runtime is None:
        return False
    return bool(getattr(runtime, "ready", True))


def _runtime_service_summary(
    name: str,
    runtime: Optional[object],
    enabled: bool,
) -> dict[str, object]:
    return {
        "service": name,
        "enabled": enabled,
        "ready": _runtime_ready(runtime) if enabled else False,
    }


def _resolve_artifact_path(
    runtime: object,
    artifact_id: str,
    key: Optional[str] = None,
) -> Path:
    artifacts = getattr(runtime, "artifacts", None)
    if artifacts is None:
        raise HTTPException(status_code=500, detail="Artifact store is unavailable")
    try:
        if key is None:
            return artifacts.resolve_file(artifact_id)
        return artifacts.resolve_file(artifact_id, key)
    except Exception as exc:
        if exc.__class__.__name__ in {"ArtifactExpiredError", "ArtifactNotFoundError"}:
            raise HTTPException(status_code=404, detail=str(exc)) from exc
        raise


def _build_hssd_candidate(candidate: object, request: Request) -> HssdCandidate:
    artifact_id = str(_candidate_value(candidate, "artifact_id"))
    return HssdCandidate(
        artifact_id=artifact_id,
        hssd_id=str(_candidate_value(candidate, "hssd_id")),
        name=str(_candidate_value(candidate, "name", "")),
        category=str(_candidate_value(candidate, "category", "")),
        similarity_score=float(_candidate_value(candidate, "similarity_score", 0.0)),
        bbox_score=float(_candidate_value(candidate, "bbox_score", 0.0)),
        size_m=tuple(_candidate_value(candidate, "size_m", (0.0, 0.0, 0.0))),
        mesh_format="glb",
        download_url=str(request.url_for("download_hssd_artifact", artifact_id=artifact_id)),
        expires_at=str(_candidate_value(candidate, "expires_at")),
    )


def _build_ambientcg_candidate(candidate: object, request: Request) -> AmbientCGCandidate:
    artifact_id = str(_candidate_value(candidate, "artifact_id"))
    return AmbientCGCandidate(
        artifact_id=artifact_id,
        material_id=str(_candidate_value(candidate, "material_id")),
        category=str(_candidate_value(candidate, "category", "")),
        tags=list(_candidate_value(candidate, "tags", [])),
        similarity_score=float(_candidate_value(candidate, "similarity_score", 0.0)),
        package_download_url=str(
            request.url_for("download_ambientcg_artifact", artifact_id=artifact_id)
        ),
        textures=AmbientCGTextureUrls(
            color_url=str(
                request.url_for(
                    "download_ambientcg_texture",
                    artifact_id=artifact_id,
                    texture_kind="color",
                )
            ),
            normal_url=str(
                request.url_for(
                    "download_ambientcg_texture",
                    artifact_id=artifact_id,
                    texture_kind="normal",
                )
            ),
            roughness_url=str(
                request.url_for(
                    "download_ambientcg_texture",
                    artifact_id=artifact_id,
                    texture_kind="roughness",
                )
            ),
        ),
        expires_at=str(_candidate_value(candidate, "expires_at")),
    )


def _build_legacy_result(candidate: object, request: Request) -> LegacySearchResult:
    artifact_id = str(_candidate_value(candidate, "artifact_id"))
    hssd_id = str(_candidate_value(candidate, "hssd_id", artifact_id))
    return LegacySearchResult(
        asset_id=hssd_id,
        similarity=float(_candidate_value(candidate, "similarity_score", 0.0)),
        caption_en=str(_candidate_value(candidate, "name", hssd_id)),
        model_url=str(request.url_for("download_hssd_artifact", artifact_id=artifact_id)),
    )


def _build_runtime_bundle(
    *,
    hssd_runtime: Optional[object],
    ambientcg_runtime: Optional[object],
    enable_hssd: bool,
    enable_ambientcg: bool,
) -> Tuple[Optional[object], Optional[object], str]:
    repo_root = _repo_root().resolve()
    if hssd_runtime is not None or ambientcg_runtime is not None:
        return hssd_runtime, ambientcg_runtime, str(repo_root)

    if enable_hssd or enable_ambientcg:
        _ensure_local_pythonpath()

    built_hssd = None
    built_ambientcg = None
    if enable_hssd:
        from hssd_service.config import HssdServiceConfig
        from hssd_service.service import HssdRuntime

        built_hssd = HssdRuntime(config=HssdServiceConfig.from_env())
    if enable_ambientcg:
        from ambientcg_service.config import AmbientCGServiceConfig
        from ambientcg_service.service import AmbientCGRuntime

        built_ambientcg = AmbientCGRuntime(config=AmbientCGServiceConfig.from_env())
    return built_hssd, built_ambientcg, str(repo_root)


def create_app(
    *,
    hssd_runtime: Optional[object] = None,
    ambientcg_runtime: Optional[object] = None,
    enable_hssd: Optional[bool] = None,
    enable_ambientcg: Optional[bool] = None,
) -> FastAPI:
    effective_enable_hssd = (
        enable_hssd
        if enable_hssd is not None
        else True
    )
    effective_enable_ambientcg = (
        enable_ambientcg
        if enable_ambientcg is not None
        else ambientcg_runtime is not None
        or _parse_env_bool("ENABLE_AMBIENTCG", False)
    )

    if not effective_enable_hssd and not effective_enable_ambientcg:
        raise ValueError("At least one SceneSmith retrieval service must be enabled")

    hssd_runtime, ambientcg_runtime, service_root = _build_runtime_bundle(
        hssd_runtime=hssd_runtime,
        ambientcg_runtime=ambientcg_runtime,
        enable_hssd=effective_enable_hssd,
        enable_ambientcg=effective_enable_ambientcg,
    )

    @asynccontextmanager
    async def lifespan(_app: FastAPI):
        started: list[object] = []
        try:
            for runtime in (hssd_runtime, ambientcg_runtime):
                if runtime is None:
                    continue
                start = getattr(runtime, "start", None)
                if callable(start):
                    start()
                    started.append(runtime)
            yield
        finally:
            for runtime in reversed(started):
                stop = getattr(runtime, "stop", None)
                if callable(stop):
                    stop()

    app = FastAPI(
        title="SceneSmith Retrieval Compatibility API",
        lifespan=lifespan,
    )

    def _summary_payload() -> dict[str, object]:
        return {
            "name": "SceneSmith Retrieval Compatibility API",
            "version": "0.2.0",
            "status": "running",
            "backend": "scenesmith",
            "service_root": service_root,
            "scenesmith_root": service_root,
            "services": {
                "hssd": _runtime_service_summary(
                    "hssd",
                    hssd_runtime,
                    effective_enable_hssd,
                ),
                "ambientcg": _runtime_service_summary(
                    "ambientcg",
                    ambientcg_runtime,
                    effective_enable_ambientcg,
                ),
            },
        }

    def _require_hssd_runtime() -> object:
        if not effective_enable_hssd or hssd_runtime is None:
            raise HTTPException(status_code=503, detail="HSSD service is disabled")
        return hssd_runtime

    def _require_ambientcg_runtime() -> object:
        if not effective_enable_ambientcg or ambientcg_runtime is None:
            raise HTTPException(status_code=503, detail="AmbientCG service is disabled")
        return ambientcg_runtime

    def _search_hssd_legacy(search_request: LegacyTextSearchRequest) -> list[object]:
        runtime = _require_hssd_runtime()
        candidates_by_id: dict[str, object] = {}
        per_type_top_k = min(search_request.top_k, 20)

        for object_type in _HSSD_OBJECT_TYPES:
            typed_request = HssdSearchRequest(
                query=search_request.query,
                object_type=object_type,
                top_k=per_type_top_k,
            )
            for candidate in runtime.search(typed_request):
                hssd_id = str(_candidate_value(candidate, "hssd_id", ""))
                similarity = float(_candidate_value(candidate, "similarity_score", 0.0))
                bbox_score = float(_candidate_value(candidate, "bbox_score", 0.0))
                current = candidates_by_id.get(hssd_id)
                if current is None:
                    candidates_by_id[hssd_id] = candidate
                    continue

                current_similarity = float(_candidate_value(current, "similarity_score", 0.0))
                current_bbox_score = float(_candidate_value(current, "bbox_score", 0.0))
                if similarity > current_similarity or (
                    similarity == current_similarity and bbox_score < current_bbox_score
                ):
                    candidates_by_id[hssd_id] = candidate

        sorted_candidates = sorted(
            candidates_by_id.values(),
            key=lambda candidate: (
                -float(_candidate_value(candidate, "similarity_score", 0.0)),
                float(_candidate_value(candidate, "bbox_score", 0.0)),
                str(_candidate_value(candidate, "hssd_id", "")),
            ),
        )
        return sorted_candidates[: search_request.top_k]

    @app.get("/")
    async def root() -> dict[str, object]:
        return _summary_payload()

    @app.get("/health")
    async def health() -> dict[str, object]:
        return _summary_payload()

    @app.get("/readyz")
    async def readyz() -> dict[str, object]:
        readiness = {
            "hssd": not effective_enable_hssd or _runtime_ready(hssd_runtime),
            "ambientcg": not effective_enable_ambientcg or _runtime_ready(ambientcg_runtime),
        }
        if not all(readiness.values()):
            raise HTTPException(status_code=503, detail={"status": "not_ready", **readiness})
        return {"status": "ready", **readiness}

    @app.get("/hssd/healthz")
    async def hssd_healthz() -> dict[str, object]:
        runtime = _require_hssd_runtime()
        return {
            "status": "ok",
            "service": "hssd",
            "ready": _runtime_ready(runtime),
        }

    @app.get("/ambientcg/healthz")
    async def ambientcg_healthz() -> dict[str, object]:
        runtime = _require_ambientcg_runtime()
        return {
            "status": "ok",
            "service": "ambientcg",
            "ready": _runtime_ready(runtime),
        }

    @app.post("/search/text", response_model=LegacyTextSearchResponse)
    async def legacy_search_text(
        search_request: LegacyTextSearchRequest,
        request: Request,
    ) -> LegacyTextSearchResponse:
        if search_request.language not in {"english", "chinese"}:
            raise HTTPException(
                status_code=400,
                detail="Language must be 'english' or 'chinese'",
            )
        if search_request.cross_modal:
            raise HTTPException(
                status_code=400,
                detail="SceneSmith compatibility service does not support cross_modal search",
            )

        try:
            candidates = _search_hssd_legacy(search_request)
        except HTTPException:
            raise
        except Exception as exc:
            logger.exception("SceneSmith legacy HSSD search failed")
            raise HTTPException(status_code=500, detail=str(exc)) from exc

        return LegacyTextSearchResponse(
            results=[_build_legacy_result(candidate, request) for candidate in candidates],
            query=search_request.query,
            algorithm="scenesmith_hssd",
            language=search_request.language,
            cross_modal=False,
            high_quality_only=search_request.high_quality_only,
        )

    @app.post("/hssd/v1/search", response_model=HssdSearchResponse)
    async def hssd_search(
        search_request: HssdSearchRequest,
        request: Request,
    ) -> HssdSearchResponse:
        runtime = _require_hssd_runtime()
        try:
            candidates = runtime.search(search_request)
        except HTTPException:
            raise
        except Exception as exc:
            logger.exception("Direct HSSD search failed")
            raise HTTPException(status_code=500, detail=str(exc)) from exc

        return HssdSearchResponse(
            query=search_request.query,
            candidates=[_build_hssd_candidate(candidate, request) for candidate in candidates],
        )

    @app.get("/hssd/v1/artifacts/{artifact_id}", name="download_hssd_artifact")
    async def download_hssd_artifact(artifact_id: str) -> FileResponse:
        runtime = _require_hssd_runtime()
        file_path = _resolve_artifact_path(runtime, artifact_id)
        return FileResponse(
            path=file_path,
            filename=file_path.name,
            media_type="model/gltf-binary",
        )

    @app.post("/ambientcg/v1/search", response_model=AmbientCGSearchResponse)
    async def ambientcg_search(
        search_request: AmbientCGSearchRequest,
        request: Request,
    ) -> AmbientCGSearchResponse:
        runtime = _require_ambientcg_runtime()
        try:
            candidates = runtime.search(search_request)
        except HTTPException:
            raise
        except Exception as exc:
            logger.exception("Direct AmbientCG search failed")
            raise HTTPException(status_code=500, detail=str(exc)) from exc

        return AmbientCGSearchResponse(
            query=search_request.query,
            candidates=[
                _build_ambientcg_candidate(candidate, request) for candidate in candidates
            ],
        )

    @app.get("/ambientcg/v1/artifacts/{artifact_id}", name="download_ambientcg_artifact")
    async def download_ambientcg_artifact(artifact_id: str) -> FileResponse:
        runtime = _require_ambientcg_runtime()
        file_path = _resolve_artifact_path(runtime, artifact_id)
        return FileResponse(
            path=file_path,
            filename=file_path.name,
            media_type="application/zip",
        )

    @app.get(
        "/ambientcg/v1/artifacts/{artifact_id}/{texture_kind}",
        name="download_ambientcg_texture",
    )
    async def download_ambientcg_texture(artifact_id: str, texture_kind: str) -> FileResponse:
        runtime = _require_ambientcg_runtime()
        if texture_kind not in {"color", "normal", "roughness"}:
            raise HTTPException(status_code=404, detail="Unknown texture kind")
        file_path = _resolve_artifact_path(runtime, artifact_id, texture_kind)
        return FileResponse(path=file_path, filename=file_path.name)

    app.state.hssd_runtime = hssd_runtime
    app.state.ambientcg_runtime = ambientcg_runtime
    return app