"""
LTX Video worker implementation for video generation inference.

This module provides the LTXVWorker class that handles video generation jobs
using the LTX Video model.
"""

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

import requests

from auror.core.worker import Worker

logger = logging.getLogger(__name__)


class LTXVWorker(Worker):
    """
    Worker for processing LTX Video generation jobs.

    This worker handles the complete pipeline for video generation:
    1. Downloads conditioning media from provided URLs
    2. Configures the LTX Video inference pipeline
    3. Runs the inference
    4. Uploads the generated video to the provided presigned URL
    """

    def __init__(
        self,
        pipeline_config: str = "configs/ltxv-13b-0.9.8-distilled.yaml",
    ):
        """
        Initialize the LTXV worker.

        Args:
            pipeline_config: Path to the LTX Video pipeline configuration file
        """
        super().__init__()
        self.pipeline_config = pipeline_config

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

        Args:
            job: Job dictionary containing:
                - id: Unique job identifier
                - input: Input parameters including:
                    - prompt: Text prompt for video generation
                    - conditioningMedias: List of URLs for conditioning media
                    - conditioningStartFrames: List of starting frames for conditioning
                    - uploadUrl: Presigned URL for uploading the result
                    - seed: Random seed for reproducibility
                    - height: Output video height
                    - width: Output video width
                    - numFrames: Number of frames to generate

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

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

            # Create temporary output directory
            output_dir = tempfile.mkdtemp(prefix=f"ltxv_{job_id}_")
            logger.info(f"[{job_id}] Output directory: {output_dir}")

            # Import ltx_video only when needed (lazy loading)
            from ltx_video.inference import infer

            # Build inference configuration
            config = self._build_inference_config(job, output_dir)

            # Run inference
            logger.info(f"[{job_id}] Running inference...")
            infer(config=config)

            # Find generated video
            logger.info(f"[{job_id}] Searching for generated video...")
            output_mp4 = self._find_generated_mp4(output_dir)
            logger.info(f"[{job_id}] Found video: {output_mp4}")

            # Upload result
            upload_url = job["input"]["uploadUrl"]
            logger.info(f"[{job_id}] Uploading result...")
            self._upload_output(output_mp4, 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 files
            if output_dir and os.path.exists(output_dir):
                self._cleanup_directory(output_dir)
                logger.info(f"[{job_id}] Cleaned up temporary files")

    def _download_conditioning_media(self, url: str) -> str:
        """
        Download conditioning media from a URL to a temporary file.

        Args:
            url: URL of the media to download

        Returns:
            Path to the downloaded temporary file

        Raises:
            requests.HTTPError: If the download fails
        """
        logger.info(f"Downloading conditioning media: {url}")

        response = requests.get(url, stream=True)
        response.raise_for_status()

        # Extract file extension from URL
        suffix = os.path.splitext(url.split("?")[0])[1]

        # Create temporary file
        fd, tmp_path = tempfile.mkstemp(suffix=suffix, prefix="ltxv_input_")

        # Download and write to file
        with os.fdopen(fd, "wb") as f:
            for chunk in response.iter_content(8192):
                f.write(chunk)

        logger.info(f"Downloaded to: {tmp_path}")
        return tmp_path

    def _build_inference_config(
        self, job: Dict[str, Any], output_dir: str
    ):
        """
        Build the LTX Video inference configuration from job parameters.

        Args:
            job: Job dictionary containing input parameters
            output_dir: Directory for output files

        Returns:
            Configured InferenceConfig instance

        Raises:
            KeyError: If required job parameters are missing
            requests.HTTPError: If media download fails
        """
        # Import ltx_video only when needed (lazy loading)
        from ltx_video.inference import InferenceConfig
        from transformers import HfArgumentParser

        inp = job["input"]

        # Download all conditioning media
        local_medias = [
            self._download_conditioning_media(url)
            for url in inp["conditioningMedias"]
        ]

        # Build configuration dictionary
        config_dict = {
            "prompt": inp["prompt"],
            "output_path": output_dir,
            "pipeline_config": self.pipeline_config,
            "seed": inp["seed"],
            "height": inp["height"],
            "width": inp["width"],
            "num_frames": inp["numFrames"],
            "conditioning_media_paths": local_medias,
            "conditioning_start_frames": inp["conditioningStartFrames"],
            "conditioning_strengths": [1.0] * len(local_medias),
        }

        # Parse configuration using HuggingFace argument parser
        parser = HfArgumentParser(InferenceConfig)
        (config,) = parser.parse_dict(config_dict)

        return config

    def _find_generated_mp4(self, output_dir: str) -> str:
        """
        Find the generated MP4 file in the output directory.

        Args:
            output_dir: Directory to search for MP4 files

        Returns:
            Path to the generated MP4 file

        Raises:
            RuntimeError: If no MP4 file is found
        """
        files = os.listdir(output_dir)
        mp4_files = [f for f in files if f.lower().endswith(".mp4")]

        if not mp4_files:
            raise RuntimeError(f"No MP4 file found in {output_dir}")

        return os.path.join(output_dir, mp4_files[0])

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

        Args:
            local_path: Path to the local video 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")

    def _cleanup_directory(self, directory: str) -> None:
        """
        Clean up a directory and all its contents.

        Args:
            directory: Path to the directory to clean up

        Note:
            This method suppresses all exceptions to ensure cleanup
            doesn't fail the job.
        """
        try:
            for filename in os.listdir(directory):
                file_path = os.path.join(directory, filename)
                try:
                    if os.path.isfile(file_path):
                        os.remove(file_path)
                except Exception as e:
                    logger.warning(f"Failed to remove {file_path}: {e}")

            os.rmdir(directory)
        except Exception as e:
            logger.warning(f"Failed to remove directory {directory}: {e}")
