src/api/websocket.py

"""WebSocket endpoint for real-time grant query streaming."""

import asyncio
import json
import logging
import time
from typing import Dict, Any

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from uuid import uuid4

from src.logging.logger import get_logger
from src.analyzer.data_loader import load_current_grants, load_past_winners
from src.analyzer.chat.chat_tools import ChatTools
from src.analyzer.chat.query_router import route
from src.analyzer.llm_client import LLMClient
from src.analyzer.config import load_config

logger = get_logger()
router = APIRouter(prefix="/ws", tags=["websocket"])

# Global cache for initialized tools
_chat_tools: ChatTools | None = None
_llm_client: LLMClient | None = None


def get_chat_tools() -> ChatTools:
    """Lazy load chat tools on first use."""
    global _chat_tools
    if _chat_tools is None:
        try:
            from pathlib import Path

            snapshots_dir = Path("data/snapshots")
            history_xlsx = Path("data/IUK-141025-InnovateUKFundedProjects-FY2015-16topresent.xlsx")

            current = load_current_grants(snapshots_dir, limit=100)
            past = load_past_winners(history_xlsx=history_xlsx)
            _chat_tools = ChatTools(current, past)
            logging.info(f"Loaded ChatTools with {len(current)} current and {len(past)} past grants")
        except Exception as e:
            logging.error(f"Failed to initialize ChatTools: {e}")
            raise
    return _chat_tools


def get_llm_client() -> LLMClient:
    """Lazy load LLM client on first use."""
    global _llm_client
    if _llm_client is None:
        try:
            config = load_config()
            _llm_client = LLMClient(config)
            logging.info("Initialized LLM client for WebSocket streaming")
        except Exception as e:
            logging.error(f"Failed to initialize LLM client: {e}")
            raise
    return _llm_client


@router.websocket("/query")
async def websocket_query_endpoint(websocket: WebSocket):
    """
    WebSocket endpoint for real-time grant query streaming.

    Protocol:
        Client sends: {"query": "user question", "session_id": "optional-id"}
        Server streams:
            - {"type": "metadata", "session_id": "...", "query": "..."}
            - {"type": "intent", "intent": "search"}
            - {"type": "token", "token": "word"}
            - {"type": "citations", "citations": [...]}
            - {"type": "done", "latency_ms": 1234}
            - {"type": "error", "error": "message"}
    """
    await websocket.accept()
    session_id = str(uuid4())

    try:
        # Initialize tools and client
        tools = get_chat_tools()
        llm_client = get_llm_client()

        logging.info(f"WebSocket connection established: {session_id}")

        while True:
            # Receive query from client
            try:
                data = await websocket.receive_json()
            except WebSocketDisconnect:
                logging.info(f"WebSocket disconnected: {session_id}")
                break

            query = data.get("query", "")
            session_id = data.get("session_id", session_id)

            if not query:
                await websocket.send_json({
                    "type": "error",
                    "error": "Query is required"
                })
                continue

            start_time = time.time()

            try:
                # Send metadata
                await websocket.send_json({
                    "type": "metadata",
                    "session_id": session_id,
                    "query": query
                })

                # Route the query
                routed = route(query, use_llm=True)
                intent = str(routed.get("intent") or "general")
                args = routed.get("args") or {}

                logging.info(f"WebSocket {session_id}: Intent={intent}, Query={query}")

                # Send intent
                await websocket.send_json({
                    "type": "intent",
                    "intent": intent
                })

                citations_list = []

                # Handle different intents
                if intent in {"search", "list", "list_grants"}:
                    # Search for grants - non-streaming response
                    keyword = args.get("keyword") or args.get("query") or ""
                    results = tools.list_grants(keyword=keyword, limit=5)

                    if results:
                        answer_text = f"Found {len(results)} grants matching your query:\n"
                        for grant in results:
                            title = grant.get("title", "Unknown")
                            grant_id = grant.get("id") or grant.get("grant_id", "")
                            answer_text += f"\n- **{title}** (ID: {grant_id})"

                            if grant_id:
                                citations_list.append({
                                    "grant_id": grant_id,
                                    "title": title,
                                    "url": grant.get("url")
                                })
                    else:
                        answer_text = f"No grants found matching '{keyword}'."

                    # Stream the complete answer token by token
                    words = answer_text.split()
                    for word in words:
                        await websocket.send_json({
                            "type": "token",
                            "token": word + " "
                        })
                        await asyncio.sleep(0.01)  # Small delay for visual effect

                elif intent == "summarize" and args.get("grant_id"):
                    # Summarize a specific grant
                    grant_id = args.get("grant_id")
                    result = tools.summarize_grant(grant_id)
                    summary = result.get("summary_md", "No summary available")

                    # Stream summary token by token
                    words = summary.split()
                    for word in words:
                        await websocket.send_json({
                            "type": "token",
                            "token": word + " "
                        })
                        await asyncio.sleep(0.01)

                    citations_list.append({
                        "grant_id": grant_id,
                        "title": result.get("title", grant_id)
                    })

                elif intent == "compare" and args.get("grant_id_a") and args.get("grant_id_b"):
                    # Compare two grants
                    result = tools.compare_grants(args["grant_id_a"], args["grant_id_b"])
                    comparison = result.get("comparison_md", "Comparison unavailable")

                    # Stream comparison token by token
                    words = comparison.split()
                    for word in words:
                        await websocket.send_json({
                            "type": "token",
                            "token": word + " "
                        })
                        await asyncio.sleep(0.01)

                    citations_list.extend([
                        {"grant_id": args["grant_id_a"], "title": f"Grant {args['grant_id_a']}"},
                        {"grant_id": args["grant_id_b"], "title": f"Grant {args['grant_id_b']}"}
                    ])

                else:
                    # Default: Use LLM streaming for general queries
                    results = tools.list_grants(keyword=query, limit=5)

                    context = f"User query: {query}\n\n"
                    if results:
                        context += "Relevant grants:\n"
                        for grant in results:
                            title = grant.get("title", "Unknown")
                            grant_id = grant.get("id") or grant.get("grant_id", "")
                            context += f"- {title} (ID: {grant_id})\n"

                            if grant_id:
                                citations_list.append({
                                    "grant_id": grant_id,
                                    "title": title,
                                    "url": grant.get("url")
                                })

                    # Stream LLM response token by token
                    messages = [
                        {"role": "system", "content": "You are a UK grant analyst. Answer questions about grants concisely and accurately."},
                        {"role": "user", "content": context}
                    ]

                    # Get streaming response from LLM
                    stream_generator = llm_client.chat(
                        messages,
                        stream=True,
                        max_tokens=1200,
                        temperature=0.3
                    )

                    # Stream each token via WebSocket
                    for token in stream_generator:
                        await websocket.send_json({
                            "type": "token",
                            "token": token
                        })
                        await asyncio.sleep(0)  # Allow other tasks to run

                # Send citations
                if citations_list:
                    await websocket.send_json({
                        "type": "citations",
                        "citations": citations_list
                    })

                # Send completion
                latency_ms = int((time.time() - start_time) * 1000)
                await websocket.send_json({
                    "type": "done",
                    "latency_ms": latency_ms
                })

            except Exception as e:
                logging.error(f"Error processing WebSocket query: {e}", exc_info=True)
                await websocket.send_json({
                    "type": "error",
                    "error": str(e)
                })

    except Exception as e:
        logging.error(f"WebSocket connection error: {e}", exc_info=True)
    finally:
        logging.info(f"WebSocket connection closed: {session_id}")