backend/app/api/routes/chat.py

"""
Chat & Query Routes — Core API endpoints for text-to-SQL.
Includes SSE streaming for real-time pipeline feedback.
"""

import json
from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from app.api.schemas import (
    GenerateSQLRequest, QueryResult,
    ExecuteQueryRequest,
    ExplainRequest, ExplainResponse,
    InsightsRequest, InsightsResponse,
)

router = APIRouter(prefix="/api/v1", tags=["Query"])


def create_chat_router(orchestrator, auth_dep, cache, rate_limiter, tracer, explainer, insights_gen, anomaly_detector, safety_validator):
    """Factory to create chat router with injected dependencies."""

    @router.post("/generate-sql", response_model=QueryResult)
    def generate_sql(request: GenerateSQLRequest, current_user: dict = Depends(auth_dep)):
        """
        Generate SQL from natural language and optionally execute it.
        The full multi-agent pipeline runs here.
        """
        # Rate limiting
        user_key = f"rl:{current_user.get('sub', 'anon')}"
        if not rate_limiter.check(user_key):
            raise HTTPException(429, "Rate limit exceeded. Please wait a moment.")

        # Check cache
        cached = cache.get(request.question, current_user.get("tenant_id", "default"))
        if cached:
            cached["trace_id"] = "cached"
            return QueryResult(**cached)

        # Run multi-agent pipeline
        result = orchestrator.process_query(
            user_query=request.question,
            conversation_history=request.history,
            tenant_id=current_user.get("tenant_id", "default"),
            user_role=current_user.get("role", "viewer"),
        )

        # Trace the query
        tracer.trace_query(result)

        # Build response
        response_data = {
            "trace_id": result.get("trace_id", ""),
            "question": request.question,
            "intent": result.get("intent"),
            "sql": result.get("sanitized_sql") or result.get("generated_sql"),
            "sql_explanation": result.get("sql_explanation"),
            "message": result.get("friendly_message", ""),
            "data": result.get("query_results", []),
            "row_count": result.get("row_count", 0),
            "column_names": result.get("column_names", []),
            "execution_time_ms": result.get("execution_time_ms", 0),
            "chart_config": result.get("chart_config"),
            "chart_type": result.get("chart_type"),
            "insights": result.get("insights", []),
            "follow_ups": result.get("follow_up_questions", []),
            "error": result.get("error"),
        }

        # Cache successful results
        if not result.get("error") and result.get("query_results"):
            cache.set(request.question, response_data, current_user.get("tenant_id", "default"))

        return QueryResult(**response_data)

    @router.post("/chat/stream")
    async def chat_stream(request: GenerateSQLRequest, current_user: dict = Depends(auth_dep)):
        """
        Stream chat responses via Server-Sent Events (SSE).
        Each agent stage emits an event as it completes.
        """
        user_key = f"rl:{current_user.get('sub', 'anon')}"
        if not rate_limiter.check(user_key):
            raise HTTPException(429, "Rate limit exceeded.")

        async def event_generator():
            async for event in orchestrator.aprocess_query_streaming(
                user_query=request.question,
                conversation_history=request.history,
                tenant_id=current_user.get("tenant_id", "default"),
                user_role=current_user.get("role", "viewer"),
            ):
                event_type = event.get("type", "message")
                yield _sse_event(event_type, event)

        return StreamingResponse(
            event_generator(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
            },
        )

    @router.post("/execute-query", response_model=QueryResult)
    def execute_query(request: ExecuteQueryRequest, current_user: dict = Depends(auth_dep)):
        """Execute a user-provided SQL query (must pass safety validation)."""
        from app.agents.sql_validation import sql_validation_node

        # Validate the SQL
        validation_state = {"generated_sql": request.sql, "retry_count": 0, "trace_id": "manual"}
        validation_result = sql_validation_node(validation_state)

        if not validation_result.get("is_valid"):
            errors = validation_result.get("validation_errors", ["Unknown validation error"])
            raise HTTPException(400, f"SQL blocked by safety layer: {', '.join(errors)}")

        # Execute
        from app.agents.execution import execution_node
        exec_state = {**validation_result, "trace_id": "manual"}
        exec_result = execution_node(exec_state, orchestrator.db_pool)

        if exec_result.get("error"):
            raise HTTPException(400, exec_result["error"])

        return QueryResult(
            trace_id="manual",
            question="Manual SQL execution",
            sql=request.sql,
            message=f"Query executed successfully. {exec_result.get('row_count', 0)} rows returned.",
            data=exec_result.get("query_results", []),
            row_count=exec_result.get("row_count", 0),
            column_names=exec_result.get("column_names", []),
            execution_time_ms=exec_result.get("execution_time_ms", 0),
        )

    @router.post("/explain", response_model=ExplainResponse)
    def explain_sql(request: ExplainRequest, current_user: dict = Depends(auth_dep)):
        """Explain a SQL query in natural language."""
        explanation = explainer.explain(request.sql, request.result_count)
        return ExplainResponse(sql=request.sql, explanation=explanation)

    @router.post("/insights", response_model=InsightsResponse)
    def get_insights(request: InsightsRequest, current_user: dict = Depends(auth_dep)):
        """Generate auto-insights and anomaly detection for data."""
        insights = insights_gen.generate(request.data, request.query)
        anomalies = anomaly_detector.detect(request.data)
        return InsightsResponse(insights=insights, anomalies=anomalies)

    return router


def _sse_event(event_type: str, data: dict) -> str:
    """Format a Server-Sent Event."""
    payload = json.dumps({"type": event_type, **data}, default=str)
    return f"event: {event_type}\ndata: {payload}\n\n"