from __future__ import annotations

import torch
import spaces

import atexit
import copy
import os
import random
import sys
import tempfile
import threading
import warnings
from pathlib import Path
from typing import Optional

import gradio as gr
import numpy as np
from PIL import Image
from huggingface_hub import snapshot_download


ROOT_DIR = Path(__file__).resolve().parent
SRC_DIR = ROOT_DIR / "src"
if str(SRC_DIR) not in sys.path:
    sys.path.insert(0, str(SRC_DIR))

from infer_runtime.model import InferenceParams, build_model
from infer_runtime.settings import load_settings
from modules.models.attention import describe_attention_backend
from modules.utils import clean_dist_env, maybe_init_distributed

warnings.filterwarnings("ignore")

for env_name, env_value in {
    "HF_HUB_DISABLE_TELEMETRY": "1",
    "GRADIO_ANALYTICS_ENABLED": "False",
    "WANDB_DISABLED": "true",
    "DO_NOT_TRACK": "1",
    "TOKENIZERS_PARALLELISM": "false",
}.items():
    os.environ.setdefault(env_name, env_value)

TRUE_VALUES = {"1", "true", "yes", "y"}
MAX_SEED = np.iinfo(np.int32).max
DEFAULT_NEGATIVE_PROMPT = ""
GPU_PRELOAD_MODES = {
    "startup",
    "startup_preload",
    "boot",
    "auto",
    "eager",
    "preload",
    "gpu",
    "gpu_preload",
    "cuda",
    "global_cuda",
}


def env_flag(name: str, default: str = "0") -> bool:
    return os.getenv(name, default).strip().lower() in TRUE_VALUES


def env_optional_int(name: str) -> int | None:
    value = os.getenv(name, "").strip()
    return int(value) if value else None


def env_optional_str(name: str) -> str | None:
    value = os.getenv(name, "").strip()
    return value or None


def is_rank0() -> bool:
    return int(os.environ.get("RANK", "0")) == 0


def resolve_ckpt_root(model_repo_id: str, explicit_ckpt_root: str | None, hf_token: str | None) -> str:
    if explicit_ckpt_root:
        return explicit_ckpt_root
    return snapshot_download(repo_id=model_repo_id, token=hf_token)


def normalize_input_image(image: Image.Image | str | dict | None) -> Image.Image | None:
    if image is None:
        return None
    if isinstance(image, Image.Image):
        return image.convert("RGB")
    if isinstance(image, str):
        return Image.open(image).convert("RGB")
    if isinstance(image, dict):
        for key in ("image", "background"):
            value = image.get(key)
            if value is None:
                continue
            if isinstance(value, Image.Image):
                return value.convert("RGB")
            if isinstance(value, str):
                return Image.open(value).convert("RGB")
    if hasattr(image, "name"):
        return Image.open(image.name).convert("RGB")
    raise TypeError(f"Unsupported image input type: {type(image)!r}")


