api/websocket.py

"""
WebSocket endpoint — receives natural-language commands, returns results.

Protocol:
  Client → Server:  {"command": "salary ko 10% badhao"}
  Server → Client:  {"status": "success|error|unresolved", "message": "...", "diff": {...}, "suggestions": [...]}
"""
from __future__ import annotations

import asyncio
from fastapi import WebSocket, WebSocketDisconnect

from core.parser.fallback import parse_intent_hybrid
from core.validator import validate_intent
from core.execution.mvcc_executor import execute_mvcc_command
from core.column_registry import column_registry
from services import session_manager, audit_service


async def websocket_handler(ws: WebSocket, session_id: str) -> None:
    await ws.accept()

    # Verify session exists
    meta = session_manager.get(session_id)
    if meta is None:
        await ws.send_json({
            "status": "error",
            "message": "Session nahi mila. Pehle file upload karo.",
        })
        await ws.close()
        return

    try:
        while True:
            # ── Receive command ─────────────────────────────────────
            data = await ws.receive_json()
            command: str = data.get("command", "").strip()
            if not command:
                await ws.send_json({
                    "status": "error",
                    "message": "Empty command",
                })
                continue

            # ── Parse intent ────────────────────────────────────────
            intent = parse_intent_hybrid(session_id, command)

            if intent is None:
                # Unresolved — suggest available columns
                cols = column_registry.get_columns(session_id)
                await ws.send_json({
                    "status": "unresolved",
                    "message": "Command samajh nahi aaya. Column ka naam check karo.",
                    "suggestions": cols,
                })
                continue

            # ── Validate ────────────────────────────────────────────
            error = validate_intent(session_id, intent)
            if error:
                await ws.send_json({
                    "status": "error",
                    "message": error,
                })
                continue

            # ── Execute (offload to thread for large files) ─────────
            try:
                result = await asyncio.to_thread(execute_mvcc_command, session_id, intent)
            except Exception as exc:
                await ws.send_json({
                    "status": "error",
                    "message": f"Execution error: {exc}",
                })
                continue

            # ── Audit log ───────────────────────────────────────────
            await audit_service.log_command(
                session_id=session_id,
                command=command,
                intent=intent,
                diff=result.get("diff"),
            )

            # Touch session
            session_manager.touch(session_id)

            # ── Send result ─────────────────────────────────────────
            await ws.send_json({
                "status": "success",
                "message": result["message"],
                "diff": result.get("diff", {}),
            })

    except WebSocketDisconnect:
        pass  # Client disconnected, clean exit