import logging
import os
import threading
import time
import traceback
from io import BytesIO
from pathlib import Path

import gradio as gr
import requests
from PIL import Image
from dotenv import load_dotenv

try:
    import spaces
except Exception:
    spaces = None

try:
    from gradio_client import Client as GradioClient
except Exception:
    GradioClient = None

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

load_dotenv()

API_TOKEN = os.environ.get("token")
API_HOST = os.environ.get("host")
GEN_IMAGE_PATH = os.environ.get("gen_image_path")
MODEL_ID = os.environ.get("model_id")

BACKEND_MODE = (os.environ.get("BACKEND_MODE") or os.environ.get("backend_mode") or "auto").strip().lower()
if BACKEND_MODE not in {"auto", "api", "local", "space"}:
    BACKEND_MODE = "auto"

FALLBACK_SPACE_ID = (
    os.environ.get("FALLBACK_SPACE_ID")
    or os.environ.get("fallback_space_id")
    or "HiDream-ai/HiDream-O1-Image-Dev-2604"
)
SPACE_ID = (os.environ.get("SPACE_ID") or "").strip()

LOCAL_MODEL_REPO_ID = "HiDream-ai/HiDream-O1-Image-Dev-2604"
LOCAL_MODEL_DIR = (Path(__file__).resolve().parent / "models" / "gen_model").resolve()


def _resolve_local_model_path() -> str:
    configured = (os.environ.get("LOCAL_MODEL_PATH") or "").strip()
    if configured:
        return configured
    if LOCAL_MODEL_DIR.exists():
        return str(LOCAL_MODEL_DIR)
    return LOCAL_MODEL_REPO_ID


LOCAL_MODEL_PATH = _resolve_local_model_path()
LOCAL_NUM_INFERENCE_STEPS = int(os.environ.get("LOCAL_NUM_INFERENCE_STEPS", "20"))
LOCAL_GUIDANCE_SCALE = float(os.environ.get("LOCAL_GUIDANCE_SCALE", "0.0"))
LOCAL_SHIFT = float(os.environ.get("LOCAL_SHIFT", "1.0"))
LOCAL_SCHEDULER_NAME = os.environ.get("LOCAL_SCHEDULER_NAME", "flow_match")
LOCAL_NOISE_SCALE_START = float(os.environ.get("LOCAL_NOISE_SCALE_START", "8.0"))
LOCAL_NOISE_SCALE_END = float(os.environ.get("LOCAL_NOISE_SCALE_END", "8.0"))
LOCAL_NOISE_CLIP_STD = float(os.environ.get("LOCAL_NOISE_CLIP_STD", "8.0"))
LOCAL_MODEL_DTYPE = os.environ.get("LOCAL_MODEL_DTYPE", "float16").strip().lower()
LOCAL_GPU_DURATION = int(os.environ.get("LOCAL_GPU_DURATION", "300"))
LOCAL_ENABLE_PROMPT_REWRITE = (os.environ.get("LOCAL_ENABLE_PROMPT_REWRITE", "1").strip().lower() in {"1", "true", "yes", "on"})
SPACE_HW_CACHE_TTL = int(os.environ.get("SPACE_HW_CACHE_TTL", "120"))

MAX_RETRY_COUNT = int(os.environ.get("MAX_RETRY_COUNT", 3))
POLL_INTERVAL = float(os.environ.get("POLL_INTERVAL", 2.0))
MAX_POLL_TIME = int(os.environ.get("MAX_POLL_TIME", 300))

DEFAULT_SEED = 42
RESOLUTION_OPTIONS = ["1024x1024", "1536x1536", "2048x2048"]
DEFAULT_RESOLUTION = os.environ.get("DEFAULT_RESOLUTION", "1024x1024")
if DEFAULT_RESOLUTION not in RESOLUTION_OPTIONS:
    DEFAULT_RESOLUTION = "1024x1024"

