main.py

import os
import time
import json
import logging
import asyncio
from contextlib import asynccontextmanager
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, FileResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware

from config import config
from schemas import ResearchRequest, ResearchResponse, HealthResponse, Source
from agent.agent import run_research
from cache import get_cached, set_cache, cache_stats, clear_cache
from database import (
    connect_db, close_db,
    save_research, get_history,
    get_by_id, search_history,
    delete_by_id
)

load_dotenv()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Starting up...")
    if not config.HF_TOKEN:
        raise RuntimeError("HF_TOKEN not set")
    await connect_db()
    logger.info("Ready.")
    yield
    await close_db()
    logger.info("Shut down.")


app = FastAPI(
    title="Research Agent API",
    description="Agentic AI that researches any topic and stores results in MongoDB.",
    version=config.API_VERSION,
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

app.mount("/static", StaticFiles(directory="static"), name="static")


@app.middleware("http")
async def log_requests(request: Request, call_next):
    start = time.time()
    response = await call_next(request)
    duration = round(time.time() - start, 2)
    logger.info(f"{request.method} {request.url.path} — {response.status_code} — {duration}s")
    return response


# ── Frontend ──────────────────────────────────────────────

@app.get("/", include_in_schema=False)
async def serve_frontend():
    return FileResponse("static/index.html")


# ── System ────────────────────────────────────────────────

@app.get("/health", response_model=HealthResponse, tags=["System"])
async def health():
    return HealthResponse(
        status="ok",
        model=config.MODEL_ID,
        version=config.API_VERSION,
        database="mongodb"
    )


@app.get("/cache/stats", tags=["System"])
async def get_cache_stats():
    """View cache usage stats."""
    return cache_stats()


@app.delete("/cache", tags=["System"])
async def flush_cache():
    """Clear all cached results."""
    count = clear_cache()
    return {"cleared": count}


# ── Agent ─────────────────────────────────────────────────

@app.post("/research", response_model=ResearchResponse, tags=["Agent"])
async def research(request: ResearchRequest):
    """Run the research agent. Returns cached result if available."""
    logger.info(f"Research: '{request.topic}' depth={request.depth}")

    # Check cache first
    cached = get_cached(request.topic, request.depth.value)
    if cached:
        logger.info(f"Cache hit for: '{request.topic}'")
        try:
            sources = [Source(**s) for s in cached.get("sources", [])]
        except Exception:
            sources = []
        return ResearchResponse(
            id=cached.get("id"),
            topic=request.topic,
            depth=request.depth.value,
            summary=cached.get("summary", ""),
            key_points=cached.get("key_points", []),
            sources=sources,
            follow_up_questions=cached.get("follow_up_questions", []),
            warning=cached.get("warning"),
            cached=True,
        )

    # Run agent
    try:
        result = run_research(topic=request.topic, depth=request.depth.value)
    except Exception as e:
        logger.error(f"Agent error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

    if "error" in result:
        raise HTTPException(status_code=500, detail=result["error"])

    try:
        sources = [Source(**s) for s in result.get("sources", [])]
    except Exception:
        sources = []

    # Save to MongoDB
    doc_id = await save_research(
        topic=request.topic,
        depth=request.depth.value,
        result=result
    )
    logger.info(f"Saved to MongoDB: {doc_id}")

    result["id"] = doc_id
    set_cache(request.topic, request.depth.value, result)

    return ResearchResponse(
        id=doc_id,
        topic=request.topic,
        depth=request.depth.value,
        summary=result.get("summary", ""),
        key_points=result.get("key_points", []),
        sources=sources,
        follow_up_questions=result.get("follow_up_questions", []),
        warning=result.get("warning"),
        cached=False,
    )


@app.post("/research/stream", tags=["Agent"])
async def research_stream(request: ResearchRequest):
    """Stream research progress as server-sent events."""

    async def generate():
        # Check cache
        cached = get_cached(request.topic, request.depth.value)
        if cached:
            yield f"data: {json.dumps({'type': 'cache_hit', 'message': 'Returning cached result...'})}\n\n"
            await asyncio.sleep(0.3)
            yield f"data: {json.dumps({'type': 'done', 'result': cached})}\n\n"
            return

        # Stream progress steps
        steps = [
            "Initializing agent...",
            "Searching the web for sources...",
            "Reading and analyzing pages...",
            "Synthesizing findings...",
            "Formatting report...",
        ]

        for i, step in enumerate(steps):
            yield f"data: {json.dumps({'type': 'progress', 'step': i+1, 'total': len(steps), 'message': step})}\n\n"
            await asyncio.sleep(0.5)

        # Run agent in thread so it doesn't block
        loop = asyncio.get_event_loop()
        try:
            result = await loop.run_in_executor(
                None,
                lambda: run_research(topic=request.topic, depth=request.depth.value)
            )
        except Exception as e:
            yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
            return

        if "error" in result:
            yield f"data: {json.dumps({'type': 'error', 'message': result['error']})}\n\n"
            return

        # Save to MongoDB
        doc_id = await save_research(
            topic=request.topic,
            depth=request.depth.value,
            result=result
        )
        result["id"] = doc_id
        set_cache(request.topic, request.depth.value, result)

        yield f"data: {json.dumps({'type': 'done', 'result': result})}\n\n"

    return StreamingResponse(generate(), media_type="text/event-stream")


# ── History ───────────────────────────────────────────────

@app.get("/history", tags=["History"])
async def history(limit: int = 20):
    return await get_history(limit=limit)


@app.get("/history/search/{query}", tags=["History"])
async def search(query: str):
    return await search_history(query)


@app.get("/history/{research_id}", tags=["History"])
async def get_one(research_id: str):
    doc = await get_by_id(research_id)
    if not doc:
        raise HTTPException(status_code=404, detail="Not found")
    return doc


@app.delete("/history/{research_id}", tags=["History"])
async def delete_one(research_id: str):
    deleted = await delete_by_id(research_id)
    if not deleted:
        raise HTTPException(status_code=404, detail="Not found")
    return {"deleted": True, "id": research_id}


# ── Error handler ─────────────────────────────────────────

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    logger.error(f"Unhandled error: {exc}")
    return JSONResponse(
        status_code=500,
        content={"detail": "Internal server error", "error": str(exc)}
    )