#!/usr/bin/env python3
"""
🎨 Figma MCP Server - Hébergé sur Hugging Face Spaces
Serveur MCP pour contrôler Figma via Claude/Cursor
"""
import gradio as gr
import asyncio
import json
import uuid
import websockets
import logging
from typing import Dict, Any, Optional, List
from PIL import Image
import base64
import io

# Configuration du logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Variables globales pour la connexion WebSocket
ws: Optional[websockets.WebSocketClientProtocol] = None
pending_requests: Dict[str, asyncio.Future] = {}
current_channel: Optional[str] = None
is_connected = False

async def connect_to_figma(port: int = 3055) -> bool:
    """Connecte au serveur WebSocket Figma"""
    global ws, is_connected
    try:
        uri = f"ws://localhost:{port}"
        ws = await websockets.connect(uri)
        is_connected = True
        logger.info(f"Connecté à Figma sur {uri}")
        
        # Démarre l'écoute des messages
        asyncio.create_task(listen_messages())
        return True
    except Exception as e:
        logger.error(f"Erreur de connexion à Figma: {e}")
        is_connected = False
        return False

async def listen_messages():
    """Écoute les messages du WebSocket"""
    global ws, is_connected
    try:
        async for message in ws:
            await handle_message(message)
    except websockets.exceptions.ConnectionClosed:
        logger.warning("Connexion WebSocket fermée")
        is_connected = False
    except Exception as e:
        logger.error(f"Erreur lors de l'écoute des messages: {e}")

async def handle_message(message: str):
    """Traite un message reçu de Figma"""
    try:
        data = json.loads(message)
        response_data = data.get("message", {})
        
        response_id = response_data.get("id")
        if response_id and response_id in pending_requests:
            future = pending_requests.pop(response_id)
            if "error" in response_data:
                future.set_exception(Exception(response_data["error"]))
            else:
                future.set_result(response_data.get("result"))
    except Exception as e:
        logger.error(f"Erreur lors du traitement du message: {e}")

async def send_command_to_figma(command: str, params: Dict[str, Any] = None) -> Any:
    """Envoie une commande à Figma et attend la réponse"""
    global ws, current_channel, is_connected
    
    if not is_connected or not ws:
        # Tente de se connecter
        await connect_to_figma()
        if not is_connected:
            raise Exception("Impossible de se connecter à Figma")
    
    if command != 'join' and not current_channel:
        raise Exception("Doit rejoindre un canal avant d'envoyer des commandes")
    
    request_id = str(uuid.uuid4())
    future = asyncio.Future()
    pending_requests[request_id] = future
    
    request = {
        "id": request_id,
        "type": "join" if command == 'join' else "message",
        "message": {
            "id": request_id,
            "command": command,
            "params": params or {}
        }
    }
    
    if command == 'join':
        request["channel"] = params.get("channel")
    else:
        request["channel"] = current_channel
    
    await ws.send(json.dumps(request))
    
    try:
        result = await asyncio.wait_for(future, timeout=30.0)
        return result
    except asyncio.TimeoutError:
        pending_requests.pop(request_id, None)
        raise Exception(f"Timeout lors de l'exécution de la commande {command}")

# === OUTILS MCP POUR FIGMA ===

def join_figma_channel(channel: str) -> str:
    """
    Rejoint un canal Figma pour commencer à communiquer.
    
    Args:
        channel (str): Nom du canal Figma à rejoindre
        
    Returns:
        str: Message de confirmation ou d'erreur
    """
    global current_channel
    try:
        result = asyncio.run(send_command_to_figma('join', {"channel": channel}))
        current_channel = channel
        return f"✅ Canal Figma rejoint avec succès: {channel}"
    except Exception as e:
        return f"❌ Erreur lors de la jointure du canal: {str(e)}"

def get_figma_document_info() -> str:
    """
    Récupère les informations détaillées du document Figma actuel.
    
    Returns:
        str: Informations du document en format JSON
    """
    try:
        result = asyncio.run(send_command_to_figma('get_document_info'))
        return json.dumps(result, indent=2, ensure_ascii=False)
    except Exception as e:
        return f"❌ Erreur: {str(e)}"

def get_figma_selection() -> str:
    """
    Récupère les informations sur la sélection actuelle dans Figma.
    
    Returns:
        str: Informations de la sélection en format JSON
    """
    try:
        result = asyncio.run(send_command_to_figma('get_selection'))
        return json.dumps(result, indent=2, ensure_ascii=False)
    except Exception as e:
        return f"❌ Erreur: {str(e)}"

