import requests
import tomllib
import os
import base64
import config
from typing import Any, Dict, List
from tools.utils import ToolError

# Paths to config files
MODEL_PRESETS_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "model_presets.toml")
RES_PRESETS_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "resolution_presets.toml")

SD_URL = "http://127.0.0.1:7860"

def load_toml(path: str) -> Dict[str, Any]:
    try:
        with open(path, "rb") as f:
            return tomllib.load(f)
    except Exception as e:
        raise ToolError(f"Failed to load config file {path}: {str(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.")

    # 2. FETCH CURRENT SERVER DEFAULTS & MAP LABELS
    try:
        info_resp = requests.get(f"{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
    label_map = {}
    label_counts = {}
    
    payload = [p["parameter_default"] for p in params_info]
    
    for idx, p in enumerate(params_info):
        label = p["label"]
        if not label or label.startswith("parameter_"):
            label = p["parameter_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] = idx

    # 3. LOAD CONFIGS
    models_cfg = load_toml(MODEL_PRESETS_PATH)
    res_cfg = load_toml(RES_PRESETS_PATH)
    
    if model_name not in models_cfg:
        raise ToolError(f"Model '{model_name}' not found in presets. Available: {', '.join(models_cfg.keys())}")
    
    model_preset = models_cfg[model_name]

    # 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_vals = res_set[res_preset_name]
                if "Width" in label_map:
                    payload[label_map["Width"]] = res_vals["width"]
                if "Height" in label_map:
                    payload[label_map["Height"]] = res_vals["height"]
            else:
                raise ToolError(f"Resolution preset '{res_preset_name}' not found for this model. Available: {', '.join(res_set.keys())}")
        else:
            raise ToolError(f"No resolution set configured for model '{model_name}'.")

    overrides = {
        "prompt": "Prompt",
        "negative_prompt": "Negative Prompt",
        "steps": "Sampling Steps",
        "cfg_scale": "CFG Scale",
        "seed": "Seed",
        "sampler": "Sampling Method"
    }
    
    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

    for _ in range(7):
        payload.insert(39, None)

    # 6. EXECUTE GENERATION
    try:
        gen_payload = {"data": payload}
        gen_resp = requests.post(f"{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
        # We strip leading slash from raw_path to avoid double slashes in the proxy URL if we want, 
        # but the current endpoint handles it.
        proxy_url = f"http://localhost:{config.PORT}/file{raw_path}"
        
        # Download the image and convert to base64 for the AI's vision
        img_resp = 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\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)}")