import os
import json
import gradio as gr
from fastapi import FastAPI, Header, HTTPException, Request
from fastapi.responses import JSONResponse
from gradio.routes import mount_gradio_app

# Import your existing Gradio app
from gradio_app import create_gradio_app

# ========= Configuration =========
PORT = int(os.environ.get("PORT", 7860))
MCP_BEARER = os.getenv("MCP_BEARER", "")  # définir dans Spaces > Settings > Variables and secrets

# ========= App FastAPI (parent) =========
api = FastAPI(title="MCP + Gradio on HF Spaces")

@api.get("/health")
def health():
    return {"ok": True}

def _check_auth(authorization: str | None):
    """Vérifie le header Authorization: Bearer <token> si MCP_BEARER est défini."""
    if not MCP_BEARER:
        return
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
    token = authorization.split(" ", 1)[1]
    if token != MCP_BEARER:
        raise HTTPException(status_code=401, detail="Unauthorized")

def handle_mcp_request(payload: dict) -> dict:
    """Handle MCP JSON-RPC 2.0 requests."""
    # Validate JSON-RPC 2.0 format
    if payload.get("jsonrpc") != "2.0":
        return {
            "jsonrpc": "2.0",
            "id": payload.get("id"),
            "error": {
                "code": -32600,
                "message": "Invalid Request",
                "data": "Missing or invalid jsonrpc version"
            }
        }
    
    request_id = payload.get("id")
    method = payload.get("method")
    params = payload.get("params", {})
    
    if not method:
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": -32600,
                "message": "Invalid Request",
                "data": "Missing method"
            }
        }
    
    # Handle MCP standard methods
    if method == "initialize":
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {
                        "listChanged": False
                    },
                    "resources": {
                        "subscribe": False,
                        "listChanged": False
                    }
                },
                "serverInfo": {
                    "name": "agricultural-mcp-server",
                    "version": "1.0.0"
                }
            }
        }
    
    elif method == "tools/list":
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "tools": [
                    {
                        "name": "analyze_weed_pressure",
                        "description": "Analyze weed pressure trends using IFT herbicide data",
                        "inputSchema": {
                            "type": "object",
                            "properties": {
                                "years": {
                                    "type": "array",
                                    "items": {"type": "integer"},
                                    "description": "Years to analyze"
                                },
                                "plots": {
                                    "type": "array", 
                                    "items": {"type": "string"},
                                    "description": "Plot names to analyze"
                                }
                            }
                        }
                    },
                    {
                        "name": "predict_future_pressure",
                        "description": "Predict future weed pressure for target years",
                        "inputSchema": {
                            "type": "object",
                            "properties": {
                                "target_years": {
                                    "type": "array",
                                    "items": {"type": "integer"},
                                    "description": "Years to predict"
                                },
                                "max_ift": {
                                    "type": "number",
                                    "description": "Maximum IFT threshold for sensitive crops"
                                }
                            }
                        }
                    },
                    {
                        "name": "analyze_crop_rotation",
                        "description": "Analyze crop rotation impact on weed pressure",
                        "inputSchema": {
                            "type": "object",
                            "properties": {}
                        }
                    }
                ]
            }
        }
    
    elif method == "tools/call":
        tool_name = params.get("name")
        tool_args = params.get("arguments", {})
        
        if tool_name == "analyze_weed_pressure":
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {
                    "content": [
                        {
                            "type": "text",
                            "text": f"Analyse de la pression adventices pour les années {tool_args.get('years', [])} et parcelles {tool_args.get('plots', [])}\n\nCette fonction analyserait normalement les données IFT herbicides de votre dataset agricultural."
                        }
                    ]
                }
            }
        
        elif tool_name == "predict_future_pressure":
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {
                    "content": [
                        {
                            "type": "text",
                            "text": f"Prédiction de pression adventices pour {tool_args.get('target_years', [])} avec seuil IFT max {tool_args.get('max_ift', 1.0)}\n\nCette fonction utiliserait vos modèles d'apprentissage automatique pour prédire les futures pressions."
                        }
                    ]
                }
            }
        
        elif tool_name == "analyze_crop_rotation":
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "result": {
                    "content": [
                        {
                            "type": "text",
                            "text": "Analyse de l'impact des rotations culturales sur la pression adventices\n\nCette fonction analyserait les patterns de rotation dans votre dataset."
                        }
                    ]
                }
            }
        
        else:
            return {
                "jsonrpc": "2.0",
                "id": request_id,
                "error": {
                    "code": -32601,
                    "message": "Method not found",
                    "data": f"Unknown tool: {tool_name}"
                }
            }
    
    elif method == "resources/list":
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "resources": [
                    {
                        "uri": "agricultural://dataset/summary",
                        "name": "Agricultural Dataset Summary",
                        "description": "Summary of the Kerguéhennec experimental station dataset",
                        "mimeType": "text/plain"
                    }
                ]
            }
        }
    
    else:
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": -32601,
                "message": "Method not found",
                "data": f"Unknown method: {method}"
            }
        }