class JoyImageEditApp:
    def __init__(
        self,
        ckpt_root: str,
        config_path: Optional[str] = None,
        *,
        rewrite_model: str = "qwen-vl-max-latest",
        hsdp_shard_dim: Optional[int] = None,
        enable_prompt_rewrite: bool = False,
        basesize: int = 1024,
        device: Optional[torch.device] = None,
        model_load_mode: str = "startup_preload",
    ):
        self.ckpt_root = str(ckpt_root)
        self.config_path = str(config_path) if config_path else None
        self.rewrite_model = str(rewrite_model)
        self.hsdp_shard_dim = hsdp_shard_dim
        self.enable_prompt_rewrite = bool(enable_prompt_rewrite)
        self.basesize = int(basesize)
        self.device = device
        self.model_load_mode = str(model_load_mode or "startup_preload").strip().lower()

        self._infer_lock = threading.Lock()
        self._model_lock = threading.Lock()
        self._scheduler_template = None
        self._preload_attempted = False
        self._preload_error = None
        self.model = None

        self.settings = load_settings(
            ckpt_root=self.ckpt_root,
            config_path=self.config_path,
            rewrite_model=self.rewrite_model,
            default_seed=0,
        )

    def _ensure_model_loaded(self) -> None:
        desired_device = torch.device("cuda")

        if self.model is not None:
            self.device = desired_device
            return

        with self._model_lock:
            if self.model is None:
                if is_rank0():
                    print(f"[Model] Building model on device: {desired_device}")

                self.model = build_model(
                    self.settings,
                    device=desired_device,
                    hsdp_shard_dim_override=self.hsdp_shard_dim,
                )
                self.device = desired_device

                try:
                    pipeline = getattr(self.model, "pipeline", None)
                    scheduler = getattr(pipeline, "scheduler", None)
                    if scheduler is not None:
                        self._scheduler_template = copy.deepcopy(scheduler)
                except Exception as exc:
                    if is_rank0():
                        print(f"[Model] Warning: failed to snapshot scheduler template: {exc}")
                    self._scheduler_template = None

                if is_rank0():
                    print(f"[Model] Device: {self.device}")
                    print("[Model] Model load complete.")

    def maybe_preload_model(self) -> None:
        if self._preload_attempted:
            return
        self._preload_attempted = True

        if self.model_load_mode in {"lazy", "disabled", "off", "false", "0"}:
            if is_rank0():
                print("[Model] Startup preload disabled; using runtime loading.")
            return

        preload_device = torch.device("cuda")

        if is_rank0():
            print(f"[Model] Attempting startup preload on {preload_device}...")
        try:
            self._ensure_model_loaded()
            self._preload_error = None
            if is_rank0():
                print(f"[Model] Startup preload succeeded on {preload_device}.")
        except Exception as exc:
            self._preload_error = str(exc)
            self.model = None
            self._scheduler_template = None
            self.device = None
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            if is_rank0():
                print(f"[Model] Startup preload failed; falling back to lazy loading: {exc}")

    def infer(
        self,
        *,
        image: Image.Image | None,
        prompt: str,
        seed: int,
        guidance_scale: float,
        steps: int,
        rewrite_prompt: bool,
        dash_api_key: str | None,
        width: int | None,
        height: int | None,
        negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
    ) -> tuple[Image.Image, str]:
        prompt = str(prompt or "").strip()
        if not prompt:
            raise gr.Error("Please enter an editing instruction.")

        input_image = normalize_input_image(image)
        with self._infer_lock:
            self._ensure_model_loaded()
            model = self.model
            if model is None:
                raise gr.Error("The model has not loaded successfully yet. Please check the logs.")

            effective_prompt = model.maybe_rewrite_prompt(
                prompt,
                input_image,
                enabled=bool(rewrite_prompt),
                dash_api_key=dash_api_key,
            )

            final_height = int(height) if height else self.basesize
            final_width = int(width) if width else self.basesize

            output_image = model.infer(
                InferenceParams(
                    prompt=effective_prompt,
                    image=input_image,
                    height=final_height,
                    width=final_width,
                    steps=int(steps),
                    guidance_scale=float(guidance_scale),
                    seed=int(seed),
                    neg_prompt=str(negative_prompt or ""),
                    basesize=int(self.basesize),
                )
            )

            return output_image, effective_prompt


def build_app() -> tuple[JoyImageEditApp, dict[str, object]]:
    model_repo_id = os.getenv("MODEL_REPO_ID", "jdopensource/JoyAI-Image-Edit")
    ckpt_root_env = env_optional_str("CKPT_ROOT")
    config_path = env_optional_str("CONFIG_PATH")
    rewrite_prompt = env_flag("REWRITE_PROMPT")
    rewrite_model = os.getenv("REWRITE_MODEL", "qwen-vl-max-latest")
    basesize = int(os.getenv("BASESIZE", "1024"))
    hsdp_shard_dim = env_optional_int("HSDP_SHARD_DIM")
    model_load_mode = os.getenv("MODEL_LOAD_MODE", "startup_preload").strip().lower()
    hf_token = env_optional_str("HF_TOKEN") or env_optional_str("HUGGING_FACE_HUB_TOKEN")

    ckpt_root = resolve_ckpt_root(model_repo_id, ckpt_root_env, hf_token)
    app = JoyImageEditApp(
        ckpt_root=ckpt_root,
        config_path=config_path,
        rewrite_model=rewrite_model,
        hsdp_shard_dim=hsdp_shard_dim,
        enable_prompt_rewrite=rewrite_prompt,
        basesize=basesize,
        device=torch.device("cuda"),
        model_load_mode=model_load_mode,
    )
    meta = {
        "model_repo_id": model_repo_id,
        "ckpt_root": ckpt_root,
        "config_path": config_path,
        "rewrite_prompt": rewrite_prompt,
        "rewrite_model": rewrite_model,
        "basesize": basesize,
        "hsdp_shard_dim": hsdp_shard_dim,
        "model_load_mode": model_load_mode,
    }
    return app, meta


def maybe_preload(app: JoyImageEditApp) -> None:
    mode = (app.model_load_mode or "").strip().lower()
    if mode in GPU_PRELOAD_MODES:
        app.maybe_preload_model()
    elif is_rank0():
        print(f"[Model] Using runtime loading mode: {mode}")



