app/main.py

import os
from pathlib import Path
from contextlib import asynccontextmanager

from dotenv import load_dotenv

# Must run before any local import that reads env vars. Resolve explicitly so
# Electron, uv, and direct uvicorn launches all see backend/.env regardless of cwd.
load_dotenv(Path(__file__).resolve().parents[1] / ".env")

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from uvicorn.protocols.utils import ClientDisconnected
from fastapi.middleware.cors import CORSMiddleware

from app.routers import agent_control, annotations, citation_graph, control_room, health, library, projects, regions, sandbox, session, review, chat, draft, wiki, voice
from app.websockets.handlers import get_connection_manager, get_db, handle_event


@asynccontextmanager
async def lifespan(app: FastAPI):
    import asyncio
    import logging
    import os
    import time
    logger = logging.getLogger(__name__)
    if os.getenv("DEPLOYMENT_ENV", "desktop") == "demo":
        print("Demo mode: skipping Cognee bootstrap")
        yield
        return

    # Cerebras has no embeddings endpoint -> Cognee's default embedding model
    # (openai/text-embedding-3-large) would otherwise be routed through the
    # OPENAI_API_BASE override above and 404. Use a local, fully offline
    # embedding model instead (matches CLAUDE.md: Cognee never writes to cloud).
    os.environ.setdefault("EMBEDDING_PROVIDER", "fastembed")
    os.environ.setdefault("EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
    os.environ.setdefault("EMBEDDING_DIMENSIONS", "384")
    skill_root = str(Path(__file__).resolve().parent / "memory_skills")
    existing_skill_roots = os.environ.get("COGNEE_SKILL_SOURCE_ROOTS", "")
    if skill_root not in existing_skill_roots.split(os.pathsep):
        os.environ["COGNEE_SKILL_SOURCE_ROOTS"] = (
            skill_root if not existing_skill_roots else existing_skill_roots + os.pathsep + skill_root
        )
    
    # Disable multi-tenant access control for Cognee 1.2+ since StudyBuddy uses a single local user
    os.environ["ENABLE_BACKEND_ACCESS_CONTROL"] = "false"

    t0 = time.perf_counter()
    import cognee
    try:
        import importlib
        from app.services.cognee_bootstrap import configure_cognee_llm

        cognee_llm_client = importlib.import_module(
            "cognee.infrastructure.llm.structured_output_framework.litellm_instructor.llm.get_llm_client"
        )

        configure_cognee_llm(
            cognee.config,
            clear_llm_client_cache=cognee_llm_client._get_llm_client_cached.cache_clear,
        )
        # Suppress expected "No data found" and "DatabaseNotCreatedError" logs on fresh installs
        logging.getLogger("cognee.shared.logging_utils").setLevel(logging.CRITICAL)
        
        root = str(Path.home() / ".studybuddy" / "cognee")
        cognee.config.data_root_directory(root)
        cognee.config.system_root_directory(f"{root}/system")
        from cognee.infrastructure.databases.relational.create_db_and_tables import create_db_and_tables
        await create_db_and_tables()
        logger.warning("Cognee mandatory bootstrap took %.2fs", time.perf_counter() - t0)

        # Bootstrap the global profile dataset. Per-project datasets are created
        # defensively by StudentMemoryService on project creation and before every
        # write/flush, because Cognee remember(session_id=...) only writes cache.
        from app.services.student_memory import StudentMemoryService
        async def _warm_profile_dataset() -> None:
            warm_t0 = time.perf_counter()
            try:
                await StudentMemoryService().ensure_profile_dataset()
                logger.warning("Cognee profile dataset warmup took %.2fs", time.perf_counter() - warm_t0)
            except Exception:
                logger.exception("Cognee profile dataset warmup failed")
        asyncio.create_task(_warm_profile_dataset())
    except Exception as e:
        print("Cognee setup error:", e)
    yield


app = FastAPI(title="ResearchMate API", lifespan=lifespan)

_origins_env = os.getenv("ALLOWED_ORIGINS", "")
_origins = (
    _origins_env.split(",")
    if _origins_env
    else [
        "http://localhost:5173",
        "http://127.0.0.1:5173",
        "http://localhost:5174",
        "http://127.0.0.1:5174",
    ]
)

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

app.include_router(health.router)
app.include_router(control_room.router)
app.include_router(agent_control.router)
app.include_router(library.router)
app.include_router(projects.router)
app.include_router(sandbox.router)
app.include_router(session.router)
app.include_router(annotations.router)
app.include_router(regions.router)
app.include_router(review.router)
app.include_router(citation_graph.router)
app.include_router(chat.router, prefix="/chat", tags=["Chat"])
app.include_router(draft.router)
app.include_router(wiki.router)
app.include_router(voice.router)


@app.websocket("/ws/{project_id}")
async def websocket_endpoint(ws: WebSocket, project_id: str):
    import asyncio
    import logging
    logger = logging.getLogger(__name__)
    cm = get_connection_manager()
    background_events: set[asyncio.Task[None]] = set()

    async def run_event(event_type: str, data: dict) -> None:
        try:
            await handle_event(project_id, event_type, data)
            if event_type in {"CHAT_TURN", "CONTEXT_CARD_REQUEST"}:
                logger.info("websocket event done project=%s type=%s", project_id, event_type)
        except (WebSocketDisconnect, ClientDisconnected, RuntimeError) as exc:
            logger.info("websocket event aborted project=%s type=%s detail=%s", project_id, event_type, exc)
        except Exception:
            logger.exception("handle_event failed for session %s event %s", project_id, event_type)
            await cm.send(project_id, "ERROR", {
                "event_type": event_type,
                "message": "Something went wrong processing that -> please try again.",
            })

    await cm.connect(project_id, ws)
    logger.info("websocket connected project=%s", project_id)
    try:
        while True:
            msg = await ws.receive_json()
            event_type = msg.get("type", "")
            if event_type in {"CHAT_TURN", "CONTEXT_CARD_REQUEST"}:
                data = msg.get("data", {}) or {}
                logger.info(
                    "websocket event start project=%s type=%s selection_len=%s has_image=%s",
                    project_id,
                    event_type,
                    len(data.get("selection_text") or ""),
                    bool(data.get("selection_image_base64")),
                )
            data = msg.get("data", {}) or {}
            if event_type in {"CHAT_TURN", "CONTEXT_CARD_REQUEST"}:
                task = asyncio.create_task(run_event(event_type, data))
                background_events.add(task)
                task.add_done_callback(background_events.discard)
            else:
                await run_event(event_type, data)
    except WebSocketDisconnect as exc:
        logger.info("websocket disconnected project=%s code=%s", project_id, getattr(exc, "code", None))
        cm.disconnect(project_id, ws)
    except RuntimeError as exc:
        logger.info("websocket closed project=%s detail=%s", project_id, exc)
        cm.disconnect(project_id, ws)
    finally:
        for task in background_events:
            task.cancel()
        if background_events:
            await asyncio.gather(*background_events, return_exceptions=True)