import logging
import os
import time
import traceback
from io import BytesIO

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

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

load_dotenv()

# API Configuration (host + gen_image_path style, aligned with test1.py)
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")

# Polling / retry configuration (with sensible defaults)
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))

# The UI no longer exposes an aspect-ratio selector, so every request goes out
# as 1:1 (matches the main flow used in test1.py).
DEFAULT_WH_RATIO = "1:1"
DEFAULT_SEED = 42

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"
)


class APIError(Exception):
    """Custom exception for API-related errors"""
    pass


# Status codes returned by the API
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 create_request(prompt, wh_ratio, enable_prompt_refine=True, seed=-1):
    """
    Submit an image generation request to the API.

    Args:
        prompt (str): Text prompt describing the image to generate
        wh_ratio (str): Aspect ratio for the output image (e.g. "1:1")
        enable_prompt_refine (bool): Whether to let the backend rewrite/expand
            the prompt before generation. Sent to the API as 0 / 1.
        seed (int): Generation seed. -1 means the backend will pick one
            randomly; any other integer fixes the seed for reproducible runs.

    Returns:
        str: Task ID

    Raises:
        APIError: If the API request fails
    """
    logger.info(
        f"Starting create_request with prompt='{prompt[:50]}...', "
        f"wh_ratio={wh_ratio}, enable_prompt_refine={enable_prompt_refine}, seed={seed}"
    )

    if not prompt or not prompt.strip():
        logger.error("Empty prompt provided to create_request")
        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,
        "wh_ratio": wh_ratio,
        "model_id": MODEL_ID,
        "n": 1,
        "enable_prompt_refine": 1 if enable_prompt_refine else 0,
        "seed": seed_int,
    }

    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:
                logger.error(f"API returned error code {code}: {message}")
                raise APIError(f"Failed to submit task (code={code}): {message}")

            task_id = response_json.get("result", {}).get("task_id")
            if not task_id:
                logger.error(f"No task ID in API response: {response_json}")
                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}"
                logger.error(f"API response error content: {error_detail}")
            except Exception:
                logger.error(f"Could not parse API error response as JSON. Raw content: {e.response.content[:500]}")

            if status_code == 401:
                logger.error(f"Authentication failed with API token. Status code: {status_code}")
                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:
                logger.error(f"Client error: {error_message}, Prompt: '{prompt[:50]}...', Status: {status_code}")
                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:
            logger.error(f"Request error: {str(e)}")
            logger.debug(f"Request error details: {traceback.format_exc()}")
            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)}")

    logger.error(f"Failed to create request after {MAX_RETRY_COUNT} retries for prompt: '{prompt[:50]}...'")
    raise APIError(f"Failed after {MAX_RETRY_COUNT} retries")


def get_results(task_id):
    """
    Check the status of an image generation task.

    Args:
        task_id (str): The task ID to check

    Returns:
        dict: Task result information (the "result" object from the response), or None on transient failure.

    Raises:
        APIError: For unrecoverable errors (e.g. authentication failure).
    """
    logger.debug(f"Checking status for task ID: {task_id}")

    if not task_id:
        logger.error("Empty task ID provided to get_results")
        raise ValueError("Task ID cannot be empty")

    url = _build_result_url(task_id)

    try:
        response = requests.get(url, headers=_headers(), timeout=10)
        logger.debug(f"Status check response code: {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:
            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
        logger.warning(f"HTTP error {status_code} when checking task {task_id}")
        try:
            error_content = e.response.json()
            logger.error(f"Error response content: {error_content}")
        except Exception:
            logger.error(f"Could not parse error response as JSON. Raw content: {e.response.content[:500]}")

        if status_code == 401:
            logger.error(f"Authentication failed when checking task {task_id}")
            raise APIError(f"Authentication failed. Please check your API token when checking task {task_id}")
        elif 400 <= status_code < 500:
            logger.error(f"Client error {status_code} when checking task {task_id}")
            return None
        else:
            logger.warning(f"Server error {status_code} when checking task {task_id}")
            return None

    except requests.exceptions.RequestException as e:
        logger.warning(f"Network error when checking task {task_id}: {str(e)}")
        logger.debug(f"Network error details: {traceback.format_exc()}")
        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):
    """
    Download an image from a URL and return it as a PIL Image.
    Converts non-PNG formats (e.g. WebP) to PNG while preserving original metadata.
    """
    logger.info(f"Starting download_image from URL: {image_url}")

    if not image_url:
        logger.error("Empty image URL provided to download_image")
        raise ValueError("Image URL cannot be empty when downloading image")

    retry_count = 0
    while retry_count < MAX_RETRY_COUNT:
        try:
            logger.info(f"Downloading image [attempt {retry_count + 1}/{MAX_RETRY_COUNT}] from {image_url}")
            response = requests.get(image_url, timeout=30)
            logger.debug(
                f"Image download response status: {response.status_code}, "
                f"Content-Type: {response.headers.get('Content-Type')}, "
                f"Content-Length: {response.headers.get('Content-Length')}"
            )
            response.raise_for_status()

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

            original_metadata = {}
            for key, value in image.info.items():
                if isinstance(key, str) and isinstance(value, str):
                    original_metadata[key] = value
            logger.debug(f"Original image metadata: {original_metadata}")

            if image.format != 'PNG':
                logger.info(f"Converting image from {image.format} to PNG format")
                png_buffer = BytesIO()
                if 'A' in image.getbands():
                    image_to_save = image
                else:
                    image_to_save = image.convert('RGB')
                image_to_save.save(png_buffer, format='PNG')
                png_buffer.seek(0)
                image = Image.open(png_buffer)
                for key, value in original_metadata.items():
                    image.info[key] = value

            logger.info(f"Successfully downloaded and processed image: {image.size[0]}x{image.size[1]}")
            return image

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

        except requests.exceptions.HTTPError as e:
            status_code = e.response.status_code
            logger.error(f"HTTP error {status_code} when downloading image from {image_url}")
            if 400 <= status_code < 500:
                raise APIError(f"HTTP error {status_code} when downloading image")
            else:
                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:
            logger.error(f"Error processing image from {image_url}: {str(e)}")
            logger.error(f"Full traceback: {traceback.format_exc()}")
            raise APIError(f"Failed to process image: {str(e)}")

    logger.error(f"Failed to download image from {image_url} after {MAX_RETRY_COUNT} retries")
    raise APIError(f"Failed to download image after {MAX_RETRY_COUNT} retries")