logger.info(
    f"API configuration loaded: HOST={API_HOST}, GEN_IMAGE_PATH={GEN_IMAGE_PATH}, MODEL_ID={MODEL_ID}"
)
logger.info(
    f"Retry configuration: MAX_RETRY_COUNT={MAX_RETRY_COUNT}, POLL_INTERVAL={POLL_INTERVAL}s, MAX_POLL_TIME={MAX_POLL_TIME}s"
)
logger.info(f"Backend mode: {BACKEND_MODE}, fallback_space_id: {FALLBACK_SPACE_ID}")
logger.info(f"Space id: {SPACE_ID or 'unknown'}, local model path: {LOCAL_MODEL_PATH}")

_SPACE_CLIENT = None
_LOCAL_RUNTIME = None
_LOCAL_RUNTIME_LOCK = threading.Lock()
_SPACE_HW_CACHE = {
    "ts": 0.0,
    "current": None,
    "requested": None,
}


class APIError(Exception):
    pass


SUCCESS_CODE = 0


def _build_request_url() -> str:
    if not API_HOST or not GEN_IMAGE_PATH:
        raise APIError(
            "API host or gen_image_path is not configured. "
            "Please set the 'host' and 'gen_image_path' environment variables."
        )
    return f"{API_HOST.rstrip('/')}{GEN_IMAGE_PATH}"


def _build_result_url(task_id: str) -> str:
    return f"{_build_request_url()}/results?task_id={task_id}"


def _headers() -> dict:
    if not API_TOKEN:
        raise APIError("API token is not configured. Please set the 'token' environment variable.")
    return {"Authorization": f"Bearer {API_TOKEN}"}


def _is_gpu_hardware(flavor: str | None) -> bool:
    if not flavor:
        return False
    return not str(flavor).lower().startswith("cpu")


def _fetch_space_hardware() -> tuple[str | None, str | None]:
    now = time.time()
    if now - _SPACE_HW_CACHE["ts"] <= SPACE_HW_CACHE_TTL:
        return _SPACE_HW_CACHE["current"], _SPACE_HW_CACHE["requested"]

    current = None
    requested = None
    if SPACE_ID:
        try:
            response = requests.get(
                f"https://huggingface.co/api/spaces/{SPACE_ID}",
                timeout=8,
            )
            response.raise_for_status()
            payload = response.json()
            runtime = payload.get("runtime") or {}
            hardware = runtime.get("hardware") or {}
            current = hardware.get("current")
            requested = hardware.get("requested")
        except Exception as exc:
            logger.warning(f"Could not resolve space hardware via API: {exc}")

    _SPACE_HW_CACHE["ts"] = now
    _SPACE_HW_CACHE["current"] = current
    _SPACE_HW_CACHE["requested"] = requested
    return current, requested


def _space_has_gpu_hardware() -> bool:
    current, requested = _fetch_space_hardware()
    return _is_gpu_hardware(current) or _is_gpu_hardware(requested)


def _resolve_backend_mode() -> str:
    if BACKEND_MODE in {"api", "local", "space"}:
        return BACKEND_MODE
    if API_TOKEN and API_HOST and GEN_IMAGE_PATH:
        return "api"
    if _space_has_gpu_hardware():
        return "local"
    return "space"


def _resolve_local_torch_dtype(torch_module):
    dtype_map = {
        "float16": torch_module.float16,
        "fp16": torch_module.float16,
        "bfloat16": torch_module.bfloat16,
        "bf16": torch_module.bfloat16,
        "float32": torch_module.float32,
        "fp32": torch_module.float32,
    }
    return dtype_map.get(LOCAL_MODEL_DTYPE, torch_module.float16)


def _add_local_special_tokens(tokenizer) -> None:
    tokenizer.boi_token = "<|boi_token|>"
    tokenizer.bor_token = "<|bor_token|>"
    tokenizer.eor_token = "<|eor_token|>"
    tokenizer.bot_token = "<|bot_token|>"
    tokenizer.tms_token = "<|tms_token|>"


