app/main.py

"""
ByteAstra backend — FastAPI application entry point and router setup.
"""
from __future__ import annotations
import torch
for i in range(1, 8):
    setattr(torch, f"int{i}", torch.int8)

import torch.utils._pytree
if not hasattr(torch.utils._pytree, "register_constant"):
    torch.utils._pytree.register_constant = lambda cls: cls

import json
import logging
from contextlib import asynccontextmanager
from typing import AsyncGenerator

from fastapi import FastAPI, Depends, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pymongo.database import Database as DBSession
from sse_starlette.sse import EventSourceResponse

from app.config import get_settings
from app.database import get_db, init_db
from app.schemas import (
    ChatRequest, ChatResponse, SessionCreate, SessionOut, SessionSummary, MessageOut
)
from app.services import sessions as session_svc
from app.services import engine
from app.services.sessions import parse_citations

logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(name)s | %(message)s")
logger = logging.getLogger(__name__)
settings = get_settings()


# ── Lifespan ───────────────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Startup and shutdown events."""
    logger.info("ByteAstra backend starting up...")
    init_db()
    logger.info("Database initialised.")
    yield
    logger.info("ByteAstra backend shutting down.")


# ── App ────────────────────────────────────────────────────────────────────────

app = FastAPI(
    title="ByteAstra API",
    description="Syllabus-grounded AI tutoring engine for Indian professional curricula.",
    version="0.1.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.cors_origins_list,
    allow_origin_regex=r"https://.*\.vercel\.app",
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── Health ─────────────────────────────────────────────────────────────────────

@app.get("/health", tags=["system"])
async def health():
    from app.services.llm import check_connection
    is_online = await check_connection()
    return {
        "status": "ok",
        "version": "0.1.0",
        "engine_mode": "hybrid" if is_online else "fallback"
    }


# ── Sessions ───────────────────────────────────────────────────────────────────

@app.post("/sessions", response_model=SessionOut, tags=["sessions"])
def create_session(body: SessionCreate, db: DBSession = Depends(get_db)):
    """Create a new tutoring session for a given domain."""
    sess = session_svc.create_session(db, domain=body.domain, title=body.title or "New Session", user_id=body.user_id)
    return _session_to_out(sess)


@app.get("/sessions", response_model=list[SessionSummary], tags=["sessions"])
def list_sessions(user_id: str | None = None, db: DBSession = Depends(get_db)):
    """List all sessions, optionally filtered by user_id."""
    return session_svc.list_sessions(db, user_id=user_id)


@app.get("/sessions/{session_id}", response_model=SessionOut, tags=["sessions"])
def get_session(session_id: str, db: DBSession = Depends(get_db)):
    sess = session_svc.get_session(db, session_id)
    if not sess:
        raise HTTPException(status_code=404, detail="Session not found")
    return _session_to_out(sess)


@app.delete("/sessions/{session_id}", tags=["sessions"])
def delete_session(session_id: str, db: DBSession = Depends(get_db)):
    deleted = session_svc.delete_session(db, session_id)
    if not deleted:
        raise HTTPException(status_code=404, detail="Session not found")
    return {"deleted": True}


@app.patch("/sessions/{session_id}/title", response_model=SessionSummary, tags=["sessions"])
def rename_session(session_id: str, title: str, db: DBSession = Depends(get_db)):
    sess = session_svc.update_session_title(db, session_id, title)
    if not sess:
        raise HTTPException(status_code=404, detail="Session not found")
    return sess


# ── Chat ───────────────────────────────────────────────────────────────────────

@app.post("/chat", tags=["chat"])
async def chat(body: ChatRequest, db: DBSession = Depends(get_db)):
    """
    Send a message and receive a tutoring response.
    If body.stream=True, returns an SSE stream.
    If body.stream=False, returns a JSON response (collects full stream internally).
    """
    sess = session_svc.get_session(db, body.session_id)
    if not sess:
        raise HTTPException(status_code=404, detail="Session not found")

    # Persist user message immediately
    session_svc.add_user_message(db, body.session_id, body.message)

    # Load conversation history for the prompt
    history = session_svc.get_recent_history(db, body.session_id)

    if body.stream:
        return EventSourceResponse(
            _stream_and_persist(db, body.session_id, sess.domain, body.message, history),
            media_type="text/event-stream",
        )
    else:
        # Collect stream into a single response
        full_content = []
        citations_data = []
        is_grounded = False
        debug_condensed_query: str | None = None

        async for event in engine.stream_answer(sess.domain, body.message, history):
            if event["type"] == "debug":
                debug_condensed_query = event.get("debug_condensed_query")
            elif event["type"] == "citations":
                citations_data = event.get("citations", [])
                is_grounded = event.get("is_grounded", False)
            elif event["type"] == "delta":
                full_content.append(event.get("content", ""))

        content_str = "".join(full_content)
        # Apply citation expansion to the stored content so MongoDB holds human-readable text
        domain_config = engine._load_domain_config(sess.domain)
        content_str = engine.make_citations_readable(content_str, domain_config)
        from app.schemas import Citation
        citations = [Citation(**c) for c in citations_data]
        msg = session_svc.add_assistant_message(db, body.session_id, content_str, citations, is_grounded)
        return ChatResponse(
            session_id=body.session_id,
            message=_msg_to_out(msg),
            debug_condensed_query=debug_condensed_query,
        )



async def _stream_and_persist(
    db: DBSession,
    session_id: str,
    domain: str,
    query: str,
    history: list[dict],
) -> AsyncGenerator[dict, None]:
    """
    Internal generator that streams engine events to SSE AND persists the
    final assistant message to the database once the stream completes.
    """
    from app.schemas import Citation
    full_content: list[str] = []
    citations_data: list[dict] = []
    is_grounded = False

    async for event in engine.stream_answer(domain, query, history):
        event_type = event.get("type")

        if event_type == "status":
            yield {"data": json.dumps(event)}

        elif event_type == "citations":
            citations_data = event.get("citations", [])
            is_grounded = event.get("is_grounded", False)
            yield {"data": json.dumps(event)}

        elif event_type == "delta":
            full_content.append(event.get("content", ""))
            yield {"data": json.dumps(event)}

        elif event_type == "done":
            # Persist the complete assistant message before emitting done.
            # Apply citation expansion so MongoDB stores human-readable text
            # (the client already received the expanded tokens via the delta events).
            content_str = "".join(full_content)
            domain_config = engine._load_domain_config(domain)
            content_str = engine.make_citations_readable(content_str, domain_config)
            citations = [Citation(**c) for c in citations_data]
            session_svc.add_assistant_message(db, session_id, content_str, citations, is_grounded)
            yield {"data": json.dumps(event)}

        elif event_type == "error":
            yield {"data": json.dumps(event)}


# ── Helpers ────────────────────────────────────────────────────────────────────

def _msg_to_out(msg) -> MessageOut:
    return MessageOut(
        id=msg.id,
        session_id=msg.session_id,
        role=msg.role.value,
        content=msg.content,
        citations=parse_citations(msg),
        is_grounded=msg.is_grounded,
        created_at=msg.created_at,
    )


def _session_to_out(sess) -> SessionOut:
    return SessionOut(
        id=sess.id,
        domain=sess.domain,
        title=sess.title,
        user_id=sess.user_id,
        subscription_tier=sess.subscription_tier,
        created_at=sess.created_at,
        updated_at=sess.updated_at,
        messages=[_msg_to_out(m) for m in sess.messages],
    )


# ── Authentication Routes ─────────────────────────────────────────────────────

from app.schemas import UserSignUp, UserLogIn, AuthResponse

@app.post("/signup", response_model=AuthResponse, tags=["auth"])
def signup(body: UserSignUp, db: DBSession = Depends(get_db)):
    import uuid
    import bcrypt
    from datetime import datetime, timezone
    
    # Check if email already exists
    existing = db.users.find_one({"email": body.email.lower().strip()})
    if existing:
        raise HTTPException(status_code=400, detail="Email already registered")
    
    # Hash password
    salt = bcrypt.gensalt()
    hashed_pwd = bcrypt.hashpw(body.password.encode('utf-8'), salt)
    
    user_id = str(uuid.uuid4())
    doc = {
        "_id": user_id,
        "name": body.name.strip(),
        "email": body.email.lower().strip(),
        "password": hashed_pwd.decode('utf-8'),
        "created_at": datetime.now(timezone.utc)
    }
    db.users.insert_one(doc)
    return AuthResponse(user_id=user_id, name=doc["name"], email=doc["email"])


@app.post("/login", response_model=AuthResponse, tags=["auth"])
def login(body: UserLogIn, db: DBSession = Depends(get_db)):
    import bcrypt
    
    user = db.users.find_one({"email": body.email.lower().strip()})
    if not user:
        raise HTTPException(status_code=400, detail="Invalid email or password")
    
    # Verify password
    hashed_pwd = user["password"].encode('utf-8')
    if not bcrypt.checkpw(body.password.encode('utf-8'), hashed_pwd):
        raise HTTPException(status_code=400, detail="Invalid email or password")
        
    return AuthResponse(user_id=user["_id"], name=user["name"], email=user["email"])


@app.get("/debug", tags=["system"])
async def debug():
    import subprocess
    import httpx
    
    # 1. Test local llama.cpp server
    llama_status = "unknown"
    llama_error = None
    try:
        async with httpx.AsyncClient(timeout=5.0) as client:
            r = await client.get("http://127.0.0.1:8001/v1/models")
            llama_status = f"HTTP {r.status_code}"
    except Exception as e:
        llama_status = "error"
        llama_error = str(e)
        
    # 2. Get processes from /proc
    processes = []
    import os
    try:
        for pid in os.listdir("/proc"):
            if pid.isdigit():
                try:
                    with open(f"/proc/{pid}/cmdline", "r") as f:
                        cmd = f.read().replace("\x00", " ").strip()
                        if cmd:
                            processes.append(f"PID {pid}: {cmd}")
                except Exception:
                    pass
    except Exception as e:
        processes = [f"Failed to read /proc: {e}"]
        
    # 3. Get ChromaDB count
    chroma_count = 0
    chroma_error = None
    try:
        from app.services.rag import get_or_create_collection
        col = get_or_create_collection("domain_ayurveda")
        chroma_count = col.count()
    except Exception as e:
        chroma_error = str(e)
        
    return {
        "llama_status": llama_status,
        "llama_error": llama_error,
        "chroma_count": chroma_count,
        "chroma_error": chroma_error,
        "processes": processes
    }


@app.get("/debug/run", tags=["system"])
async def debug_run(cmd: str):
    import subprocess
    try:
        res = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=10.0)
        return {
            "returncode": res.returncode,
            "stdout": res.stdout,
            "stderr": res.stderr
        }
    except Exception as e:
        return {"error": str(e)}