"""
cosmos → gemini cross-org demo.

demonstrates:
1. registering cosmos and gemini agents in the registry
2. resolving gemini@google via registry lookup
3. sending an a2a envelope from cosmos to gemini
4. gemini processing and responding with structured content

run this script with the fastapi app running on localhost:8000.
"""

import asyncio
import uuid
from typing import Any, Dict

import httpx

from synqed.agent_email.addressing import AgentId
from synqed.agent_email.inbox.api import LocalAgentRuntime, register_agent_runtime
from synqed.agent_email.registry.api import get_registry
from synqed.agent_email.registry.models import AgentRegistryEntry


# configuration
BASE_URL = "http://localhost:8000"
COSMOS_AGENT_ID = "agent://futurehouse/cosmos"
COSMOS_EMAIL = "cosmos@futurehouse"
GEMINI_AGENT_ID = "agent://google/gemini"
GEMINI_EMAIL = "gemini@google"


class CosmosRuntime(LocalAgentRuntime):
    """
    mock runtime for cosmos agent.
    
    cosmos analyzes startup ideas and sends them to other agents for review.
    """
    
    async def handle_a2a_envelope(
        self,
        sender: str,
        recipient: str,
        envelope: Dict[str, Any],
    ) -> Dict[str, Any] | None:
        """handle incoming envelope (cosmos mostly sends, doesn't receive much)."""
        print(f"\n[cosmos] received message from {sender}")
        print(f"[cosmos] content: {envelope.get('content', '')}")
        
        # cosmos could respond, but for this demo it's the initiator
        return None


class GeminiRuntime(LocalAgentRuntime):
    """
    mock runtime for gemini agent.
    
    gemini provides structured analysis of startup ideas.
    """
    
    async def handle_a2a_envelope(
        self,
        sender: str,
        recipient: str,
        envelope: Dict[str, Any],
    ) -> Dict[str, Any] | None:
        """handle incoming envelope and generate analysis."""
        print(f"\n[gemini] received message from {sender}")
        print(f"[gemini] thread_id: {envelope.get('thread_id', 'unknown')}")
        print(f"[gemini] content: {envelope.get('content', '')}")
        
        # extract the startup idea from content
        content = envelope.get("content", "")
        
        # generate structured analysis (mock)
        analysis = f"""
structured analysis of the proposal:

1. problem statement: {content[:100]}...
2. market potential: high demand in academic and biotech sectors
3. technical feasibility: requires robust nlp and reasoning capabilities
4. competitive advantage: integration with a2a protocol for agent collaboration
5. recommendation: proceed with mvp development

this analysis was generated by gemini@google via the agent email layer.
        """.strip()
        
        # create response envelope
        response = {
            "thread_id": envelope.get("thread_id", str(uuid.uuid4())),
            "role": "assistant",
            "content": analysis,
            "tool_calls": [],
        }
        
        print(f"[gemini] sending response with analysis")
        
        return response


async def setup_local_runtimes() -> None:
    """register local runtimes for cosmos and gemini."""
    print("\n=== setting up local agent runtimes ===")
    
    cosmos_runtime = CosmosRuntime()
    gemini_runtime = GeminiRuntime()
    
    register_agent_runtime(COSMOS_AGENT_ID, cosmos_runtime)
    register_agent_runtime(GEMINI_AGENT_ID, gemini_runtime)
    
    print(f"registered cosmos runtime: {COSMOS_AGENT_ID}")
    print(f"registered gemini runtime: {GEMINI_AGENT_ID}")


async def register_agents_via_api() -> None:
    """register cosmos and gemini in the registry via http api."""
    print("\n=== registering agents in the registry ===")
    
    async with httpx.AsyncClient() as client:
        # register cosmos
        cosmos_entry = {
            "email_like": COSMOS_EMAIL,
            "inbox_url": f"{BASE_URL}/v1/a2a/inbox",
            "capabilities": ["a2a/1.0", "startup-analysis"],
            "metadata": {"org": "Future House", "description": "Research agent for biomedical literature"},
        }
        
        response = await client.post(
            f"{BASE_URL}/v1/agents",
            json=cosmos_entry,
        )
        response.raise_for_status()
        print(f"registered cosmos: {response.json()['agent_id']}")
        
        # register gemini
        gemini_entry = {
            "email_like": GEMINI_EMAIL,
            "inbox_url": f"{BASE_URL}/v1/a2a/inbox",
            "capabilities": ["a2a/1.0", "general-reasoning", "analysis"],
            "metadata": {"org": "Google", "description": "Gemini LLM agent"},
        }
        
        response = await client.post(
            f"{BASE_URL}/v1/agents",
            json=gemini_entry,
        )
        response.raise_for_status()
        print(f"registered gemini: {response.json()['agent_id']}")