def get_figma_node_info(node_id: str) -> str:
    """
    Récupère les informations détaillées d'un nœud spécifique dans Figma.
    
    Args:
        node_id (str): ID du nœud Figma à analyser
        
    Returns:
        str: Informations du nœud en format JSON
    """
    try:
        result = asyncio.run(send_command_to_figma('get_node_info', {"nodeId": node_id}))
        return json.dumps(result, indent=2, ensure_ascii=False)
    except Exception as e:
        return f"❌ Erreur: {str(e)}"

def create_figma_rectangle(x: str, y: str, width: str, height: str, name: str = "Rectangle", parent_id: str = "") -> str:
    """
    Crée un nouveau rectangle dans Figma.
    
    Args:
        x (str): Position X du rectangle
        y (str): Position Y du rectangle  
        width (str): Largeur du rectangle
        height (str): Hauteur du rectangle
        name (str): Nom optionnel du rectangle
        parent_id (str): ID du nœud parent (optionnel)
        
    Returns:
        str: Informations du rectangle créé
    """
    try:
        params = {
            "x": float(x),
            "y": float(y), 
            "width": float(width),
            "height": float(height),
            "name": name
        }
        if parent_id:
            params["parentId"] = parent_id
            
        result = asyncio.run(send_command_to_figma('create_rectangle', params))
        return f"✅ Rectangle créé: {json.dumps(result, indent=2)}"
    except Exception as e:
        return f"❌ Erreur: {str(e)}"

def create_figma_frame(x: str, y: str, width: str, height: str, name: str = "Frame", parent_id: str = "") -> str:
    """
    Crée un nouveau frame (conteneur) dans Figma.
    
    Args:
        x (str): Position X du frame
        y (str): Position Y du frame
        width (str): Largeur du frame
        height (str): Hauteur du frame
        name (str): Nom optionnel du frame
        parent_id (str): ID du nœud parent (optionnel)
        
    Returns:
        str: Informations du frame créé avec son ID
    """
    try:
        params = {
            "x": float(x),
            "y": float(y),
            "width": float(width), 
            "height": float(height),
            "name": name
        }
        if parent_id:
            params["parentId"] = parent_id
            
        result = asyncio.run(send_command_to_figma('create_frame', params))
        return f"✅ Frame créé: {json.dumps(result, indent=2)}"
    except Exception as e:
        return f"❌ Erreur: {str(e)}"

def create_figma_text(x: str, y: str, text: str, font_size: str = "14", name: str = "Text", parent_id: str = "") -> str:
    """
    Crée un nouvel élément texte dans Figma.
    
    Args:
        x (str): Position X du texte
        y (str): Position Y du texte
        text (str): Contenu du texte
        font_size (str): Taille de la police (défaut: 14)
        name (str): Nom optionnel de l'élément texte
        parent_id (str): ID du nœud parent (optionnel)
        
    Returns:
        str: Informations du texte créé
    """
    try:
        params = {
            "x": float(x),
            "y": float(y),
            "text": text,
            "fontSize": float(font_size),
            "name": name
        }
        if parent_id:
            params["parentId"] = parent_id
            
        result = asyncio.run(send_command_to_figma('create_text', params))
        return f"✅ Texte créé: {json.dumps(result, indent=2)}"
    except Exception as e:
        return f"❌ Erreur: {str(e)}"

def set_figma_fill_color(node_id: str, r: str, g: str, b: str, a: str = "1.0") -> str:
    """
    Définit la couleur de remplissage d'un nœud Figma.
    
    Args:
        node_id (str): ID du nœud à modifier
        r (str): Composant rouge (0-1)
        g (str): Composant vert (0-1)
        b (str): Composant bleu (0-1)
        a (str): Composant alpha/transparence (0-1, défaut: 1.0)
        
    Returns:
        str: Confirmation de la modification
    """
    try:
        params = {
            "nodeId": node_id,
            "color": {
                "r": float(r),
                "g": float(g), 
                "b": float(b),
                "a": float(a)
            }
        }
        result = asyncio.run(send_command_to_figma('set_fill_color', params))
        return f"✅ Couleur de remplissage définie: {json.dumps(result, indent=2)}"
    except Exception as e:
        return f"❌ Erreur: {str(e)}"