# ========= Endpoint MCP (conforme JSON-RPC 2.0) =========
@api.post("/mcp")
async def mcp_endpoint(request: Request, authorization: str | None = Header(None)):
    _check_auth(authorization)
    try:
        payload = await request.json()
    except Exception:
        return JSONResponse(
            status_code=400,
            content={
                "jsonrpc": "2.0",
                "id": None,
                "error": {
                    "code": -32700,
                    "message": "Parse error",
                    "data": "Invalid JSON"
                }
            }
        )
    
    response = handle_mcp_request(payload)
    return JSONResponse(response)

# ========= UI Gradio =========
# Use your existing comprehensive agricultural analysis interface
demo = create_gradio_app()

# Monte Gradio sous la racine "/"
app = mount_gradio_app(api, demo, path="/")

# ========= Entrée (pour exécution locale éventuelle) =========
if __name__ == "__main__":
    # En local uniquement ; sur Spaces, le runner est géré par la plateforme.
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=PORT)

# ========= Tests curl (exemples conformes JSON-RPC 2.0) =========
# Healthcheck (public, GET)
# curl -s https://hackathoncra-mcp.hf.space/health

# Test d'initialisation MCP
# curl -s -X POST https://hackathoncra-mcp.hf.space/mcp \
#   -H "Content-Type: application/json" \
#   -d '{"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}}}'

# Liste des outils disponibles
# curl -s -X POST https://hackathoncra-mcp.hf.space/mcp \
#   -H "Content-Type: application/json" \
#   -d '{"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}}'

# Appel d'outil - Analyse pression adventices
# curl -s -X POST https://hackathoncra-mcp.hf.space/mcp \
#   -H "Content-Type: application/json" \
#   -d '{"jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": {"name": "analyze_weed_pressure", "arguments": {"years": [2020, 2021, 2022], "plots": ["P1", "P2"]}}}'

# Appel d'outil - Prédiction future
# curl -s -X POST https://hackathoncra-mcp.hf.space/mcp \
#   -H "Content-Type: application/json" \
#   -d '{"jsonrpc": "2.0", "id": 4, "method": "tools/call", "params": {"name": "predict_future_pressure", "arguments": {"target_years": [2025, 2026], "max_ift": 1.0}}}'

# Appel d'outil - Analyse rotation
# curl -s -X POST https://hackathoncra-mcp.hf.space/mcp \
#   -H "Content-Type: application/json" \
#   -d '{"jsonrpc": "2.0", "id": 5, "method": "tools/call", "params": {"name": "analyze_crop_rotation", "arguments": {}}}'

# Liste des ressources
# curl -s -X POST https://hackathoncra-mcp.hf.space/mcp \
#   -H "Content-Type: application/json" \
#   -d '{"jsonrpc": "2.0", "id": 6, "method": "resources/list", "params": {}}'

# Avec authentification Bearer (si MCP_BEARER défini)
# curl -s -X POST https://hackathoncra-mcp.hf.space/mcp \
#   -H "Authorization: Bearer VOTRE_TOKEN" \
#   -H "Content-Type: application/json" \
#   -d '{"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}}}'