"""
Stable Diffusion 3.5 worker implementation for image generation inference.

This module provides the SD35Worker class that handles image generation jobs
using the Stable Diffusion 3.5 model.
"""

import logging
import os
import tempfile
from typing import Any, Dict, Optional

import requests
import torch
from diffusers import StableDiffusionPipeline

from auror.core.worker import Worker

logger = logging.getLogger(__name__)


class SD35Worker(Worker):
    """
    Worker for processing Stable Diffusion 3.5 image generation jobs.

    This worker handles the complete pipeline for image generation:
    1. Loads the Stable Diffusion 3.5 model on initialization
    2. Generates images based on text prompts
    3. Uploads the generated image to the provided presigned URL
    """

    def __init__(
        self,
        model_id: str = "stabilityai/stable-diffusion-3.5-large",
        torch_dtype: torch.dtype = torch.float16,
        device: str = "cuda",
    ):
        """
        Initialize the SD35 worker and load the model.

        Args:
            model_id: HuggingFace model identifier for SD 3.5
            torch_dtype: Torch data type for model weights (default: float16)
            device: Device to run inference on (default: cuda)
        """
        super().__init__()
        self.model_id = model_id
        self.torch_dtype = torch_dtype
        self.device = device
        self.pipe: Optional[StableDiffusionPipeline] = None

        logger.info(f"Initializing SD35Worker with model: {model_id}")
        self._load_pipeline()

    def _load_pipeline(self) -> None:
        """
        Load the Stable Diffusion pipeline.

        Raises:
            Exception: If pipeline loading fails
        """
        try:
            logger.info("Loading Stable Diffusion 3.5 pipeline...")
            self.pipe = StableDiffusionPipeline.from_pretrained(
                self.model_id, torch_dtype=self.torch_dtype
            )
            self.pipe = self.pipe.to(self.device)
            logger.info("Pipeline loaded successfully")
        except Exception as e:
            logger.error(f"Failed to load pipeline: {e}", exc_info=True)
            raise

    def process(self, job: Dict[str, Any]) -> None:
        """
        Process an image generation job.

        Args:
            job: Job dictionary containing:
                - id: Unique job identifier
                - input: Input parameters including:
                    - prompt: Text prompt for image generation
                    - uploadUrl: Presigned URL for uploading the result
                    - height: Output image height (optional, default: 1024)
                    - width: Output image width (optional, default: 1024)
                    - numInferenceSteps: Number of denoising steps (optional, default: 50)
                    - seed: Random seed for reproducibility (optional)

        Raises:
            Exception: If job processing fails at any stage
        """
        job_id = job["id"]
        output_path = None

        try:
            logger.info(f"[{job_id}] Starting SD 3.5 job")

            inp = job["input"]
            prompt = inp["prompt"]
            upload_url = inp["uploadUrl"]
            height = inp.get("height", 1024)
            width = inp.get("width", 1024)
            num_inference_steps = inp.get("numInferenceSteps", 50)
            seed = inp.get("seed")

            # Set seed for reproducibility if provided
            if seed is not None:
                generator = torch.Generator(device=self.device).manual_seed(seed)
                logger.info(f"[{job_id}] Using seed: {seed}")
            else:
                generator = None

            # Generate image
            logger.info(
                f"[{job_id}] Generating image with prompt: '{prompt[:50]}...'"
            )
            logger.info(
                f"[{job_id}] Parameters: {width}x{height}, {num_inference_steps} steps"
            )

            image = self.pipe(
                prompt,
                height=height,
                width=width,
                num_inference_steps=num_inference_steps,
                generator=generator,
            ).images[0]

            # Save to temporary file
            fd, output_path = tempfile.mkstemp(suffix=".png", prefix=f"sd35_{job_id}_")
            os.close(fd)
            image.save(output_path)
            logger.info(f"[{job_id}] Image saved to: {output_path}")

            # Upload result
            logger.info(f"[{job_id}] Uploading result...")
            self._upload_output(output_path, upload_url)

            logger.info(f"[{job_id}] Job completed successfully")

        except Exception as e:
            logger.error(f"[{job_id}] Job failed: {e}", exc_info=True)
            raise

        finally:
            # Cleanup temporary file
            if output_path and os.path.exists(output_path):
                try:
                    os.remove(output_path)
                    logger.info(f"[{job_id}] Cleaned up temporary file")
                except Exception as e:
                    logger.warning(f"[{job_id}] Failed to remove {output_path}: {e}")

    def _upload_output(self, local_path: str, upload_url: str) -> None:
        """
        Upload the generated image to a presigned URL.

        Args:
            local_path: Path to the local image file
            upload_url: Presigned URL for upload

        Raises:
            requests.HTTPError: If the upload fails
        """
        logger.info(f"Uploading {local_path} to presigned URL")

        with open(local_path, "rb") as f:
            response = requests.put(upload_url, data=f)
            response.raise_for_status()

        logger.info("Upload complete")