"""Hugging Face ZeroGPU entry point for dots.tts.

This Space runs an *unmodified*, pip-installed ``dots.tts`` package. All ZeroGPU
specifics live in the local :mod:`dots_tts_zerogpu` patch layer:

* the model source is never forked -- AOTI acceleration is applied by
  intercepting upstream's single compile seam;
* generation buckets are warmed via upstream's own ``run_warmup``.

Startup flow:

1. Apply the AOTI compile-seam patches (before any runtime is built).
2. Preload the runtime on the root process with warmup suppressed (root has only
   CUDA emulation, so no real compile can happen here).
3. Inside a ``@spaces.GPU`` call, run warmup under an AOTI capture session so
   every reachable step module is exported + AOTI-compiled, then hand the
   artifacts back and install them into the serving process.
4. Wrap ``service.generate`` with ``@spaces.GPU`` and launch Gradio.
"""

from __future__ import annotations

import os
import sys
from pathlib import Path
from typing import Any, Callable

REPO_ROOT = Path(__file__).resolve().parent
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))


class _SpacesFallback:
    """No-op ``spaces`` shim for local runs outside Hugging Face Spaces."""

    @staticmethod
    def GPU(*decorator_args, **_decorator_kwargs):
        if decorator_args and callable(decorator_args[0]):
            return decorator_args[0]

        def decorate(fn: Callable[..., Any]) -> Callable[..., Any]:
            return fn

        return decorate


try:
    import spaces  # type: ignore
except Exception:  # pragma: no cover - only used outside Hugging Face Spaces.
    spaces = _SpacesFallback()  # type: ignore


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


def _env_int(name: str, default: int) -> int:
    value = os.environ.get(name)
    if value is None or not value.strip():
        return default
    return int(value)


def _configure_zero_gpu_environment() -> None:
    # optimize=True on CUDA would otherwise select the flex-attention backend,
    # which relies on module-level torch.compile that ZeroGPU cannot run.
    os.environ.setdefault("DOTS_TTS_DELAYED_DIT_BACKEND", "sdpa")
    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")


def _suppress_root_warmup() -> None:
    """Force ``DotsTtsRuntime.from_pretrained`` to skip the in-process warmup.

    The service loads the runtime lazily via ``from_pretrained`` and does not
    expose ``warmup_on_optimize``. On the root process there is no real GPU
    (only CUDA emulation), so any warmup here would run eagerly and slowly --
    real compilation is done separately inside ``@spaces.GPU``.
    """
    from dots_tts.runtime import DotsTtsRuntime

    original_from_pretrained = DotsTtsRuntime.from_pretrained.__func__

    def _from_pretrained(cls, *args, **kwargs):
        kwargs.setdefault("warmup_on_optimize", False)
        return original_from_pretrained(cls, *args, **kwargs)

    DotsTtsRuntime.from_pretrained = classmethod(_from_pretrained)


def _preload_runtime(app_service, app_config):
    runtime, resolved_model_name_or_path = app_service._get_runtime(  # noqa: SLF001
        app_config.default_model_name_or_path,
    )
    return runtime, resolved_model_name_or_path


def main() -> None:
    _configure_zero_gpu_environment()

    from loguru import logger

    from dots_tts_zerogpu import apply_zero_gpu_patches
    from dots_tts_zerogpu.aoti import AOTI_REGISTRY
    from dots_tts_zerogpu.warmup import compile_runtime_aoti

    # Patch the compile seam before any runtime (and thus any step module) exists.
    apply_zero_gpu_patches()
    _suppress_root_warmup()

    import gradio as gr

    from apps.gradio.app import PLAYGROUND_CSS, build_demo, build_playground_theme
    from apps.gradio.service import GradioAppService, build_gradio_app_config
    from dots_tts.utils.logging import configure_logging

    host = os.environ.get("DOTS_TTS_HOST", "0.0.0.0")
    port = _env_int("DOTS_TTS_PORT", 7860)
    model_name_or_path = os.environ.get(
        "DOTS_TTS_MODEL_NAME_OR_PATH",
        "dots-studio/dots.tts-soar",
    )
    model_revision = os.environ.get("DOTS_TTS_MODEL_REVISION") or None
    precision = os.environ.get("DOTS_TTS_PRECISION", "bfloat16")
    execution_mode = os.environ.get("DOTS_TTS_EXECUTION_MODE", "generate")
    max_generate_length = _env_int("DOTS_TTS_MAX_GENERATE_LENGTH", 500)
    default_num_steps = _env_int("DOTS_TTS_DEFAULT_NUM_STEPS", 16)
    optimize = _env_bool("DOTS_TTS_OPTIMIZE", True)
    enable_aoti = _env_bool("DOTS_TTS_ENABLE_AOTI", True)
    startup_compile = _env_bool("DOTS_TTS_AOTI_COMPILE_ON_STARTUP", True)
    generation_duration = _env_int("DOTS_TTS_ZERO_GPU_DURATION", 60)
    compile_duration = _env_int("DOTS_TTS_ZERO_GPU_COMPILE_DURATION", 1500)
    output_dir = Path(os.environ.get("DOTS_TTS_OUTPUT_DIR", "/data/generated"))
    log_file = Path(os.environ.get("DOTS_TTS_LOG_FILE", "/tmp/dots_tts_gradio.log"))

    configure_logging(log_file=log_file)
    logger.info(
        "Space app starting: model={} execution_mode={} precision={} optimize={} "
        "enable_aoti={} startup_compile={} max_generate_length={}",
        model_name_or_path,
        execution_mode,
        precision,
        optimize,
        enable_aoti,
        startup_compile,
        max_generate_length,
    )

    app_config = build_gradio_app_config(
        host=host,
        port=port,
        execution_mode=execution_mode,
        precision=precision,
        optimize=optimize,
        model_name_or_path=model_name_or_path,
        output_dir=output_dir,
        max_generate_length=max_generate_length,
        default_num_steps=default_num_steps,
        default_max_generate_length=max_generate_length,
        repo_root=REPO_ROOT,
        model_revision=model_revision,
    )
    app_service = GradioAppService(app_config)
    runtime, resolved_model_name_or_path = _preload_runtime(app_service, app_config)

    if enable_aoti and startup_compile and optimize:

        @spaces.GPU(duration=compile_duration)
        def compile_aoti_cache():
            child_runtime, _ = _preload_runtime(app_service, app_config)
            return compile_runtime_aoti(child_runtime)

        compiled_models = compile_aoti_cache()
        installed = AOTI_REGISTRY.install(compiled_models or {})
        logger.info("AOTI startup compile completed: installed_target_count={}", installed)

    app_service.generate = spaces.GPU(duration=generation_duration)(app_service.generate)

    demo = build_demo(gr, app_config, app_service)
    logger.info(
        "Space app ready: host={} port={} resolved_model={} aoti_target_count={}",
        app_config.host,
        app_config.port,
        resolved_model_name_or_path,
        len(AOTI_REGISTRY),
    )
    demo.launch(
        server_name=app_config.host,
        server_port=app_config.port,
        theme=build_playground_theme(gr),
        css=PLAYGROUND_CSS,
    )


if __name__ == "__main__":
    main()