chatkit/backend/app/main.py

"""FastAPI entrypoint for the ClassLens backend (v2 - 3-step workflow)."""

from __future__ import annotations

from pathlib import Path
from contextlib import asynccontextmanager
from typing import Optional

from fastapi import FastAPI, Depends, File, Form, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, Response, FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from .database import (
    init_database,
    create_session,
    get_sessions,
    get_session,
    update_session_status,
    get_parsed_data,
    save_prompt,
    get_prompts,
    get_all_prompts,
    get_prompt,
    update_prompt,
    delete_prompt,
    save_report,
    get_reports,
    get_report,
)
from .auth import get_current_user, register_user, login_user
from .file_processor import process_uploaded_files
from .report_generator import generate_report, get_default_prompt

STATIC_DIR = Path(__file__).parent.parent / "static"


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize database on startup."""
    await init_database()
    yield


app = FastAPI(
    title="ClassLens API",
    description="AI-powered exam analysis for teachers (v2)",
    version="2.0.0",
    lifespan=lifespan,
)

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


# =============================================================================
# Auth Endpoints
# =============================================================================

class AuthRequest(BaseModel):
    email: str
    password: str


@app.post("/api/auth/register")
async def api_register(body: AuthRequest):
    user = await register_user(body.email, body.password)
    from .auth import create_access_token
    token = create_access_token({"sub": str(user["id"])})
    return {"token": token, "user": {"id": user["id"], "email": user["email"], "display_name": user["display_name"]}}


@app.post("/api/auth/login")
async def api_login(body: AuthRequest):
    result = await login_user(body.email, body.password)
    return result


@app.get("/api/auth/me")
async def api_me(user=Depends(get_current_user)):
    return {"id": user["id"], "email": user["email"], "display_name": user["display_name"]}


# =============================================================================
# Session Endpoints
# =============================================================================

class CreateSessionRequest(BaseModel):
    title: str = "Untitled Session"


@app.post("/api/sessions")
async def api_create_session(body: CreateSessionRequest, user=Depends(get_current_user)):
    session_id = await create_session(user["id"], body.title)
    session = await get_session(session_id)
    return session


@app.get("/api/sessions")
async def api_list_sessions(user=Depends(get_current_user)):
    sessions = await get_sessions(user["id"])
    return {"sessions": sessions}


@app.get("/api/sessions/{session_id}")
async def api_get_session(session_id: int, user=Depends(get_current_user)):
    session = await get_session(session_id)
    if not session or session["user_id"] != user["id"]:
        raise HTTPException(status_code=404, detail="Session not found")
    return session


# =============================================================================
# File Upload & Processing Endpoints
# =============================================================================

@app.post("/api/sessions/{session_id}/upload")
async def api_upload_files(
    session_id: int,
    data_type: str = Form(...),
    files: list[UploadFile] = File(...),
    description: str = Form(""),
    model: str = Form("gpt-5.4"),
    user=Depends(get_current_user),
):
    """Upload files for a session. data_type: 'questions', 'student_answers', or 'teacher_answers'."""
    session = await get_session(session_id)
    if not session or session["user_id"] != user["id"]:
        raise HTTPException(status_code=404, detail="Session not found")

    if data_type not in ("questions", "student_answers", "teacher_answers"):
        raise HTTPException(status_code=400, detail="data_type must be 'questions', 'student_answers', or 'teacher_answers'")

    try:
        structured = await process_uploaded_files(files, data_type, session_id, description=description, model=model)
        await update_session_status(session_id, "processed")
        return {"status": "ok", "data_type": data_type, "data": structured}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Processing failed: {str(e)}")


@app.get("/api/sessions/{session_id}/parsed-data")
async def api_get_parsed_data(session_id: int, user=Depends(get_current_user)):
    session = await get_session(session_id)
    if not session or session["user_id"] != user["id"]:
        raise HTTPException(status_code=404, detail="Session not found")
    data = await get_parsed_data(session_id)
    # Group by data_type
    grouped = {}
    for item in data:
        dt = item["data_type"]
        if dt not in grouped:
            grouped[dt] = []
        grouped[dt].append(item)
    return {"parsed_data": grouped}


# =============================================================================
# Prompt Endpoints
# =============================================================================

class PromptRequest(BaseModel):
    name: str
    content: str


@app.get("/api/prompts")
async def api_list_prompts(user=Depends(get_current_user)):
    prompts = await get_prompts(user["id"])
    return {"prompts": prompts}


@app.get("/api/prompts/all")
async def api_list_all_prompts(user=Depends(get_current_user)):
    """List all prompts from all users (read others' prompts)."""
    prompts = await get_all_prompts()
    return {"prompts": prompts}


@app.get("/api/prompts/default")
async def api_get_default_prompt():
    """Get the default analysis prompt template."""
    return {"content": get_default_prompt()}


@app.post("/api/prompts")
async def api_save_prompt(body: PromptRequest, user=Depends(get_current_user)):
    prompt_id = await save_prompt(user["id"], body.name, body.content)
    prompt = await get_prompt(prompt_id)
    return prompt


@app.get("/api/prompts/{prompt_id}")
async def api_get_prompt(prompt_id: int, user=Depends(get_current_user)):
    prompt = await get_prompt(prompt_id)
    if not prompt:
        raise HTTPException(status_code=404, detail="Prompt not found")
    return prompt


@app.put("/api/prompts/{prompt_id}")
async def api_update_prompt(prompt_id: int, body: PromptRequest, user=Depends(get_current_user)):
    prompt = await get_prompt(prompt_id)
    if not prompt or prompt["user_id"] != user["id"]:
        raise HTTPException(status_code=404, detail="Prompt not found")
    await update_prompt(prompt_id, body.name, body.content)
    return await get_prompt(prompt_id)


@app.delete("/api/prompts/{prompt_id}")
async def api_delete_prompt(prompt_id: int, user=Depends(get_current_user)):
    prompt = await get_prompt(prompt_id)
    if not prompt or prompt["user_id"] != user["id"]:
        raise HTTPException(status_code=404, detail="Prompt not found")
    await delete_prompt(prompt_id)
    return {"status": "deleted"}


# =============================================================================
# Report Generation Endpoints
# =============================================================================

class GenerateReportRequest(BaseModel):
    prompt_id: Optional[int] = None
    prompt_content: Optional[str] = None
    model: str = "gpt-5.4"


@app.post("/api/sessions/{session_id}/generate-report")
async def api_generate_report(
    session_id: int, body: GenerateReportRequest, user=Depends(get_current_user)
):
    session = await get_session(session_id)
    if not session or session["user_id"] != user["id"]:
        raise HTTPException(status_code=404, detail="Session not found")

    # Load parsed data
    data = await get_parsed_data(session_id)
    questions = {}
    student_answers = {}
    teacher_answers = {}
    for item in data:
        if item["data_type"] == "questions":
            questions = item["structured_data"]
        elif item["data_type"] == "student_answers":
            student_answers = item["structured_data"]
        elif item["data_type"] == "teacher_answers":
            teacher_answers = item["structured_data"]

    if not questions:
        raise HTTPException(status_code=400, detail="No questions data found. Upload and process files first.")

    # Get prompt content
    prompt_content = body.prompt_content
    if not prompt_content and body.prompt_id:
        prompt = await get_prompt(body.prompt_id)
        if prompt:
            prompt_content = prompt["content"]
    if not prompt_content:
        prompt_content = get_default_prompt()

    try:
        html = await generate_report(questions, student_answers, teacher_answers, prompt_content, body.model)
        report_id = await save_report(session_id, body.prompt_id, html)
        return {"report_id": report_id, "html_content": html}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Report generation failed: {str(e)}")


@app.get("/api/sessions/{session_id}/reports")
async def api_list_reports(session_id: int, user=Depends(get_current_user)):
    session = await get_session(session_id)
    if not session or session["user_id"] != user["id"]:
        raise HTTPException(status_code=404, detail="Session not found")
    reports = await get_reports(session_id)
    return {"reports": reports}


@app.get("/api/reports/{report_id}")
async def api_get_report(report_id: int, user=Depends(get_current_user)):
    report = await get_report(report_id)
    if not report:
        raise HTTPException(status_code=404, detail="Report not found")
    return report


# =============================================================================
# Health Check
# =============================================================================

@app.get("/api/health")
@app.get("/health")
async def health_check():
    return {"status": "healthy", "service": "ClassLens", "version": "2.0.0"}


# =============================================================================
# Static File Serving (Production)
# =============================================================================

if STATIC_DIR.exists() and (STATIC_DIR / "index.html").exists():
    if (STATIC_DIR / "assets").exists():
        app.mount("/assets", StaticFiles(directory=STATIC_DIR / "assets"), name="assets")

    @app.get("/favicon.ico")
    async def favicon():
        favicon_path = STATIC_DIR / "favicon.ico"
        if favicon_path.exists():
            return FileResponse(favicon_path)
        return Response(status_code=404)

    @app.get("/")
    async def serve_spa():
        return FileResponse(STATIC_DIR / "index.html")

    @app.get("/{full_path:path}")
    async def serve_spa_routes(full_path: str):
        if full_path.startswith("api/") or full_path == "health":
            return Response(status_code=404)
        file_path = STATIC_DIR / full_path
        if file_path.exists() and file_path.is_file():
            return FileResponse(file_path)
        return FileResponse(STATIC_DIR / "index.html")
else:
    @app.get("/")
    async def root():
        return {
            "name": "ClassLens API",
            "version": "2.0.0",
            "mode": "development",
            "description": "AI-powered exam analysis for teachers",
        }