"""Image Upscaler Service - Topaz Labs API

Available Models:
- proteus: General enhancement with fine-tuning parameters (default)
- artemis: Detail enhancement and noise reduction
- gaia: Specialized for HD/4K upscaling
- iris: Noise and compression artifact reduction
- nyx: Low light and high ISO recovery
- rhea: Detail recovery for older/degraded images
- theia: High-fidelity upscaling

Output Options:
- Scale: 2x, 4x, 6x, 8x (up to 16K)
- Output formats: png, jpg, tiff
- Face enhancement: auto-detect and enhance faces
- Noise reduction: 0-100
- Sharpening: 0-100
- Grain recovery: preserve film grain
"""
import httpx
import os
from uuid import uuid4
from datetime import datetime
import asyncio
from typing import Optional, Dict, Any

from app.database import SessionLocal
from app.models.job import Job
from app.models.asset import Asset
from app.config import settings
import logging

logger = logging.getLogger(__name__)


# Topaz enhancement models with their specialties
TOPAZ_MODELS = {
    "Standard V2": {
        "name": "Standard V2",
        "description": "General-purpose model balancing detail, sharpness, and noise reduction",
        "parameters": ["face_enhancement"],
        "best_for": "General purpose"
    },
    "High Fidelity V2": {
        "name": "High Fidelity V2",
        "description": "Ideal for high-quality images, preserving intricate details",
        "parameters": ["face_enhancement"],
        "best_for": "High quality sources"
    },
    "Low Resolution V2": {
        "name": "Low Resolution V2",
        "description": "Designed for enhancing clarity and detail in low-resolution images",
        "parameters": ["face_enhancement"],
        "best_for": "Low resolution / compressed"
    },
    "CGI": {
        "name": "CGI",
        "description": "Optimized for computer-generated imagery and digital illustrations",
        "parameters": ["face_enhancement"],
        "best_for": "Digital art, 3D renders"
    },
    "Text Refine": {
        "name": "Text Refine",
        "description": "Optimized for images containing text",
        "parameters": [],
        "best_for": "Documents, screenshots"
    },
    "Enhance Generative": {
        "name": "Enhance Generative",
        "description": "Generative model for high quality and creative detail (slower)",
        "parameters": ["face_enhancement"],
        "best_for": "High quality creative upscaling"
    }
}