def move_figma_node(node_id: str, x: str, y: str) -> str:
    """
    Déplace un nœud Figma vers une nouvelle position.
    
    Args:
        node_id (str): ID du nœud à déplacer
        x (str): Nouvelle position X
        y (str): Nouvelle position Y
        
    Returns:
        str: Confirmation du déplacement
    """
    try:
        params = {
            "nodeId": node_id,
            "x": float(x),
            "y": float(y)
        }
        result = asyncio.run(send_command_to_figma('move_node', params))
        return f"✅ Nœud déplacé: {json.dumps(result, indent=2)}"
    except Exception as e:
        return f"❌ Erreur: {str(e)}"

def resize_figma_node(node_id: str, width: str, height: str) -> str:
    """
    Redimensionne un nœud Figma.
    
    Args:
        node_id (str): ID du nœud à redimensionner
        width (str): Nouvelle largeur
        height (str): Nouvelle hauteur
        
    Returns:
        str: Confirmation du redimensionnement
    """
    try:
        params = {
            "nodeId": node_id,
            "width": float(width),
            "height": float(height)
        }
        result = asyncio.run(send_command_to_figma('resize_node', params))
        return f"✅ Nœud redimensionné: {json.dumps(result, indent=2)}"
    except Exception as e:
        return f"❌ Erreur: {str(e)}"

def delete_figma_node(node_id: str) -> str:
    """
    Supprime un nœud dans Figma.
    
    Args:
        node_id (str): ID du nœud à supprimer
        
    Returns:
        str: Confirmation de la suppression
    """
    try:
        result = asyncio.run(send_command_to_figma('delete_node', {"nodeId": node_id}))
        return f"✅ Nœud supprimé: {node_id}"
    except Exception as e:
        return f"❌ Erreur: {str(e)}"

def get_figma_styles() -> str:
    """
    Récupère tous les styles du document Figma.
    
    Returns:
        str: Liste des styles en format JSON
    """
    try:
        result = asyncio.run(send_command_to_figma('get_styles'))
        return json.dumps(result, indent=2, ensure_ascii=False)
    except Exception as e:
        return f"❌ Erreur: {str(e)}"

def get_figma_components() -> str:
    """
    Récupère tous les composants locaux du document Figma.
    
    Returns:
        str: Liste des composants en format JSON
    """
    try:
        result = asyncio.run(send_command_to_figma('get_local_components'))
        return json.dumps(result, indent=2, ensure_ascii=False)
    except Exception as e:
        return f"❌ Erreur: {str(e)}"

# === INTERFACE GRADIO (SIMPLE) ===

def health_check() -> str:
    """
    Vérifie l'état de santé du serveur MCP Figma.
    
    Returns:
        str: État de la connexion
    """
    global is_connected, current_channel
    if is_connected:
        return f"✅ Serveur MCP Figma actif - Canal: {current_channel or 'Aucun'}"
    else:
        return "❌ Serveur MCP Figma non connecté"

# Interface Gradio simple pour le monitoring
with gr.Blocks(title="🎨 Figma MCP Server") as demo:
    gr.Markdown("# 🎨 Serveur MCP Figma")
    gr.Markdown("Serveur MCP hébergé sur Hugging Face Spaces pour contrôler Figma via Claude/Cursor")
    
    with gr.Row():
        status_btn = gr.Button("Vérifier l'état", variant="primary")
        status_output = gr.Textbox(label="État du serveur", interactive=False)
    
    gr.Markdown("""
    ## 🔗 Utilisation
    
    Ce serveur MCP expose les outils Figma suivants :
    - `join_figma_channel` - Rejoindre un canal Figma
    - `get_figma_document_info` - Infos du document
    - `get_figma_selection` - Sélection actuelle  
    - `create_figma_rectangle` - Créer un rectangle
    - `create_figma_frame` - Créer un frame
    - `create_figma_text` - Créer du texte
    - `set_figma_fill_color` - Définir une couleur
    - `move_figma_node` - Déplacer un élément
    - `resize_figma_node` - Redimensionner un élément
    - `delete_figma_node` - Supprimer un élément
    
    **URL du serveur MCP:** `https://votre-space.hf.space/gradio_api/mcp/sse`
    """)
    
    status_btn.click(health_check, outputs=[status_output])

if __name__ == "__main__":
    # Lance le serveur MCP selon les recommandations Gradio
    demo.launch(
        mcp_server=True,  # Active le serveur MCP 
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )