"""
FastAPI server for memory graph visualization and API.

Provides REST API endpoints for memory operations and serves
the interactive visualization interface.
"""
import asyncio
from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
from dotenv import load_dotenv
import os
import sys
from pathlib import Path
from typing import Optional, List, Dict, Any
from datetime import datetime

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent))
from memory import TemporalSemanticMemory

import logging

load_dotenv()
logging.basicConfig(level=logging.INFO)

app = FastAPI(title="Memory Graph API", version="1.0.0")

# Mount static files
app.mount("/static", StaticFiles(directory="web/static"), name="static")


class SearchRequest(BaseModel):
    """Request model for search endpoint."""
    query: str
    agent_id: str = "default"
    thinking_budget: int = 100
    top_k: int = 10
    mmr_lambda: float = 0.5
    trace: bool = False


class MemoryItem(BaseModel):
    """Single memory item for batch put."""
    content: str
    event_date: Optional[datetime] = None
    context: Optional[str] = None


class BatchPutRequest(BaseModel):
    """Request model for batch put endpoint."""
    agent_id: str
    items: List[MemoryItem]
    document_id: Optional[str] = None
    document_metadata: Optional[Dict[str, Any]] = None
    upsert: bool = False


class ThinkRequest(BaseModel):
    """Request model for think endpoint."""
    query: str
    agent_id: str = "default"
    thinking_budget: int = 50
    top_k: int = 10


class ThinkResponse(BaseModel):
    """Response model for think endpoint."""
    text: str
    based_on: Dict[str, List[Dict[str, Any]]]  # {"world": [...], "agent": [...], "opinion": [...]}
    new_opinions: List[str] = []  # List of newly formed opinions


memory = TemporalSemanticMemory()

@app.on_event("startup")
async def startup_event():
    """Initialize memory system on startup."""
    await memory.initialize()
    logging.info("Memory system initialized")

@app.on_event("shutdown")
async def shutdown_event():
    """Cleanup memory system on shutdown."""
    await memory.close()
    logging.info("Memory system closed")

@app.get("/")
async def index():
    """Serve the visualization page."""
    return FileResponse("web/templates/index.html")


@app.get("/api/graph")
async def api_graph(agent_id: Optional[str] = None, fact_type: Optional[str] = None):
    """Get graph data from database, optionally filtered by agent_id and fact_type."""
    try:
        data = await memory.get_graph_data(agent_id, fact_type)
        return data
    except Exception as e:
        import traceback
        error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        print(f"Error in /api/graph: {error_detail}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/search")
async def api_search(request: SearchRequest):
    """Run a search and return results with trace."""
    try:
        # Initialize memory system
        # Run search with tracing
        results, trace = await memory.search_async(
            agent_id=request.agent_id,
            query=request.query,
            thinking_budget=request.thinking_budget,
            top_k=request.top_k,
            enable_trace=request.trace,
            mmr_lambda=request.mmr_lambda
        )

        # Convert trace to dict
        trace_dict = trace.to_dict() if trace else None

        return {
            'results': results,
            'trace': trace_dict
        }
    except Exception as e:
        import traceback
        error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        print(f"Error in /api/search: {error_detail}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/world_search")
async def api_world_search(request: SearchRequest):
    """Search only world facts (general knowledge about the world)."""
    try:
        # Run search with fact_type filter for 'world'
        results, trace = await memory.search_async(
            agent_id=request.agent_id,
            query=request.query,
            thinking_budget=request.thinking_budget,
            top_k=request.top_k,
            enable_trace=request.trace,
            mmr_lambda=request.mmr_lambda,
            fact_type='world'
        )

        # Convert trace to dict
        trace_dict = trace.to_dict() if trace else None

        return {
            'results': results,
            'trace': trace_dict
        }
    except Exception as e:
        import traceback
        error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        print(f"Error in /api/world_search: {error_detail}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/agent_search")
async def api_agent_search(request: SearchRequest):
    """Search only agent facts (facts about what the agent did)."""
    try:
        # Run search with fact_type filter for 'agent'
        results, trace = await memory.search_async(
            agent_id=request.agent_id,
            query=request.query,
            thinking_budget=request.thinking_budget,
            top_k=request.top_k,
            enable_trace=request.trace,
            mmr_lambda=request.mmr_lambda,
            fact_type='agent'
        )

        # Convert trace to dict
        trace_dict = trace.to_dict() if trace else None

        return {
            'results': results,
            'trace': trace_dict
        }
    except Exception as e:
        import traceback
        error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        print(f"Error in /api/agent_search: {error_detail}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/opinion_search")