def print_startup_info(meta: dict[str, object]) -> None:
    if not is_rank0():
        return
    print(f"[Info] Attention backend: {describe_attention_backend()}")
    print(f"[Info] MODEL_REPO_ID: {meta['model_repo_id']}")
    print(f"[Info] CKPT_ROOT: {meta['ckpt_root']}")
    print(f"[Info] CONFIG_PATH: {meta['config_path'] or '(auto)'}")
    print(f"[Info] REWRITE_PROMPT(default): {meta['rewrite_prompt']}")
    print(f"[Info] REWRITE_MODEL: {meta['rewrite_model']}")
    print(f"[Info] BASESIZE: {meta['basesize']}")
    print(f"[Info] MODEL_LOAD_MODE: {meta['model_load_mode']}")


DIST_INITIALIZED = maybe_init_distributed()
if DIST_INITIALIZED:
    atexit.register(clean_dist_env)

APP, APP_META = build_app()
print_startup_info(APP_META)
maybe_preload(APP)


@spaces.GPU(duration=40)
def run_inference(
    image,
    prompt,
    seed,
    randomize_seed,
    guidance_scale,
    num_inference_steps,
    height,
    width,
    rewrite_prompt,
    dash_api_key,
    progress=gr.Progress(track_tqdm=True),
):
    del progress
    if randomize_seed:
        seed = random.randint(0, int(MAX_SEED))

    output_image, effective_prompt = APP.infer(
        image=image,
        prompt=prompt,
        seed=int(seed),
        guidance_scale=float(guidance_scale),
        steps=int(num_inference_steps),
        rewrite_prompt=bool(rewrite_prompt),
        dash_api_key=str(dash_api_key or "").strip(),
        width=int(width) if width not in (None, 0) else None,
        height=int(height) if height not in (None, 0) else None,
    )
    return output_image, int(seed), effective_prompt


CSS = """
#col-container {
    margin: 0 auto;
    max-width: 1100px;
}
#title {
    text-align: center;
    margin-bottom: 10px;
}
#subtitle {
    text-align: center;
    margin-top: -6px;
    margin-bottom: 18px;
    color: var(--body-text-color-subdued);
}
"""


with gr.Blocks(css=CSS, title="JoyAI-Image-Edit Demo") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""
        <div id="title">
          <h1>JoyAI-Image-Edit</h1>
        </div>
        """)
        gr.Markdown(
            "Model weights are available on [Model](https://huggingface.co/jdopensource/JoyAI-Image-Edit), and project code plus additional information are available on [GitHub](https://github.com/jd-opensource/JoyAI-Image)."
        )

        with gr.Row():
            with gr.Column():
                input_image = gr.Image(
                    label="Input Image",
                    type="pil",
                    image_mode="RGB",
                    sources=["upload", "clipboard"],
                    height=420,
                )
            with gr.Column():
                result = gr.Image(
                    label="Result",
                    type="pil",
                    image_mode="RGB",
                    height=420,
                )

        with gr.Row():
            prompt = gr.Textbox(
                label="Prompt",
                show_label=False,
                placeholder="Describe the edit instruction. Leave the image empty to try text-to-image.",
                lines=2,
            )
            run_button = gr.Button("Edit!", variant="primary")

        with gr.Accordion("Advanced Settings", open=False):
            seed = gr.Slider(label="Seed", minimum=0, maximum=int(MAX_SEED), step=1, value=0)
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            with gr.Row():
                guidance_scale = gr.Slider(label="Guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=4.0)
                num_inference_steps = gr.Slider(label="Number of inference steps", minimum=1, maximum=60, step=1, value=30)
            with gr.Row():
                height = gr.Slider(label="Height (for T2I)", minimum=256, maximum=2048, step=8, value=1024)
                width = gr.Slider(label="Width (for T2I)", minimum=256, maximum=2048, step=8, value=1024)
            with gr.Row():
                rewrite_prompt = gr.Checkbox(
                    label="Rewrite prompt",
                    value=True,
                    interactive=True,
                    info="Use DashScope Qwen-VL to rewrite the edit instruction before inference.",
                )
                dash_api_key = gr.Textbox(
                    label="DASH_API_KEY",
                    type="password",
                    placeholder="Enter DASH_API_KEY for prompt rewrite",
                    value=os.getenv("DASH_API_KEY", ""),
                    container=True,
                    scale=2,
                )
            final_prompt = gr.Textbox(label="Final prompt used", lines=3, interactive=False)

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=run_inference,
        inputs=[
            input_image,
            prompt,
            seed,
            randomize_seed,
            guidance_scale,
            num_inference_steps,
            height,
            width,
            rewrite_prompt,
            dash_api_key,
        ],
        outputs=[result, seed, final_prompt],
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=20).launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        allowed_paths=[
            str(Path(tempfile.gettempdir()).resolve()),
            str((ROOT_DIR / "test_images").resolve()),
            str((ROOT_DIR / "images").resolve()),
        ],
    )