"""Hunyuan3D-2.1 3D model generation."""

# CRITICAL: Import spaces BEFORE torch/CUDA packages
import spaces

import torch
from pathlib import Path
from gradio_client import Client, handle_file
import httpx
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type

from core.config import HUNYUAN_SETTINGS, QualityPreset
from utils.memory import MemoryManager


class HunyuanGenerator:
    """Generates 3D models using Hunyuan3D-2.1."""
    
    def __init__(self):
        self.memory_manager = MemoryManager()
    
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception_type((httpx.TimeoutException, httpx.NetworkError))
    )
    def _call_api(self, client: Client, **kwargs):
        """Call Hunyuan3D API with automatic retry."""
        return client.predict(**kwargs)
    
    @spaces.GPU(duration=90)
    def generate(
        self,
        image_path: Path,
        preset: QualityPreset,
        output_dir: Path
    ) -> Path:
        """Generate 3D model from 2D image."""
        try:
            print(f"[Hunyuan3D] Generating 3D model: {preset.name} quality")
            print(f"[Hunyuan3D] Input image: {image_path}")
            print(f"[Hunyuan3D] Settings: steps={preset.hunyuan_steps}, guidance={preset.hunyuan_guidance}, octree={preset.octree_resolution}")
            
            # Validate input image exists
            if not image_path.exists():
                raise FileNotFoundError(f"Input image not found: {image_path}")
            
            # Connect to API
            print(f"[Hunyuan3D] Connecting to {HUNYUAN_SETTINGS['space_id']}...")
            client = Client(
                HUNYUAN_SETTINGS["space_id"],
                httpx_kwargs={
                    "timeout": httpx.Timeout(
                        HUNYUAN_SETTINGS["timeout"],
                        connect=HUNYUAN_SETTINGS["connect_timeout"]
                    )
                }
            )
            print(f"[Hunyuan3D] Connected successfully")
            
            # Call API (returns tuple: file, output, mesh_stats, seed)
            print(f"[Hunyuan3D] Calling API with parameters...")
            result = self._call_api(
                client,
                image=handle_file(str(image_path)),
                mv_image_front=None,
                mv_image_back=None,
                mv_image_left=None,
                mv_image_right=None,
                steps=preset.hunyuan_steps,
                guidance_scale=preset.hunyuan_guidance,
                seed=1234,
                octree_resolution=preset.octree_resolution,
                check_box_rembg=True,
                num_chunks=preset.num_chunks,
                randomize_seed=True,
                api_name="/shape_generation"
            )
            print(f"[Hunyuan3D] API call completed")
            
            # Extract GLB path from tuple response
            # API returns: (file, output, mesh_stats, seed)
            print(f"[Hunyuan3D] Raw result type: {type(result)}")
            print(f"[Hunyuan3D] Raw result length: {len(result) if isinstance(result, (tuple, list)) else 'N/A'}")
            
            if not isinstance(result, tuple):
                raise ValueError(
                    f"Unexpected result type from Hunyuan3D API: {type(result)}. "
                    f"Expected tuple of (file, output, mesh_stats, seed)."
                )
            
            if len(result) != 4:
                raise ValueError(
                    f"Unexpected result length from Hunyuan3D API: {len(result)}. "
                    f"Expected 4 elements (file, output, mesh_stats, seed), got {len(result)}."
                )
            
            # Extract GLB file path (first element)
            file_data, html_output, mesh_stats, used_seed = result
            
            print(f"[Hunyuan3D] file_data type: {type(file_data)}")
            print(f"[Hunyuan3D] mesh_stats: {mesh_stats}")
            print(f"[Hunyuan3D] used_seed: {used_seed}")
            
            # Extract path from file_data
            if file_data is None:
                raise ValueError(
                    "Hunyuan3D API returned None for file. "
                    "This usually means the generation failed on the server side. "
                    "Possible causes:\n"
                    "  - Invalid image input\n"
                    "  - API timeout\n"
                    "  - Server overload\n"
                    "Try again with a different image or quality setting."
                )
            
            # Handle different file_data formats
            if isinstance(file_data, dict):
                print(f"[Hunyuan3D] file_data is dict with keys: {file_data.keys()}")
                if 'path' in file_data:
                    glb_path = file_data['path']
                elif 'value' in file_data:
                    glb_path = file_data['value']
                elif 'name' in file_data:
                    glb_path = file_data['name']
                else:
                    raise ValueError(
                        f"Unexpected dict format from Hunyuan3D API. "
                        f"Keys: {list(file_data.keys())}"
                    )
            elif isinstance(file_data, str):
                glb_path = file_data
            else:
                raise ValueError(
                    f"Unexpected file_data type: {type(file_data)}. "
                    f"Expected dict or str."
                )
            
            print(f"[Hunyuan3D] Extracted GLB path: {glb_path}")
            
            # Validate path exists
            if not glb_path or glb_path == "None":
                raise ValueError(
                    "Hunyuan3D API returned invalid path. "
                    "The generation may have failed on the server side."
                )
            
            if not Path(glb_path).exists():
                raise ValueError(
                    f"GLB file not found at path: {glb_path}. "
                    f"The file may not have been generated or saved correctly."
                )
            
            print(f"[Hunyuan3D] Model generated: {glb_path}")
            
            # Cleanup
            del client
            import gc
            gc.collect()
            torch.cuda.empty_cache()
            
            return Path(glb_path)
            
        except Exception as e:
            import traceback
            error_details = traceback.format_exc()
            print(f"[Hunyuan3D] ERROR: {e}")
            print(f"[Hunyuan3D] Full traceback:\n{error_details}")
            
            # Provide helpful error message
            if "list index out of range" in str(e):
                raise ValueError(
                    f"Hunyuan3D API returned unexpected result format. "
                    f"This usually means the generation failed on the server side. "
                    f"Please try again with a different prompt or quality setting."
                ) from e
            elif "timeout" in str(e).lower():
                raise TimeoutError(
                    f"Hunyuan3D generation timed out. "
                    f"Try using a lower quality preset (Fast or Balanced)."
                ) from e
            else:
                raise RuntimeError(
                    f"Hunyuan3D generation failed: {e}. "
                    f"Check logs for details."
                ) from e