import requests
import base64
import asyncio
import config

from typing import Any, Dict, List
from tools.utils import ToolError, load_toml

async def ensure_model_state(model_name: str, model_preset: Dict[str, Any]):
    """
    Checks the current server state and switches model/modules if they differ from the preset.
    """
    try:
        config_resp = await asyncio.to_thread(requests.get, f"{config.SD_URL}/config", timeout=10)
        config_resp.raise_for_status()
        cfg_data = config_resp.json()
        components = cfg_data.get("components", [])
        
        # Extract current state from components
        current_state = {}
        for comp in components:
            elem_id = comp.get("props", {}).get("elem_id")
            if elem_id:
                current_state[elem_id] = comp.get("props", {}).get("value")
        
        # 1. Check and change Checkpoint
        active_ckpt = current_state.get("setting_sd_model_checkpoint")
        # Use 'filename' from TOML if available, otherwise fall back to the model_name key
        target_ckpt = model_preset.get("filename", model_name)
        
        if active_ckpt != target_ckpt:
            preset = model_preset.get("preset", "xl")
            await asyncio.to_thread(
                requests.post,
                f"{config.SD_URL}/api/checkpoint_change", 
                json={"data": [target_ckpt, preset]}, 
                timeout=30
            )
        
        # 2. Check and change VAE / Text Encoders
        active_modules = current_state.get("setting_sd_modules", [])
        target_modules = model_preset.get("modules", [])
        
        if active_modules != target_modules:
            preset = model_preset.get("preset", "xl")
            await asyncio.to_thread(
                requests.post,
                f"{config.SD_URL}/api/modules_change", 
                json={"data": [target_modules, preset]}, 
                timeout=30
            )
            
    except Exception as e:
        # We log this but don't necessarily raise a ToolError unless the generation itself fails,
        # as the server might still be able to generate if it's just a state-check failure.
        print(f"Warning: Failed to sync model state: {e}")