async def upscale(job_id: str):
    """Upscale image using Topaz Labs API

    Input parameters:
    - scale: Upscale factor (2, 4, 6, 8)
    - model: Enhancement model (see TOPAZ_MODELS)
    - output_format: 'png', 'jpg', 'tiff' (default: png)
    - face_enhancement: Boolean to enable face detection and enhancement
    - noise_reduction: 0-100, amount of noise removal
    - sharpening: 0-100, output sharpening level
    - compression_recovery: 0-100, recover compression artifacts
    - detail_enhancement: 0-100, enhance fine details
    - preserve_grain: Boolean to preserve film grain
    - output_quality: 1-100 for jpg output (default: 95)
    """
    db = SessionLocal()
    try:
        job = db.query(Job).filter(Job.id == job_id).first()
        if not job:
            return

        input_data = job.input_data
        input_asset_ids = job.input_asset_ids

        if not input_asset_ids:
            raise ValueError("No input asset provided")

        # Get input asset
        input_asset = db.query(Asset).filter(Asset.id == input_asset_ids[0]).first()
        if not input_asset:
            raise ValueError("Input asset not found")

        # Extract parameters
        scale = input_data.get("scale", 2)
        model = input_data.get("model", "auto")
        output_format = input_data.get("output_format", "png")
        face_enhancement = input_data.get("face_enhancement", False)
        noise_reduction = input_data.get("noise_reduction")
        sharpening = input_data.get("sharpening")
        compression_recovery = input_data.get("compression_recovery")
        detail_enhancement = input_data.get("detail_enhancement")
        preserve_grain = input_data.get("preserve_grain", False)
        output_quality = input_data.get("output_quality", 95)

        job.progress = 10
        job.api_provider = "topaz"
        job.api_model = model
        db.commit()

        # Read input image
        with open(input_asset.file_path, "rb") as f:
            image_data = f.read()

        # Ensure dimensions are set - extract if missing
        if not input_asset.width or not input_asset.height:
            from PIL import Image
            import io
            img = Image.open(io.BytesIO(image_data))
            original_width, original_height = img.size

            # Update asset with correct dimensions
            input_asset.width = original_width
            input_asset.height = original_height
            db.commit()
        else:
            original_width = input_asset.width
            original_height = input_asset.height
            
        logger.info(f"Topaz Image Setup: {original_width}x{original_height}, Scale: {scale}, Model: {model}")

        # Calculate output dimensions
        output_width = original_width * scale
        output_height = original_height * scale

        job.progress = 20
        db.commit()

        # Build enhancement parameters matching simple working PHP version
        enhance_params: Dict[str, Any] = {
            "output_height": str(output_height),
            "output_format": output_format if output_format in ["jpeg", "jpg", "png", "tiff", "tif"] else "png",
            "crop_to_fill": "true" if input_data.get("crop_to_fill") else "false",
            "face_enhancement": "true" if input_data.get("face_enhancement") else "false"
        }
        
        # Handle 'auto' model: simplest is to OMIT parameter if auto
        if model.lower() != "auto":
             enhance_params["model"] = model

        # Strict PHP parity for reliability:
        # process.php ONLY sent: image, output_height, output_format, crop_to_fill, face_enhancement, model.
        # It did NOT send denoise, sharpen, creativity, etc.
        # To avoid 422 errors or timeouts, we omit them until verified.

        # Call Topaz API
        async with httpx.AsyncClient(timeout=600) as client:
            # Start async enhancement
            response = await client.post(
                "https://api.topazlabs.com/image/v1/enhance/async",
                headers={
                    "X-API-Key": settings.topaz_api_key,
                    "Accept": "application/json"
                },
                files={"image": (input_asset.original_filename, image_data, input_asset.mime_type)},
                data=enhance_params
            )
            response.raise_for_status()
            result = response.json()

            request_id = result.get("id") or result.get("requestId") or result.get("process_id")

            job.progress = 40
            job.api_request_id = request_id
            db.commit()

            # Poll for completion
            output_url = None
            polling_interval = 2
            max_attempts = 180
            status_data = {}
            
            for i in range(max_attempts):
                await asyncio.sleep(polling_interval)

                try:
                    status_response = await client.get(
                        f"https://api.topazlabs.com/image/v1/status/{request_id}",
                        headers={"X-API-Key": settings.topaz_api_key}
                    )
                    
                    if status_response.status_code != 200:
                        logger.warning(f"Topaz Status Check Failed: {status_response.text}")
                        continue
                        
                    status_data = status_response.json()
                    status = status_data.get("status", "")
                    
                    # Topaz uses different status values and field names
                    topaz_status = status.lower() if status else ""
                    
                    # Log status occasionally
                    if i % 5 == 0:
                        logger.info(f"Topaz Job {request_id} Status: {topaz_status} (Attempt {i}/{max_attempts})")


                    if topaz_status == "completed":
                        logger.info(f"Topaz job completed, fetching download URL...")
                        break
                    elif topaz_status == "failed":
                        error_msg = status_data.get("error") or "Unknown error"
                        raise ValueError(f"Topaz enhancement failed: {error_msg}")

                    # Update progress
                    # Start at 40, go up to 85.
                    progress_increment = (85 - 40) / max_attempts
                    current_progress = 40 + (i * progress_increment)
                    
                    if i % 2 == 0:
                         job.progress = int(current_progress)
                         db.commit()

                except Exception as loop_e:
                     logger.error(f"Error in polling loop: {loop_e}")
                     continue

            if topaz_status != "completed":
                 raise TimeoutError("Topaz upscaling timed out or did not complete.")

            # Call the download endpoint to get the download URL
            logger.info(f"Calling download endpoint for request_id: {request_id}")
            download_response = await client.get(
                f"https://api.topazlabs.com/image/v1/download/{request_id}",
                headers={"X-API-Key": settings.topaz_api_key}
            )
            download_response.raise_for_status()
            download_data = download_response.json()
            
            output_url = download_data.get("download_url")
            if not output_url:
                raise ValueError(f"No download_url in response: {download_data}")

            logger.info(f"Topaz download URL received: {output_url[:100] if output_url else 'None'}")
            # Download result
            img_response = await client.get(output_url)
            upscaled_data = img_response.content
            logger.info(f"Downloaded upscaled image: {len(upscaled_data)} bytes")

            job.progress = 90
            db.commit()

            # Determine output extension
            ext_map = {"png": ".png", "jpg": ".jpg", "jpeg": ".jpg", "tiff": ".tiff"}
            ext = ext_map.get(output_format, ".png")
            mime_map = {"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg", "tiff": "image/tiff"}
            mime = mime_map.get(output_format, "image/png")

            # Save output
            # Save output
            base_name = os.path.splitext(input_asset.original_filename)[0]
            # Clean base name: replace spaces with underscores
            clean_base_name = base_name.replace(" ", "_")
            # Clean model name: remove spaces, lowercase, etc - User wants specific format "Proteus" etc.
            # actually user said "Model name in original case or uppercase".
            # The model var comes from input_data.get("model").
            # Let's keep it as is but replace spaces with underscores if any.
            clean_model = model.replace(" ", "_")
            
            filename = f"{clean_base_name}_{scale}X_{clean_model}{ext}"
            # Ensure unique filename if needed in future, but for now this naming convention is requested
            # If we really need uniqueness, maybe append a short hash, but user asked specifically for this format
            # Let's check if we should add a short suffix to avoid collisions if they run it twice
            # or rely on duplicate handling?
            # User request: "like 2X-Auto or 2X-CGI or 2X-Standard-V2"
            # So we prioritize readability.
            # filename = f"upscaled_{scale}x_{model}_{uuid4()}{ext}" <-- OLD
            storage_path = os.path.join(settings.storage_path, "images")
            os.makedirs(storage_path, exist_ok=True)
            file_path = os.path.join(storage_path, filename)

            with open(file_path, "wb") as f:
                f.write(upscaled_data)

            # Create output asset
            output_asset = Asset(
                user_id=job.user_id,
                project_id=job.project_id,
                original_filename=filename,
                stored_filename=filename,
                file_path=file_path,
                file_type="image",
                mime_type=mime,
                file_size_bytes=len(upscaled_data),
                width=output_width,
                height=output_height,
                source_module="image_upscaler",
                source_job_id=job.id,
                parent_asset_id=input_asset.id,
                asset_metadata={
                    "scale": scale,
                    "model": model,
                    "face_enhancement": face_enhancement,
                    "noise_reduction": noise_reduction,
                    "sharpening": sharpening,
                    "original_dimensions": f"{original_width}x{original_height}",
                    "output_dimensions": f"{output_width}x{output_height}"
                }
            )
            db.add(output_asset)
            db.commit()
            db.refresh(output_asset)

            job.output_asset_ids = [output_asset.id]
            job.output_data = {"asset_id": str(output_asset.id), "file_path": file_path}
            logger.info(f"✓ Topaz upscale completed: Asset {output_asset.id} created")
            job.progress = 100
            job.status = "completed"
            job.completed_at = datetime.utcnow()
            db.commit()

    except Exception as e:
        job.status = "failed"
        job.error_message = str(e)
        db.commit()
    finally:
        db.close()


def get_available_models() -> Dict[str, Any]:
    """Get all available Topaz upscaling models and their capabilities"""
    return TOPAZ_MODELS