def _parse_resolution(resolution: str) -> tuple[int, int]:
    try:
        width_str, height_str = resolution.lower().split("x")
        width = int(width_str)
        height = int(height_str)
        if width <= 0 or height <= 0:
            raise ValueError
        return width, height
    except Exception:
        raise ValueError(f"Invalid resolution value: {resolution}")


def _ensure_local_runtime():
    global _LOCAL_RUNTIME

    with _LOCAL_RUNTIME_LOCK:
        if _LOCAL_RUNTIME is not None:
            return _LOCAL_RUNTIME

        try:
            import torch
            from transformers import AutoProcessor, PreTrainedTokenizerBase
            from models.qwen3_vl_transformers import Qwen3VLForConditionalGeneration
            from models.pipeline import generate_image, DEFAULT_TIMESTEPS
        except Exception as exc:
            raise APIError(f"Local backend dependencies are unavailable: {exc}") from exc

        if not torch.cuda.is_available():
            raise APIError(
                "Local backend requires CUDA GPU hardware. "
                "Switch Space hardware to ZeroGPU or a paid GPU tier."
            )

        logger.info(
            f"Initializing local backend model from {LOCAL_MODEL_PATH} "
            f"(dtype={LOCAL_MODEL_DTYPE}, steps={LOCAL_NUM_INFERENCE_STEPS})"
        )
        processor = AutoProcessor.from_pretrained(LOCAL_MODEL_PATH)
        model = (
            Qwen3VLForConditionalGeneration.from_pretrained(
                LOCAL_MODEL_PATH,
                torch_dtype=_resolve_local_torch_dtype(torch),
            )
            .eval()
            .to("cuda")
        )

        tokenizer = processor if isinstance(processor, PreTrainedTokenizerBase) else processor.tokenizer
        _add_local_special_tokens(tokenizer)

        _LOCAL_RUNTIME = {
            "model": model,
            "processor": processor,
            "generate_image": generate_image,
            "default_timesteps": DEFAULT_TIMESTEPS,
        }
        return _LOCAL_RUNTIME


def _maybe_refine_prompt_for_local(prompt: str, use_rewrite: bool) -> str:
    if not use_rewrite:
        return prompt
    if not LOCAL_ENABLE_PROMPT_REWRITE:
        return prompt
    return prompt


def _run_local_generation_impl(prompt: str, seed: int, width: int, height: int):
    runtime = _ensure_local_runtime()
    return runtime["generate_image"](
        model=runtime["model"],
        processor=runtime["processor"],
        prompt=prompt,
        ref_image_paths=[],
        height=height,
        width=width,
        num_inference_steps=LOCAL_NUM_INFERENCE_STEPS,
        guidance_scale=LOCAL_GUIDANCE_SCALE,
        shift=LOCAL_SHIFT,
        timesteps_list=runtime["default_timesteps"],
        scheduler_name=LOCAL_SCHEDULER_NAME,
        seed=int(seed),
        keep_original_aspect=False,
        layout_bboxes=None,
        noise_scale_start=LOCAL_NOISE_SCALE_START,
        noise_scale_end=LOCAL_NOISE_SCALE_END,
        noise_clip_std=LOCAL_NOISE_CLIP_STD,
    )


if spaces is not None:
    @spaces.GPU(duration=LOCAL_GPU_DURATION)
    def _run_local_generation(prompt: str, seed: int, width: int, height: int):
        return _run_local_generation_impl(prompt, seed, width, height)
else:
    def _run_local_generation(prompt: str, seed: int, width: int, height: int):
        return _run_local_generation_impl(prompt, seed, width, height)


def create_request_via_local_backend(prompt: str, seed: int, use_rewrite: bool, width: int, height: int) -> Image.Image:
    logger.info("Using local backend for generation")
    final_prompt = _maybe_refine_prompt_for_local(prompt, use_rewrite)
    try:
        return _run_local_generation(final_prompt, int(seed), int(width), int(height))
    except Exception as exc:
        raise APIError(f"Local backend generation failed: {exc}") from exc


