#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Demo A2A Agent for Issue #840 Testing.

This script creates a simple A2A agent with calculator and weather tools,
runs it on an available port, and registers it with ContextForge.

Usage:
    uv run python scripts/demo_a2a_agent.py

The agent supports queries like:
    - "calc: 5*10+2"
    - "weather: Dallas"

Press Ctrl+C to stop the server and unregister the agent.
"""

import atexit
import json
import os
import random
import signal
import socket
import sys
from contextlib import closing
import httpx
import jwt
import uvicorn
from fastapi import FastAPI, Request
from pydantic import BaseModel

# ============================================================================
# A2A Agent Implementation (from Issue #840)
# ============================================================================


def calculator(expression: str) -> str:
    """Evaluate a math expression safely using ast module."""
    import ast
    import operator

    # Supported operators for safe evaluation
    operators = {
        ast.Add: operator.add,
        ast.Sub: operator.sub,
        ast.Mult: operator.mul,
        ast.Div: operator.truediv,
        ast.USub: operator.neg,
        ast.UAdd: operator.pos,
    }

    def safe_eval(node):
        """Recursively evaluate an AST node safely."""
        if isinstance(node, ast.Constant):  # Python 3.8+
            if isinstance(node.value, (int, float)):
                return node.value
            raise ValueError(f"Invalid constant type: {type(node.value)}")
        elif isinstance(node, ast.BinOp):
            if type(node.op) not in operators:
                raise ValueError(f"Unsupported operator: {type(node.op).__name__}")
            left = safe_eval(node.left)
            right = safe_eval(node.right)
            return operators[type(node.op)](left, right)
        elif isinstance(node, ast.UnaryOp):
            if type(node.op) not in operators:
                raise ValueError(f"Unsupported operator: {type(node.op).__name__}")
            operand = safe_eval(node.operand)
            return operators[type(node.op)](operand)
        elif isinstance(node, ast.Expression):
            return safe_eval(node.body)
        else:
            raise ValueError(f"Unsupported expression type: {type(node).__name__}")

    try:
        # Parse the expression into an AST
        tree = ast.parse(expression, mode="eval")
        result = safe_eval(tree)
        return str(result)
    except (SyntaxError, ValueError) as e:
        return f"Error: {e}"
    except ZeroDivisionError:
        return "Error: Division by zero"
    except Exception as e:
        return f"Error: {e}"


def weather(city: str) -> str:
    """Mock weather lookup tool."""
    conditions = ["sunny", "rainy", "cloudy", "stormy"]
    temp = random.randint(10, 35)
    return f"The weather in {city} is {random.choice(conditions)}, {temp}C"


class SimpleAgent:
    """Simple A2A agent that routes queries to tools."""

    def __init__(self, name: str = "Agent"):
        self.name = name
        self.tools = {
            "calculator": calculator,
            "weather": weather,
        }

    def run(self, query: str) -> str:
        """Process a query and route to appropriate tool."""
        if "calc:" in query.lower():
            expr = query.lower().replace("calc:", "").strip()
            return self.tools["calculator"](expr)
        elif "weather:" in query.lower():
            city = query.lower().replace("weather:", "").strip()
            return self.tools["weather"](city.title())
        else:
            return f"{self.name} received: {query}. Try 'calc: 5*10' or 'weather: Dallas'"


# ============================================================================
# FastAPI Application
# ============================================================================

app = FastAPI(title="Demo A2A Agent", description="Calculator and Weather Agent for Issue #840")
agent = SimpleAgent("Demo-A2A-Agent")


class Response(BaseModel):
    """Response model for agent results."""

    response: str


@app.post("/run")
async def run_agent(request: Request) -> Response:
    """Execute a query against the agent.

    Supports multiple formats:
    - JSONRPC: {"jsonrpc": "2.0", "method": "...", "params": {"query": "..."}}
    - A2A protocol: {"parameters": {"query": "..."}}
    - Simple: {"query": "..."}
    """
    # Parse request body
    body = await request.body()
    try:
        body_dict = json.loads(body)
    except (json.JSONDecodeError, ValueError):
        return Response(response="Error: invalid JSON in request body")

    query_text = ""

    # Handle JSONRPC format (ContextForge sends this for agents with URLs ending in /)
    if "jsonrpc" in body_dict:
        params = body_dict.get("params", {})

        # Handle nested message structure from Admin UI test
        if "message" in params and isinstance(params["message"], dict):
            message_obj = params["message"]
            # Extract text from parts array
            if "parts" in message_obj and isinstance(message_obj["parts"], list):
                for part in message_obj["parts"]:
                    if isinstance(part, dict) and part.get("kind") == "text":
                        query_text = part.get("text", "")
                        break
            print("Extracted query from JSONRPC message.parts")
        else:
            # Simple query or message string
            query_text = params.get("query") or params.get("message", "")
            print("Extracted query from JSONRPC params")
    # Handle A2A protocol format
    elif "parameters" in body_dict and isinstance(body_dict["parameters"], dict):
        params = body_dict["parameters"]
        query_text = params.get("query") or params.get("message", "")
        print("Extracted query from A2A parameters")
    # Handle simple format
    elif "query" in body_dict:
        query_text = body_dict["query"]
        print("Extracted query from query field")
    elif "message" in body_dict:
        query_text = body_dict["message"]
        print("Extracted query from message field")

    if not query_text:
        query_text = "Hello"
        print("No query found, using default: Hello")

    response = agent.run(query_text)
    return Response(response=response)


@app.get("/health")
def health():
    """Health check endpoint."""
    return {"status": "healthy", "agent": agent.name}


# ============================================================================
# ContextForge Registration
# ============================================================================

# Configuration from environment with fallbacks for local development
CONTEXTFORGE_URL = os.environ.get("CONTEXTFORGE_URL", "http://localhost:8000")
JWT_SECRET = os.environ.get("JWT_SECRET_KEY", "my-test-key-but-now-longer-than-32-bytes")  # noqa: S105 - default for demo only
AGENT_ID = None


def create_jwt_token(username: str = os.environ.get("PLATFORM_ADMIN_EMAIL", "admin@example.com")) -> str:
    """Create a JWT token for ContextForge authentication."""
    import datetime

    payload = {
        "sub": username,
        "email": username,
        "iat": int(datetime.datetime.now(datetime.timezone.utc).timestamp()),
        "exp": int((datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(hours=1)).timestamp()),
        "iss": "mcpgateway",
        "aud": "mcpgateway-api",
        "teams": [],
    }
    return jwt.encode(payload, JWT_SECRET, algorithm="HS256")


def register_agent(port: int) -> str | None:
    """Register the A2A agent with ContextForge."""
    global AGENT_ID

    token = create_jwt_token()
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {token}",
    }

    agent_data = {
        "agent": {
            "name": "Demo Calculator Agent",
            "description": "Demo A2A Agent with calculator and weather tools (Issue #840)",
            "endpoint_url": f"http://localhost:{port}/run",
            "agent_type": "custom",
            "protocol_version": "1.0",
            "capabilities": {"tools": ["calculator", "weather"]},
            "config": {},
            "tags": ["demo", "calculator", "weather", "issue-840"],
        },
        "visibility": "public",
    }

    try:
        with httpx.Client(timeout=10) as client:
            response = client.post(f"{CONTEXTFORGE_URL}/a2a", headers=headers, json=agent_data)

            if response.status_code == 201:
                data = response.json()
                AGENT_ID = data.get("id")
                print(f"Registered A2A agent with ContextForge: {AGENT_ID}")
                print(f"  Name: {data.get('name')}")
                print(f"  Endpoint: {data.get('endpointUrl')}")
                return AGENT_ID
            else:
                print(f"Failed to register agent: {response.status_code}")
                print(f"  Response: {response.text}")
                return None
    except Exception as e:
        print(f"Error registering agent: {e}")
        return None


def unregister_agent():
    """Unregister the A2A agent from ContextForge."""
    if not AGENT_ID:
        return

    token = create_jwt_token()
    headers = {"Authorization": f"Bearer {token}"}

    try:
        with httpx.Client(timeout=10) as client:
            response = client.delete(f"{CONTEXTFORGE_URL}/a2a/{AGENT_ID}", headers=headers)
            if response.status_code in (200, 204):
                print(f"Unregistered A2A agent: {AGENT_ID}")
            else:
                print(f"Failed to unregister agent: {response.status_code}")
    except Exception as e:
        print(f"Error unregistering agent: {e}")


def find_available_port(start: int = 9100, end: int = 9200) -> int:
    """Find an available port in the given range."""
    for port in range(start, end):
        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
            if sock.connect_ex(("localhost", port)) != 0:
                return port
    raise RuntimeError(f"No available port found in range {start}-{end}")


# ============================================================================
# Main Entry Point
# ============================================================================


def main():
    """Run the demo A2A agent."""
    # Find available port
    port = find_available_port()
    print(f"\n{'='*60}")
    print("Demo A2A Agent for Issue #840")
    print(f"{'='*60}")
    print(f"Starting agent on port {port}...")
    print("\nSupported queries:")
    print("  - calc: 5*10+2")
    print("  - weather: Dallas")
    print("\nPress Ctrl+C to stop\n")

    # Register cleanup handler
    atexit.register(unregister_agent)

    def signal_handler(sig, frame):
        print("\nShutting down...")
        unregister_agent()
        sys.exit(0)

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    # Register with ContextForge
    register_agent(port)

    # Start the server
    uvicorn.run(app, host="0.0.0.0", port=port, log_level="info")


if __name__ == "__main__":
    main()