async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
    """
    Generates an image using the specified model and parameters.
    """
    # 1. REQUIRED ARGUMENTS
    prompt = args.get("prompt")
    if not prompt:
        raise ToolError("The 'prompt' argument is required.")
    
    model_name = args.get("model_name")
    if not model_name:
        raise ToolError("The 'model_name' argument is required. Use get_model_info to see available models.")

    res_preset_name = args.get("resolution_preset")
    if not res_preset_name:
        raise ToolError("The 'resolution_preset' argument is required. Use get_model_info to see available resolutions for the chosen model.")

    # 2. FETCH CURRENT SERVER DEFAULTS & MAP LABELS
    try:
        info_resp = await asyncio.to_thread(requests.get, f"{config.SD_URL}/info", timeout=10)
        info_resp.raise_for_status()
        info_data = info_resp.json()
        params_info = info_data["named_endpoints"]["/txt2img"]["parameters"]
    except Exception as e:
        raise ToolError(f"Failed to connect to Stable Diffusion server: {str(e)}")

    # Build the label-to-index map and a sparse payload
    label_map = {}
    label_counts = {}
    
    # Find the maximum param index to determine payload size
    max_param_idx = 0
    for p in params_info:
        name = p.get("parameter_name", "")
        if name.startswith("param_"):
            try:
                idx = int(name.replace("param_", ""))
                max_param_idx = max(max_param_idx, idx)
            except ValueError:
                pass

    # Initialize payload with Nones (size is max_idx + 1)
    payload = [None] * (max_param_idx + 1)
    
    for idx_in_list, p in enumerate(params_info):
        name = p.get("parameter_name", "")
        label = p.get("label")
        default = p.get("parameter_default")
        
        # Determine the absolute index in the payload
        if name == "id_task":
            abs_idx = 0
        elif name.startswith("param_"):
            try:
                abs_idx = int(name.replace("param_", ""))
            except ValueError:
                continue
        else:
            # Fallback for unexpected names, though unlikely
            continue

        # Set the default value at the absolute index
        payload[abs_idx] = default
        
        # Handle label mapping for AI overrides
        if not label or label.startswith("parameter_"):
            label = name
            
        if label in label_counts:
            label_counts[label] += 1
            mapped_label = f"{label} [{label_counts[label]}]"
        else:
            label_counts[label] = 0
            mapped_label = label
            
        label_map[mapped_label] = abs_idx

    # 3. LOAD CONFIGS
    models_cfg = load_toml(config.MODEL_PRESETS_PATH)
    res_cfg = load_toml(config.RES_PRESETS_PATH)
    
    if model_name not in models_cfg:
        raise ToolError(f"Model '{model_name}' not found in presets. Please call get_model_info to see available models and their correct names.")
    
    model_preset = models_cfg[model_name]

    # Ensure server state (Checkpoint, VAE, etc.) matches the preset before generating
    await ensure_model_state(model_name, model_preset)

    # 4. MERGE PIPELINE
    for key, value in model_preset.items():
        if key in ["description", "guide", "Resolution Set"]:
            continue
        if key in label_map:
            payload[label_map[key]] = value
        elif key.startswith("param_"):
            try:
                idx = int(key.replace("param_", ""))
                if 0 <= idx < len(payload):
                    payload[idx] = value
            except ValueError:
                pass

    res_preset_name = args.get("resolution_preset")
    if res_preset_name:
        res_set_name = model_preset.get("Resolution Set")
        if res_set_name and res_set_name in res_cfg:
            res_set = res_cfg[res_set_name]
            if res_preset_name in res_set:
                res_val_str = res_set[res_preset_name]
                try:
                    w, h = map(int, res_val_str.split('x'))
                    if "Width" in label_map:
                        payload[label_map["Width"]] = w
                    if "Height" in label_map:
                        payload[label_map["Height"]] = h
                except (ValueError, AttributeError):
                    raise ToolError(f"Invalid resolution format for preset '{res_preset_name}': {res_val_str}. Expected 'WidthxHeight'.")
            else:
                raise ToolError(f"Resolution preset '{res_preset_name}' not found for this model. You MUST call get_model_info with this model_name to see the available resolution presets.")
        else:
            raise ToolError(f"No resolution set configured for model '{model_name}'.")

    overrides = {
        "prompt": "Prompt",
        "negative_prompt": "Negative Prompt",
        "cfg_scale": "CFG Scale",
    }
    
    for arg_key, label in overrides.items():
        if arg_key in args and label in label_map:
            payload[label_map[label]] = args[arg_key]
    
    if "Prompt" in label_map:
        payload[label_map["Prompt"]] = prompt

    if "Seed" in label_map:
        payload[label_map["Seed"]] = -1

    # The "Magic Number" splice is no longer needed as gaps are 
    # automatically filled by the sparse-to-dense mapping.

    # 6. EXECUTE GENERATION
    try:
        gen_payload = {"data": payload}
        gen_resp = await asyncio.to_thread(requests.post, f"{config.SD_URL}/api/txt2img", json=gen_payload, timeout=300)
        gen_resp.raise_for_status()
        res_data = gen_resp.json()
        
        gallery_data = res_data["data"][0]["value"]
        if not gallery_data:
            raise ToolError("Server returned successfully but no image was generated.")
            
        sd_img_url = gallery_data[0]["image"]["url"]
        
        # Extract raw path from SD URL (e.g. http://.../file=/tmp/gradio/abc.png -> /tmp/gradio/abc.png)
        if "/file=" in sd_img_url:
            raw_path = sd_img_url.split("/file=")[1]
        else:
            # Fallback if the URL format changes
            raise ToolError(f"Could not extract file path from SD URL: {sd_img_url}")

        # Construct proxy URL through our MCP server
        # Use the captured request host to ensure links work across the network
        host = config.request_host.get()
        proxy_url = f"http://{host}/file{raw_path}"
        
        # Download the image and convert to base64 for the AI's vision
        img_resp = await asyncio.to_thread(requests.get, sd_img_url, timeout=30)
        img_resp.raise_for_status()
        b64_data = base64.b64encode(img_resp.content).decode('utf-8')
        
        return [
            {
                "type": "text", 
                "text": f"Image successfully generated using model '{model_name}'.\n\nLocal path: {raw_path}\nTo show the image to the user, you can include this markdown link in your response: ![Generated Image]({proxy_url})"
            },
            {
                "type": "image",
                "data": b64_data,
                "mimeType": "image/png"
            }
        ]
        
    except requests.exceptions.HTTPError as e:
        raise ToolError(f"Server error during generation: {str(e)}")
    except Exception as e:
        raise ToolError(f"Unexpected error during generation: {str(e)}")