def _get_space_client():
    global _SPACE_CLIENT

    if not FALLBACK_SPACE_ID:
        raise APIError("Fallback space id is empty. Set FALLBACK_SPACE_ID.")

    if SPACE_ID and FALLBACK_SPACE_ID.strip().lower() == SPACE_ID.lower():
        raise APIError("Fallback space id cannot be the current space id.")

    if GradioClient is None:
        raise APIError(
            "gradio_client is not available. Please add it to requirements or set private API secrets."
        )

    if _SPACE_CLIENT is None:
        logger.info(f"Initializing fallback Gradio client for: {FALLBACK_SPACE_ID}")
        _SPACE_CLIENT = GradioClient(FALLBACK_SPACE_ID)

    return _SPACE_CLIENT


def _open_local_image(image_path: str) -> Image.Image:
    if not image_path:
        raise APIError("Fallback backend did not return an image path.")

    path_obj = Path(image_path)
    if not path_obj.exists():
        raise APIError(f"Fallback backend returned missing file path: {image_path}")

    try:
        with Image.open(path_obj) as image:
            return image.copy()
    except Exception as exc:
        raise APIError(f"Failed to open fallback image: {exc}") from exc


def _convert_fallback_result_to_image(result) -> Image.Image:
    if isinstance(result, str):
        return _open_local_image(result)

    if isinstance(result, dict):
        if result.get("path"):
            return _open_local_image(result["path"])
        if result.get("url"):
            return download_image(result["url"])

    if isinstance(result, (list, tuple)) and result:
        return _convert_fallback_result_to_image(result[0])

    raise APIError(f"Unsupported fallback result format: {type(result)}")


def create_request_via_fallback_space(prompt: str, seed: int, use_rewrite: bool) -> Image.Image:
    logger.info(f"Using fallback space backend: {FALLBACK_SPACE_ID}")
    client = _get_space_client()

    try:
        result = client.predict(
            prompt=prompt,
            seed=seed,
            use_rewrite=bool(use_rewrite),
            api_name="/text_to_image",
        )
    except Exception as exc:
        raise APIError(f"Fallback space request failed: {exc}") from exc

    return _convert_fallback_result_to_image(result)


def create_request(prompt, width, height, enable_prompt_refine=True, seed=-1):
    logger.info(
        f"Starting create_request with prompt='{prompt[:50]}...', width={width}, height={height}, "
        f"enable_prompt_refine={enable_prompt_refine}, seed={seed}"
    )

    if not prompt or not prompt.strip():
        raise ValueError("Prompt cannot be empty")

    try:
        seed_int = int(seed)
    except (TypeError, ValueError):
        logger.warning(f"Invalid seed value '{seed}', falling back to -1 (random)")
        seed_int = -1

    model_params = {
        "prompt": prompt,
        "model_id": MODEL_ID,
        "n": 1,
        "enable_prompt_refine": 1 if enable_prompt_refine else 0,
        "seed": seed_int,
        "width": int(width),
        "height": int(height),
    }

    url = _build_request_url()

    retry_count = 0
    while retry_count < MAX_RETRY_COUNT:
        try:
            logger.info(
                f"Sending API request [attempt {retry_count + 1}/{MAX_RETRY_COUNT}] for prompt: '{prompt[:50]}...'"
            )
            response = requests.post(url, json=model_params, headers=_headers(), timeout=15)
            logger.info(f"API request response status: {response.status_code}")
            response.raise_for_status()

            response_json = response.json()
            code = response_json.get("code")
            message = response_json.get("message", "")
            if code != SUCCESS_CODE:
                raise APIError(f"Failed to submit task (code={code}): {message}")

            task_id = response_json.get("result", {}).get("task_id")
            if not task_id:
                raise APIError(f"No task ID returned from API: {response_json}")

            logger.info(f"Successfully created task with ID: {task_id}")
            return task_id

        except requests.exceptions.Timeout:
            retry_count += 1
            logger.warning(f"Request timed out. Retrying ({retry_count}/{MAX_RETRY_COUNT})...")
            time.sleep(1)

        except requests.exceptions.HTTPError as e:
            status_code = e.response.status_code
            error_message = f"HTTP error {status_code}"
            try:
                error_detail = e.response.json()
                error_message += f": {error_detail}"
            except Exception:
                pass

            if status_code == 401:
                raise APIError("Authentication failed. Please check your API token.")
            elif status_code == 429:
                retry_count += 1
                wait_time = min(2 ** retry_count, 10)
                logger.warning(f"Rate limit exceeded. Waiting {wait_time}s before retry ({retry_count}/{MAX_RETRY_COUNT})...")
                time.sleep(wait_time)
            elif 400 <= status_code < 500:
                raise APIError(error_message)
            else:
                retry_count += 1
                logger.warning(f"Server error: {error_message}. Retrying ({retry_count}/{MAX_RETRY_COUNT})...")
                time.sleep(1)

        except requests.exceptions.RequestException as e:
            raise APIError(f"Failed to connect to API: {str(e)}")

        except APIError:
            raise

        except Exception as e:
            logger.error(f"Unexpected error in create_request: {str(e)}")
            logger.error(f"Full traceback: {traceback.format_exc()}")
            raise APIError(f"Unexpected error: {str(e)}")

    raise APIError(f"Failed after {MAX_RETRY_COUNT} retries")