async def api_opinion_search(request: SearchRequest):
    """Search only opinion facts (agent's formed opinions and perspectives)."""
    try:
        # Run search with fact_type filter for 'opinion'
        results, trace = await memory.search_async(
            agent_id=request.agent_id,
            query=request.query,
            thinking_budget=request.thinking_budget,
            top_k=request.top_k,
            enable_trace=request.trace,
            mmr_lambda=request.mmr_lambda,
            fact_type='opinion'
        )

        # Convert trace to dict
        trace_dict = trace.to_dict() if trace else None

        return {
            'results': results,
            'trace': trace_dict
        }
    except Exception as e:
        import traceback
        error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        print(f"Error in /api/opinion_search: {error_detail}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/think")
async def api_think(request: ThinkRequest):
    """
    Think and formulate an answer using agent identity, world facts, and opinions.

    This endpoint:
    1. Retrieves agent facts (agent's identity)
    2. Retrieves world facts relevant to the query
    3. Retrieves existing opinions (agent's perspectives)
    4. Uses Groq LLM to formulate an answer
    5. Extracts and stores any new opinions formed
    6. Returns plain text answer, the facts used, and new opinions
    """
    try:
        # Use the memory system's think_async method
        result = await memory.think_async(
            agent_id=request.agent_id,
            query=request.query,
            thinking_budget=request.thinking_budget,
            top_k=request.top_k
        )

        return ThinkResponse(
            text=result["text"],
            based_on=result["based_on"],
            new_opinions=result.get("new_opinions", [])
        )

    except Exception as e:
        import traceback
        error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        print(f"Error in /api/think: {error_detail}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/agents")
async def api_agents():
    """Get list of available agents from database."""
    try:
        agent_list = await memory.list_agents()
        return {"agents": agent_list}
    except Exception as e:
        import traceback
        error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        print(f"Error in /api/agents: {error_detail}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/memories/batch")
async def api_batch_put(request: BatchPutRequest):
    """
    Store multiple memories in batch.

    This endpoint calls put_batch_async to efficiently store multiple memory items.
    Supports document tracking and upsert operations.

    Example request:
    {
        "agent_id": "user123",
        "items": [
            {"content": "Alice works at Google", "context": "work"},
            {"content": "Bob went hiking yesterday", "event_date": "2024-01-15T10:00:00Z"}
        ],
        "document_id": "conversation_123",
        "upsert": false
    }
    """
    try:
        # Validate agent_id - prevent writing to reserved agents
        RESERVED_AGENT_IDS = {"locomo"}
        if request.agent_id in RESERVED_AGENT_IDS:
            raise HTTPException(
                status_code=403,
                detail=f"Cannot write to reserved agent_id '{request.agent_id}'. Reserved agents: {', '.join(RESERVED_AGENT_IDS)}"
            )

        # Initialize memory system


        # Prepare contents for put_batch_async
        contents = []
        for item in request.items:
            content_dict = {"content": item.content}
            if item.event_date:
                content_dict["event_date"] = item.event_date
            if item.context:
                content_dict["context"] = item.context
            contents.append(content_dict)

        # Call put_batch_async
        result = await memory.put_batch_async(
            agent_id=request.agent_id,
            contents=contents,
            document_id=request.document_id,
            document_metadata=request.document_metadata,
            upsert=request.upsert
        )

        return {
            "success": True,
            "message": f"Successfully stored {len(contents)} memory items",
            "agent_id": request.agent_id,
            "document_id": request.document_id,
            "items_count": len(contents)
        }
    except Exception as e:
        import traceback
        error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        print(f"Error in /api/memories/batch: {error_detail}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/locomo")
async def api_locomo():
    """Get Locomo benchmark results."""
    import json
    try:
        results_path = Path(__file__).parent.parent / "benchmarks" / "locomo" / "benchmark_results.json"
        if not results_path.exists():
            raise HTTPException(status_code=404, detail="Benchmark results not found")

        with open(results_path, 'r') as f:
            data = json.load(f)
        return data
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="Benchmark results not found")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    import uvicorn
    print("\n" + "=" * 80)
    print("Memory Graph API Server")
    print("=" * 80)
    print("\nStarting server at http://localhost:8080")
    print("\nEndpoints:")
    print("  GET  /                    - Visualization UI")
    print("  GET  /api/graph           - Get graph data")
    print("  POST /api/search          - Run search with trace")
    print("  POST /api/memories/batch  - Store multiple memories in batch")
    print("  GET  /api/agents          - List available agents")
    print("\n" + "=" * 80 + "\n")

    uvicorn.run("server:app", host="0.0.0.0", port=8080, reload=True)