import requests
import json
import time
import logging
from typing import Any, Dict, Optional, Sequence
from urllib.parse import quote

from asset_processor import get_image_metadata

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("ComfyUIClient")

class ComfyUIClient:
    def __init__(self, base_url):
        self.base_url = base_url
        self.available_models = self._get_available_models()
    
    def refresh_models(self):
        """Re-fetch available models and update available_models list."""
        self.available_models = self._get_available_models()

    def _get_available_models(self):
        """Fetch list of available checkpoint models from ComfyUI"""
        try:
            response = requests.get(f"{self.base_url}/object_info/CheckpointLoaderSimple", timeout=10)
            if response.status_code != 200:
                logger.warning("Failed to fetch model list; using default handling")
                return []
            data = response.json()
            # Safe dictionary access with proper error handling
            try:
                checkpoint_info = data.get("CheckpointLoaderSimple", {})
                if not isinstance(checkpoint_info, dict):
                    logger.warning("Unexpected CheckpointLoaderSimple structure")
                    return []
                input_info = checkpoint_info.get("input", {})
                if not isinstance(input_info, dict):
                    logger.warning("Unexpected input structure")
                    return []
                required_info = input_info.get("required", {})
                if not isinstance(required_info, dict):
                    logger.warning("Unexpected required structure")
                    return []
                ckpt_name_info = required_info.get("ckpt_name", [])
                if not isinstance(ckpt_name_info, list) or len(ckpt_name_info) == 0:
                    logger.warning("No checkpoint models found in API response")
                    return []
                models = ckpt_name_info[0] if isinstance(ckpt_name_info[0], list) else ckpt_name_info
                logger.info(f"Available models: {models}")
                return models
            except (KeyError, IndexError, TypeError) as e:
                logger.warning(f"Unexpected API response structure: {e}")
                return []
        except requests.RequestException as e:
            logger.warning(f"Error fetching models: {e}")
            return []

    def run_custom_workflow(self, workflow: Dict[str, Any], preferred_output_keys: Sequence[str] | None = None, max_attempts: int = 30):
        if preferred_output_keys is None:
            preferred_output_keys = ("images", "image", "gifs", "gif", "audio", "audios", "files")

        prompt_id = self._queue_workflow(workflow)
        outputs = self._wait_for_prompt(prompt_id, max_attempts=max_attempts)
        
        # Extract asset info (filename, subfolder, type) - stable identity
        asset_info = self._extract_first_asset_info(outputs, preferred_output_keys)
        asset_url = asset_info["asset_url"]
        
        # Extract asset metadata (pass workflow to extract dimensions from it)
        asset_metadata = self._get_asset_metadata(asset_url, outputs, preferred_output_keys, workflow)
        
        # Get full history snapshot for this prompt
        try:
            history = self.get_history(prompt_id)
            comfy_history = history.get(prompt_id, {}) if history else {}
        except Exception as e:
            logger.warning(f"Failed to fetch history snapshot for {prompt_id}: {e}")
            comfy_history = None
        
        return {
            "asset_url": asset_url,
            "filename": asset_info["filename"],
            "subfolder": asset_info["subfolder"],
            "folder_type": asset_info["type"],
            "prompt_id": prompt_id,
            "raw_outputs": outputs,
            "asset_metadata": asset_metadata,
            "comfy_history": comfy_history,
            "submitted_workflow": workflow
        }
    
    def _get_asset_metadata(self, asset_url: str, outputs: Dict[str, Any], preferred_output_keys: Sequence[str], workflow: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Extract metadata about the generated asset"""
        metadata = {
            "mime_type": None,
            "width": None,
            "height": None,
            "bytes_size": None
        }
        
        # Try to extract from outputs first
        for node_id, node_output in outputs.items():
            if not isinstance(node_output, dict):
                continue
            for key in preferred_output_keys:
                assets = node_output.get(key)
                if assets and isinstance(assets, list) and len(assets) > 0:
                    asset = assets[0]
                    if isinstance(asset, dict):
                        # Infer mime type from filename extension
                        filename = asset.get("filename", "")
                        if filename.endswith((".png", ".PNG")):
                            metadata["mime_type"] = "image/png"
                        elif filename.endswith((".jpg", ".jpeg", ".JPG", ".JPEG")):
                            metadata["mime_type"] = "image/jpeg"
                        elif filename.endswith((".webp", ".WEBP")):
                            metadata["mime_type"] = "image/webp"
                        elif filename.endswith((".mp3", ".MP3")):
                            metadata["mime_type"] = "audio/mpeg"
                        elif filename.endswith((".mp4", ".MP4")):
                            metadata["mime_type"] = "video/mp4"
                        elif filename.endswith((".gif", ".GIF")):
                            metadata["mime_type"] = "image/gif"
                        break
        
        # Extract dimensions from workflow (EmptyLatentImage node) - much more efficient than analyzing image
        if workflow and (metadata["width"] is None or metadata["height"] is None):
            for node_id, node_data in workflow.items():
                if not isinstance(node_data, dict):
                    continue
                if node_data.get("class_type") == "EmptyLatentImage":
                    inputs = node_data.get("inputs", {})
                    if "width" in inputs and metadata["width"] is None:
                        metadata["width"] = inputs["width"]
                    if "height" in inputs and metadata["height"] is None:
                        metadata["height"] = inputs["height"]
                    if metadata["width"] and metadata["height"]:
                        break
        
        # Try to fetch headers to get size (non-blocking, best effort)
        try:
            response = requests.head(asset_url, timeout=5)
            if response.status_code == 200:
                content_length = response.headers.get("Content-Length")
                if content_length:
                    metadata["bytes_size"] = int(content_length)
                content_type = response.headers.get("Content-Type")
                if content_type and not metadata["mime_type"]:
                    metadata["mime_type"] = content_type.split(";")[0].strip()
        except Exception as e:
            logger.debug(f"Could not fetch asset metadata: {e}")
        
        # Fallback: Extract image dimensions by analyzing image bytes (only if not found in workflow)
        # This should rarely be needed now, but kept as a fallback
        if metadata["mime_type"] and metadata["mime_type"].startswith("image/") and (metadata["width"] is None or metadata["height"] is None):
            try:
                # Fetch image bytes to extract dimensions
                img_response = requests.get(asset_url, timeout=10)
                if img_response.status_code == 200:
                    image_bytes = img_response.content
                    # Update bytes_size if we got it from the full response
                    if not metadata["bytes_size"]:
                        metadata["bytes_size"] = len(image_bytes)
                    # Extract dimensions
                    img_metadata = get_image_metadata(image_bytes)
                    if img_metadata.get("width") and img_metadata.get("height"):
                        metadata["width"] = img_metadata["width"]
                        metadata["height"] = img_metadata["height"]
            except Exception as e:
                logger.debug(f"Could not extract image dimensions: {e}")
        
        return metadata

    def _queue_workflow(self, workflow: Dict[str, Any]):
        logger.info("Submitting workflow to ComfyUI...")
        response = requests.post(f"{self.base_url}/prompt", json={"prompt": workflow}, timeout=30)
        if response.status_code != 200:
            raise Exception(f"Failed to queue workflow: {response.status_code} - {response.text}")
        try:
            response_data = response.json()
            prompt_id = response_data.get("prompt_id")
            if not prompt_id:
                raise Exception("Response missing prompt_id")
        except (KeyError, ValueError) as e:
            raise Exception(f"Invalid response format from ComfyUI: {e}")
        logger.info(f"Queued workflow with prompt_id: {prompt_id}")
        return prompt_id

    def _wait_for_prompt(self, prompt_id: str, max_attempts: int = 30):
        for attempt in range(max_attempts):
            try:
                # Try both the specific prompt_id endpoint and the full history endpoint
                response = requests.get(f"{self.base_url}/history/{prompt_id}", timeout=10)
                # If that doesn't work, we can also try: f"{self.base_url}/history"
                if response.status_code != 200:
                    logger.warning("History endpoint returned %s on attempt %s", response.status_code, attempt + 1)
                    time.sleep(1)
                    continue
                
                history = response.json()
                if not isinstance(history, dict):
                    logger.warning("Invalid history response format on attempt %s", attempt + 1)
                    time.sleep(1)
                    continue
                
                if prompt_id not in history:
                    # Workflow might still be running, wait and retry
                    if attempt < max_attempts - 1:
                        time.sleep(1)
                        continue
                    else:
                        # Last attempt - check if there's any history at all
                        logger.warning("Prompt ID not found in history. Available IDs: %s", list(history.keys())[:10])
                        time.sleep(1)
                        continue
                
                prompt_data = history[prompt_id]
                if not isinstance(prompt_data, dict):
                    logger.warning("Prompt data is not a dict on attempt %s", attempt + 1)
                    time.sleep(1)
                    continue
                
                # Check for workflow errors
                if "error" in prompt_data:
                    error_info = prompt_data["error"]
                    raise Exception(f"Workflow failed with error: {json.dumps(error_info, indent=2)}")
                
                # Check if workflow status indicates failure
                status = prompt_data.get("status", {})
                if isinstance(status, dict) and status.get("completed") == False:
                    error_msg = status.get("messages", ["Workflow failed"])
                    raise Exception(f"Workflow failed: {error_msg}")
                
                # Get outputs
                if "outputs" not in prompt_data:
                    # Check status to see if workflow completed
                    status = prompt_data.get("status", [])
                    if isinstance(status, list) and len(status) > 0:
                        last_status = status[-1]
                        if isinstance(last_status, list) and len(last_status) > 0:
                            status_type = last_status[0]
                            if status_type == "execution_success":
                                # Workflow completed successfully but outputs not yet available
                                # Wait a bit longer for outputs to be written, especially for cached executions
                                logger.info("Workflow execution succeeded, waiting for outputs to be available...")
                                time.sleep(3)  # Give ComfyUI time to write outputs (longer for cached)
                                # Try fetching full history to see if outputs appear there
                                try:
                                    full_history_response = requests.get(f"{self.base_url}/history", timeout=10)
                                    if full_history_response.status_code == 200:
                                        full_history = full_history_response.json()
                                        if prompt_id in full_history:
                                            full_prompt_data = full_history[prompt_id]
                                            if "outputs" in full_prompt_data and full_prompt_data["outputs"]:
                                                logger.info("Found outputs in full history endpoint")
                                                return full_prompt_data["outputs"]
                                except Exception as e:
                                    logger.debug("Could not fetch full history: %s", e)
                                continue
                    logger.warning("Prompt data missing outputs on attempt %s. Full data: %s", attempt + 1, json.dumps(prompt_data, indent=2))
                    time.sleep(1)
                    continue
                
                outputs = prompt_data["outputs"]
                if not outputs or not isinstance(outputs, dict):
                    # Check if workflow actually succeeded
                    status = prompt_data.get("status", [])
                    if isinstance(status, list):
                        status_messages = [s[0] if isinstance(s, list) else str(s) for s in status]
                        if "execution_success" in status_messages:
                            # Workflow succeeded but no outputs - might need to check queue or wait longer
                            logger.warning("Workflow succeeded but outputs empty. Status: %s. Waiting longer...", status_messages)
                            time.sleep(2)
                            continue
                        raise Exception(f"Workflow completed but produced no outputs. Status: {status_messages}")
                    logger.warning("Outputs is empty or not a dict. Prompt data: %s", json.dumps(prompt_data, indent=2))
                    raise Exception("Workflow completed but produced no outputs. Check ComfyUI logs for errors.")
                
                logger.info("Workflow completed. Output nodes: %s", list(outputs.keys()))
                logger.debug("Full workflow outputs: %s", json.dumps(outputs, indent=2))
                logger.debug("Full prompt data: %s", json.dumps(prompt_data, indent=2))
                return outputs
            except requests.RequestException as e:
                logger.warning("Request error on attempt %s: %s", attempt + 1, e)
                time.sleep(1)
                continue
            except (ValueError, KeyError) as e:
                logger.warning("JSON parsing error on attempt %s: %s", attempt + 1, e)
                time.sleep(1)
                continue
        
        raise Exception(f"Workflow {prompt_id} didn't complete within {max_attempts} seconds")

    def _extract_first_asset_url(self, outputs: Dict[str, Any], preferred_output_keys: Sequence[str]):
        # Log available outputs for debugging
        logger.debug("Available output keys in workflow: %s", list(outputs.keys()))
        for node_id, node_output in outputs.items():
            if not isinstance(node_output, dict):
                logger.debug("Node %s output is not a dict: %s", node_id, type(node_output))
                continue
            logger.debug("Node %s has keys: %s", node_id, list(node_output.keys()))
            for key in preferred_output_keys:
                assets = node_output.get(key)
                if assets and isinstance(assets, list) and len(assets) > 0:
                    asset = assets[0]
                    if not isinstance(asset, dict):
                        logger.debug("Asset in node %s, key %s is not a dict", node_id, key)
                        continue
                    filename = asset.get("filename")
                    if not filename:
                        logger.debug("Asset in node %s, key %s missing filename", node_id, key)
                        continue
                    subfolder = asset.get("subfolder", "")
                    output_type = asset.get("type", "output")
                    logger.info("Found asset: filename=%s, subfolder=%s, type=%s", filename, subfolder, output_type)
                    return f"{self.base_url}/view?filename={filename}&subfolder={subfolder}&type={output_type}"
        
        # Enhanced error message with actual output structure
        logger.error("No outputs matched preferred keys: %s", preferred_output_keys)
        logger.error("Actual outputs structure: %s", json.dumps(outputs, indent=2))
        raise Exception(
            f"No outputs matched preferred keys: {preferred_output_keys}. "
            f"Available outputs: {json.dumps({k: list(v.keys()) if isinstance(v, dict) else type(v).__name__ for k, v in outputs.items()}, indent=2)}"
        )
    
    def _extract_first_asset_info(self, outputs: Dict[str, Any], preferred_output_keys: Sequence[str]) -> Dict[str, Any]:
        """Extract first asset info (filename, subfolder, type) from outputs.
        
        Returns dict with 'filename', 'subfolder', 'type', and 'asset_url'.
        """
        logger.debug("Available output keys in workflow: %s", list(outputs.keys()))
        for node_id, node_output in outputs.items():
            if not isinstance(node_output, dict):
                continue
            for key in preferred_output_keys:
                assets = node_output.get(key)
                if assets and isinstance(assets, list) and len(assets) > 0:
                    asset = assets[0]
                    if not isinstance(asset, dict):
                        continue
                    filename = asset.get("filename")
                    if not filename:
                        continue
                    subfolder = asset.get("subfolder", "")
                    output_type = asset.get("type", "output")
                    
                    # URL encode for special characters
                    base_url = self.base_url.rstrip('/')
                    encoded_filename = quote(filename, safe='')
                    encoded_subfolder = quote(subfolder, safe='') if subfolder else ''
                    
                    if encoded_subfolder:
                        asset_url = f"{base_url}/view?filename={encoded_filename}&subfolder={encoded_subfolder}&type={output_type}"
                    else:
                        asset_url = f"{base_url}/view?filename={encoded_filename}&type={output_type}"
                    
                    return {
                        "filename": filename,
                        "subfolder": subfolder,
                        "type": output_type,
                        "asset_url": asset_url
                    }
        
        raise Exception(
            f"No outputs matched preferred keys: {preferred_output_keys}. "
            f"Available outputs: {json.dumps({k: list(v.keys()) if isinstance(v, dict) else type(v).__name__ for k, v in outputs.items()}, indent=2)}"
        )
    
    def get_queue(self) -> Dict[str, Any]:
        """Get current queue status from ComfyUI.
        
        Returns the full /queue endpoint response.
        """
        try:
            response = requests.get(f"{self.base_url}/queue", timeout=10)
            response.raise_for_status()
            return response.json()
        except requests.RequestException as e:
            logger.error(f"Failed to get queue status: {e}")
            raise Exception(f"Failed to get queue status: {e}")
    
    def get_history(self, prompt_id: Optional[str] = None) -> Dict[str, Any]:
        """Get history from ComfyUI.
        
        Args:
            prompt_id: Optional specific prompt ID. If None, returns full history.
        
        Returns:
            History dict. If prompt_id provided, returns {prompt_id: {...}} or {} if not found.
        """
        try:
            if prompt_id:
                url = f"{self.base_url}/history/{prompt_id}"
            else:
                url = f"{self.base_url}/history"
            response = requests.get(url, timeout=10)
            response.raise_for_status()
            return response.json()
        except requests.RequestException as e:
            logger.error(f"Failed to get history: {e}")
            raise Exception(f"Failed to get history: {e}")
    
    def cancel_prompt(self, prompt_id: str) -> Dict[str, Any]:
        """Cancel a queued or running prompt.
        
        Args:
            prompt_id: The prompt ID to cancel.
        
        Returns:
            Response from ComfyUI cancel endpoint.
        """
        try:
            response = requests.post(
                f"{self.base_url}/queue",
                json={"delete": [prompt_id]},
                timeout=10
            )
            response.raise_for_status()
            return response.json()
        except requests.RequestException as e:
            logger.error(f"Failed to cancel prompt {prompt_id}: {e}")
            raise Exception(f"Failed to cancel prompt: {e}")