async def cosmos_sends_to_gemini() -> None:
    """demonstrate cosmos resolving and sending to gemini."""
    print("\n=== cosmos → gemini demo flow ===")
    
    async with httpx.AsyncClient() as client:
        # step 1: cosmos only knows "gemini@google" as a string
        print(f"\n[cosmos] wants to send message to: {GEMINI_EMAIL}")
        
        # step 2: resolve gemini@google via registry
        print(f"[cosmos] resolving {GEMINI_EMAIL} via registry...")
        response = await client.get(f"{BASE_URL}/v1/agents/by-email/{GEMINI_EMAIL}")
        response.raise_for_status()
        gemini_entry = response.json()
        
        print(f"[cosmos] resolved to agent_id: {gemini_entry['agent_id']}")
        print(f"[cosmos] inbox_url: {gemini_entry['inbox_url']}")
        print(f"[cosmos] capabilities: {gemini_entry['capabilities']}")
        
        # step 3: build a2a envelope
        thread_id = str(uuid.uuid4())
        envelope = {
            "thread_id": thread_id,
            "role": "user",
            "content": (
                "analyze the following startup idea: "
                "we are building an ai agent that can read and analyze "
                "biomedical research papers, extract key findings, and "
                "collaborate with other agents to answer complex research questions. "
                "the agent uses the a2a protocol for inter-agent communication "
                "and can be deployed across different organizations."
            ),
            "tool_calls": [],
        }
        
        # step 4: send to gemini's inbox
        inbox_request = {
            "sender": COSMOS_AGENT_ID,
            "recipient": gemini_entry["agent_id"],
            "message": envelope,
            "token": None,  # no auth for demo
        }
        
        print(f"\n[cosmos] sending a2a envelope to gemini's inbox...")
        response = await client.post(
            gemini_entry["inbox_url"],
            json=inbox_request,
        )
        response.raise_for_status()
        inbox_response = response.json()
        
        # step 5: display response
        print(f"\n=== inbox response ===")
        print(f"status: {inbox_response['status']}")
        print(f"message_id: {inbox_response.get('message_id', 'N/A')}")
        
        if inbox_response.get("response_envelope"):
            print(f"\n=== gemini's response envelope ===")
            resp_env = inbox_response["response_envelope"]
            print(f"thread_id: {resp_env.get('thread_id', 'N/A')}")
            print(f"role: {resp_env.get('role', 'N/A')}")
            print(f"\ncontent:\n{resp_env.get('content', 'N/A')}")
        
        if inbox_response.get("error"):
            print(f"\nerror: {inbox_response['error']}")


async def demo_in_process() -> None:
    """
    run the full demo in-process (without needing separate server).
    
    uses direct registry and runtime access instead of http calls.
    """
    print("\n" + "="*70)
    print("COSMOS → GEMINI DEMO (In-Process)")
    print("="*70)
    
    # setup local runtimes directly
    await setup_local_runtimes()
    
    # register agents directly in registry
    print("\n=== registering agents directly in registry ===")
    registry = get_registry()
    
    cosmos_entry = AgentRegistryEntry(
        agent_id=COSMOS_AGENT_ID,
        email_like=COSMOS_EMAIL,
        inbox_url=f"{BASE_URL}/v1/a2a/inbox",
        capabilities=["a2a/1.0", "startup-analysis"],
        metadata={"org": "Future House"},
    )
    registry.register(cosmos_entry)
    print(f"registered cosmos: {COSMOS_AGENT_ID}")
    
    gemini_entry = AgentRegistryEntry(
        agent_id=GEMINI_AGENT_ID,
        email_like=GEMINI_EMAIL,
        inbox_url=f"{BASE_URL}/v1/a2a/inbox",
        capabilities=["a2a/1.0", "general-reasoning"],
        metadata={"org": "Google"},
    )
    registry.register(gemini_entry)
    print(f"registered gemini: {GEMINI_AGENT_ID}")
    
    # demonstrate resolution and messaging
    print(f"\n[cosmos] resolving {GEMINI_EMAIL}...")
    resolved_entry = registry.get_by_email(GEMINI_EMAIL)
    print(f"[cosmos] resolved to: {resolved_entry.agent_id}")
    print(f"[cosmos] inbox: {resolved_entry.inbox_url}")
    
    # create envelope
    thread_id = str(uuid.uuid4())
    envelope = {
        "thread_id": thread_id,
        "role": "user",
        "content": (
            "analyze the following startup idea: "
            "we are building an ai agent that can read and analyze "
            "biomedical research papers, extract key findings, and "
            "collaborate with other agents to answer complex research questions."
        ),
        "tool_calls": [],
    }
    
        # send directly to runtime (simulating inbox endpoint)
    from synqed.agent_email.inbox.api import get_agent_runtime
    
    gemini_runtime = get_agent_runtime(GEMINI_AGENT_ID)
    if gemini_runtime:
        print(f"\n[cosmos] sending envelope to gemini runtime...")
        response_envelope = await gemini_runtime.handle_a2a_envelope(
            sender=COSMOS_AGENT_ID,
            recipient=GEMINI_AGENT_ID,
            envelope=envelope,
        )
        
        if response_envelope:
            print(f"\n=== gemini's response ===")
            print(f"thread_id: {response_envelope.get('thread_id', 'N/A')}")
            print(f"\ncontent:\n{response_envelope.get('content', 'N/A')}")
    else:
        print("[error] gemini runtime not found")
    
    print("\n" + "="*70)
    print("DEMO COMPLETE")
    print("="*70)


async def demo_via_http() -> None:
    """
    run the full demo via http api.
    
    requires the demo server to be running on localhost:8000.
    start the server with: python demos/demo_server.py
    """
    print("\n" + "="*70)
    print("COSMOS → GEMINI DEMO (via HTTP)")
    print("="*70)
    
    try:
        # note: local runtimes are registered in the demo_server.py process,
        # not here. the demo script only makes http requests.
        
        # register agents via http
        await register_agents_via_api()
        
        # run the main demo flow
        await cosmos_sends_to_gemini()
        
        print("\n" + "="*70)
        print("DEMO COMPLETE")
        print("="*70)
        
    except httpx.ConnectError:
        print("\n[error] could not connect to http://localhost:8000")
        print("please start the demo server first with:")
        print("  python demos/demo_server.py")
        print("\nalternatively, run the in-process demo:")
        print("  python demos/cosmos_gemini_demo.py --in-process")


async def main() -> None:
    """main entry point."""
    import sys
    
    if len(sys.argv) > 1 and sys.argv[1] == "--in-process":
        # run in-process demo
        await demo_in_process()
    else:
        # run via http (requires server to be running)
        await demo_via_http()


if __name__ == "__main__":
    asyncio.run(main())

