"""Space entry point. Load both Hub models before serving GPU requests."""

import argparse
import logging
import os

import spaces  # Initialize ZeroGPU before importing PyTorch or Diffusers.

from gradio_app import (
    DEFAULT_MODEL, EXAMPLES, MODEL_IDS, UI_ASSETS,
    ModelRunner, build_app, load_pipeline,
)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--preview", action="store_true", help="Preview UI without downloading weights")
    parser.add_argument("--port", type=int, default=7860)
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO)
    model = os.environ.get("MOSS_MODEL_ID", DEFAULT_MODEL)
    if model not in MODEL_IDS:
        parser.error("MOSS_MODEL_ID must be OpenMOSS-Team/MOSS-Image-Base or OpenMOSS-Team/MOSS-Image-SFT")
    runner = ModelRunner("outputs/gradio")
    if not args.preview:
        shared_pipe = load_pipeline(model)
        runner.pipelines[model] = shared_pipe
        for repo in MODEL_IDS:
            if repo != model:
                runner.pipelines[repo] = load_pipeline(repo, shared_pipe=shared_pipe)
        if os.environ.get("MOSS_AOTI", "1") == "1":
            from aoti import prepare_aoti
            prepare_aoti(runner.pipelines.values())
    demo = build_app(model=model, runner=runner)
    generate_fn = next(fn for fn in demo.fns.values() if fn.api_name == "generate")
    logging.info(
        "Generation queue: SPACES_ZERO_GPU=%r, concurrency_limit=%s",
        os.environ.get("SPACES_ZERO_GPU"), generate_fn.concurrency_limit,
    )
    demo.launch(
        server_name="0.0.0.0", server_port=args.port, share=False,
        allowed_paths=[str(runner.output_dir), str(EXAMPLES.resolve()), str(UI_ASSETS.resolve())],
    )


if __name__ == "__main__":
    main()