def get_results(task_id):
    if not task_id:
        raise ValueError("Task ID cannot be empty")

    url = _build_result_url(task_id)

    try:
        response = requests.get(url, headers=_headers(), timeout=10)
        response.raise_for_status()
        response_json = response.json()

        code = response_json.get("code")
        message = response_json.get("message", "")
        if code != SUCCESS_CODE:
            logger.warning(f"API returned non-success code {code} for task {task_id}: {message}")
            return None

        return response_json.get("result")

    except requests.exceptions.Timeout:
        logger.warning(f"Request timed out when checking task {task_id}")
        return None
    except requests.exceptions.HTTPError as e:
        status_code = e.response.status_code
        if status_code == 401:
            raise APIError(f"Authentication failed. Please check your API token when checking task {task_id}")
        return None
    except requests.exceptions.RequestException:
        return None
    except Exception as e:
        logger.error(f"Unexpected error when checking task {task_id}: {str(e)}")
        logger.error(f"Full traceback: {traceback.format_exc()}")
        return None


def download_image(image_url):
    logger.info(f"Starting download_image from URL: {image_url}")

    if not image_url:
        raise ValueError("Image URL cannot be empty when downloading image")

    retry_count = 0
    while retry_count < MAX_RETRY_COUNT:
        try:
            response = requests.get(image_url, timeout=30)
            response.raise_for_status()

            image = Image.open(BytesIO(response.content))
            logger.info(f"Image opened successfully. Format: {image.format}, Size: {image.size}, Mode: {image.mode}")

            if image.format != "PNG":
                png_buffer = BytesIO()
                image_to_save = image if "A" in image.getbands() else image.convert("RGB")
                image_to_save.save(png_buffer, format="PNG")
                png_buffer.seek(0)
                image = Image.open(png_buffer)

            return image

        except requests.exceptions.Timeout:
            retry_count += 1
            time.sleep(1)
        except requests.exceptions.RequestException as e:
            retry_count += 1
            logger.warning(f"Network error during image download: {str(e)}. Retrying ({retry_count}/{MAX_RETRY_COUNT})...")
            time.sleep(1)
        except Exception as e:
            raise APIError(f"Failed to process image: {str(e)}")

    raise APIError(f"Failed to download image after {MAX_RETRY_COUNT} retries")