def text_to_image(prompt: str, seed: int, use_rewrite: bool):
    """Submit a text-to-image request and return the generated PIL image.

    Mirrors the signature used in ref.py (prompt, seed, use_rewrite).
    `use_rewrite` is forwarded to the backend as `enable_prompt_refine`.
    """
    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:
        task_id = create_request(
            prompt,
            DEFAULT_WH_RATIO,
            enable_prompt_refine=bool(use_rewrite),
            seed=seed_int,
        )
    except APIError as exc:
        logger.error(f"API error while submitting task: {exc}")
        raise gr.Error(str(exc))
    except ValueError as exc:
        logger.error(f"Value error while submitting task: {exc}")
        raise gr.Error(str(exc))
    except Exception as exc:
        logger.error(f"Unexpected error while submitting task: {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:
            logger.error(f"Task completed but no sub_task_results returned. Task ID: {task_id}")
            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:
                logger.error(f"No image URL in successful response. Sub result: {sub}")
                raise gr.Error("No image URL in response.")
            logger.info(f"Downloading image - Task ID: {task_id}, URL: {image_url}")
            try:
                image = download_image(image_url)
            except APIError as exc:
                raise gr.Error(f"Failed to download generated image: {exc}")
            if image is None:
                raise gr.Error("Failed to download generated image.")
            return image

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

        time.sleep(POLL_INTERVAL)

    logger.error(f"Timeout waiting for task completion - Task ID: {task_id}, Max time: {MAX_POLL_TIME}s")
    raise gr.Error(f"Timed out after {MAX_POLL_TIME}s waiting for image generation.")


CUSTOM_CSS = """
.page-footer {
    margin-top: 32px;
    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="HiDream-O1-Image-Dev-2604", css=CUSTOM_CSS) as demo:
    gr.Markdown("# HiDream-O1-Image-Dev-2604\nA minimal text-to-image demo.")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                lines=6,
                placeholder="Describe the image you want to generate...",
            )
            seed = gr.Number(
                label="Seed",
                value=DEFAULT_SEED,
                precision=0,
            )
            use_rewrite = gr.Checkbox(
                label="Rewrite prompt before generation",
                value=True,
            )

            run_btn = gr.Button("Generate", variant="primary")

        with gr.Column():
            output_image = gr.Image(label="Output", type="pil")

    gr.HTML(
        """
        <div class="page-footer">
            <div class="footer-links">
                <a href="https://huggingface.co/HiDream-ai/HiDream-O1-Image-Dev-2604" target="_blank">HuggingFace</a>
                <a href="https://github.com/HiDream-ai/HiDream-O1-Image" target="_blank">GitHub</a>
                <a href="https://x.com/vivago_ai" target="_blank">Twitter</a>
            </div>
            <div class="tagline">
                For more features and the full experience, visit
                <a href="https://vivago.ai/" target="_blank">vivago.ai</a>.
            </div>
        </div>
        """
    )

    run_btn.click(
        fn=text_to_image,
        inputs=[prompt, seed, use_rewrite],
        outputs=[output_image],
    )
    prompt.submit(
        fn=text_to_image,
        inputs=[prompt, seed, use_rewrite],
        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")