def text_to_image(prompt: str, seed: int, use_rewrite: bool, resolution: str):
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    try:
        seed_int = int(seed) if seed is not None else -1
    except (TypeError, ValueError):
        seed_int = -1

    try:
        width, height = _parse_resolution(resolution)
    except ValueError as exc:
        raise gr.Error(str(exc))

    selected_backend = _resolve_backend_mode()
    logger.info(f"Selected backend: {selected_backend}, resolution={resolution}")

    if selected_backend == "local":
        try:
            return create_request_via_local_backend(
                prompt=prompt,
                seed=seed_int,
                use_rewrite=bool(use_rewrite),
                width=width,
                height=height,
            )
        except APIError as exc:
            raise gr.Error(str(exc))

    if selected_backend == "space":
        try:
            return create_request_via_fallback_space(
                prompt=prompt,
                seed=seed_int,
                use_rewrite=bool(use_rewrite),
            )
        except APIError as exc:
            raise gr.Error(str(exc))

    try:
        task_id = create_request(
            prompt,
            width,
            height,
            enable_prompt_refine=bool(use_rewrite),
            seed=seed_int,
        )
    except APIError as exc:
        raise gr.Error(str(exc))
    except ValueError as exc:
        raise gr.Error(str(exc))
    except Exception as exc:
        logger.error(traceback.format_exc())
        raise gr.Error(f"Unexpected error: {exc}")

    start_time = time.time()
    logger.info(f"Polling for results - Task ID: {task_id}")

    while time.time() - start_time < MAX_POLL_TIME:
        result = get_results(task_id)
        if not result:
            time.sleep(POLL_INTERVAL)
            continue

        overall_status = result.get("status")
        sub_results = result.get("sub_task_results", []) or []

        if overall_status != 1:
            time.sleep(POLL_INTERVAL)
            continue

        if not sub_results:
            raise gr.Error("Task completed but no results returned.")

        sub = sub_results[0]
        sub_status = sub.get("task_status")

        if sub_status == 1:
            image_url = sub.get("url")
            if not image_url:
                raise gr.Error("No image URL in response.")
            try:
                image = download_image(image_url)
            except APIError as exc:
                raise gr.Error(f"Failed to download generated image: {exc}")
            return image

        if sub_status == 3:
            error_msg = sub.get("task_error") or sub.get("message") or "Unknown error"
            raise gr.Error(f"Task failed: {error_msg}")

        time.sleep(POLL_INTERVAL)

    raise gr.Error(f"Timed out after {MAX_POLL_TIME}s waiting for image generation.")


CUSTOM_CSS = """
.page-footer {
    margin-top: 0px;
    padding: 20px 0 8px 0;
    border-top: 1px solid var(--border-color-primary, #e5e7eb);
    text-align: center;
}
.page-footer .footer-links a {
    margin: 0 12px;
    text-decoration: none;
    font-weight: 500;
}
.page-footer .tagline {
    margin-top: 8px;
    font-size: 0.9em;
    opacity: 0.75;
}
"""


with gr.Blocks(title="", css=CUSTOM_CSS) as demo:
    gr.Markdown("")
    with gr.Column():
        with gr.Column():
            prompt = gr.Textbox(
                label="",
                lines=3,
                placeholder="",
            )
            seed = gr.Number(
                label="",
                value=0,
                precision=0,
            )
            use_rewrite = gr.Checkbox(
                label="",
                value=True,
            )
            resolution = gr.Dropdown(
                choices=RESOLUTION_OPTIONS,
                value=DEFAULT_RESOLUTION,
                label="",
            )
            run_btn = gr.Button("", variant="primary")

        with gr.Column():
            output_image = gr.Image(label="", type="pil", format="png", height=1200)

    gr.HTML(""" """)

    run_btn.click(
        fn=text_to_image,
        inputs=[prompt, seed, use_rewrite, resolution],
        outputs=[output_image],
    )
    prompt.submit(
        fn=text_to_image,
        inputs=[prompt, seed, use_rewrite, resolution],
        outputs=[output_image],
    )


if __name__ == "__main__":
    logger.info("Starting HiDream-O1-Image-Dev Generator application")
    demo.queue(max_size=50, default_concurrency_limit=4).launch()
    logger.info("Application shutdown")