Nat20-Notes/backend/app/routers/sessions.py

import json
import os
import shutil
from pathlib import Path

from fastapi import APIRouter, UploadFile, File, Form, HTTPException

from .. import database as db, config, jobs
from ..config import merge_campaign_settings_with_env
from ..errors import PipelineError
from ..pipeline.audio import extract_audio, generate_playback_audio
from ..pipeline.transcribe import transcribe_and_diarize


def _dedup_path(path: Path) -> Path:
    """Append a counter suffix like ` (1)`, ` (2)` if the file exists."""
    if not path.exists():
        return path
    stem = path.stem
    suffix = path.suffix
    parent = path.parent
    counter = 1
    while True:
        new = parent / f"{stem} ({counter}){suffix}"
        if not new.exists():
            return new
        counter += 1

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

# All uploaded/derived files live under config.UPLOAD_DIR / config.AUDIO_DIR /
# config.TRANSCRIPT_DIR (all under config.DATA_DIR, the /data volume) - never
# a hardcoded /app/storage path, which isn't backed by any volume the rest of
# the app agrees on.


def _row_to_dict(row) -> dict:
    return dict(row)


@router.get("")
def list_sessions(campaign_id: str = None):
    if campaign_id:
        rows = db.get_conn().execute(
            "SELECT * FROM sessions WHERE campaign_id = ? ORDER BY created_at DESC", (campaign_id,)
        ).fetchall()
    else:
        rows = db.get_conn().execute("SELECT * FROM sessions ORDER BY created_at DESC").fetchall()
    sessions = []
    for r in rows:
        s = _row_to_dict(r)
        job = db.get_conn().execute(
            "SELECT * FROM jobs WHERE session_id = ? ORDER BY created_at DESC LIMIT 1", (s["id"],)
        ).fetchone()
        s["latest_job"] = _row_to_dict(job) if job else None
        sessions.append(s)
    return sessions


@router.get("/{session_id}")
def get_session(session_id: str):
    row = db.get_conn().execute("SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone()
    if not row:
        raise HTTPException(404, "Session not found")
    session = _row_to_dict(row)
    job = db.get_conn().execute(
        "SELECT * FROM jobs WHERE session_id = ? ORDER BY created_at DESC LIMIT 1", (session_id,)
    ).fetchone()
    session["latest_job"] = _row_to_dict(job) if job else None

    # Stats computed from transcript
    if session.get("transcript_path"):
        tp = Path(session["transcript_path"])
        if tp.exists():
            try:
                data = json.loads(tp.read_text())
                segments = data.get("segments", [])
                if segments:
                    session["audio_duration"] = segments[-1]["end"]
                    session["word_count"] = sum(len(s.get("text", "").split()) for s in segments)
                    session["language"] = data.get("language")
            except (json.JSONDecodeError, OSError):
                pass

    return session


def _run_transcription(session_id: str, video_path: Path, audio_path: Path, transcript_path: Path, campaign_id: str | None = None):
    cid = campaign_id or "default"
    settings = merge_campaign_settings_with_env(db.get_campaign_settings(cid))

    def run(progress_cb):
        db.get_conn()  # ensure thread-local connection exists in this worker thread
        with db.tx() as conn:
            conn.execute("UPDATE sessions SET status = 'transcribing' WHERE id = ?", (session_id,))
        progress_cb("Extracting audio...")
        try:
            extract_audio(video_path, audio_path)
            generate_playback_audio(video_path, audio_path.with_suffix(".opus"))
            transcribe_and_diarize(
                audio_path,
                transcript_path,
                model_size=settings.get("whisper_model", "medium"),
                compute_type=settings.get("whisper_compute_type", "int8"),
                hf_token=settings.get("hf_token", ""),
                progress_cb=progress_cb,
            )
        except PipelineError:
            with db.tx() as conn:
                conn.execute("UPDATE sessions SET status = 'uploaded' WHERE id = ?", (session_id,))
            raise
        with db.tx() as conn:
            conn.execute(
                "UPDATE sessions SET audio_path = ?, transcript_path = ?, status = 'transcribed' WHERE id = ?",
                (str(audio_path), str(transcript_path), session_id),
            )

        # Original upload no longer needed — Opus serves playback, WAV serves
        # re-transcription.  Reclaim storage.
        if video_path.exists():
            try:
                video_path.unlink()
                with db.tx() as conn:
                    conn.execute("UPDATE sessions SET video_path = NULL WHERE id = ?", (session_id,))
            except OSError:
                pass  # non-fatal

    job_id = db.create_job(session_id, "transcribe")
    jobs.submit(job_id, run, requires_gpu=True)
    return job_id


@router.post("")
async def create_session(
    name: str = Form(...),
    file: UploadFile = File(None),
    upload_path: str = Form(None),
    campaign_id: str = Form(None),
):
    session_id = db.new_id()

    if campaign_id:
        upload_dir = config.campaign_dir(campaign_id)
        audio_dir = config.campaign_audio_dir(campaign_id)
        transcript_dir = config.campaign_transcript_dir(campaign_id)
    else:
        upload_dir = config.UPLOAD_DIR
        audio_dir = config.AUDIO_DIR
        transcript_dir = config.TRANSCRIPT_DIR

    if upload_path:
        src = Path(upload_path)
        if not src.exists():
            raise HTTPException(400, f"Uploaded file not found at {upload_path}")
        safe_filename = os.path.basename(src)
        video_path = upload_dir / f"{session_id}_{safe_filename}"
        shutil.move(str(src), str(video_path))
    elif file:
        safe_filename = os.path.basename(file.filename or f"{session_id}")
        video_path = upload_dir / f"{session_id}_{safe_filename}"
        with open(video_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
    else:
        raise HTTPException(400, "Either a file upload or upload_path is required")

    with db.tx() as conn:
        conn.execute(
            "INSERT INTO sessions (id, name, original_filename, video_path, campaign_id, status, created_at) "
            "VALUES (?, ?, ?, ?, ?, 'uploaded', ?)",
            (session_id, name, safe_filename, str(video_path), campaign_id, db.now()),
        )

    audio_path = audio_dir / f"{session_id}.wav"
    transcript_path = transcript_dir / f"{session_id}.json"
    job_id = _run_transcription(session_id, video_path, audio_path, transcript_path, campaign_id)

    return {"session_id": session_id, "job_id": job_id}


@router.post("/{session_id}/transcribe")
def retry_transcription(session_id: str):
    """Re-run transcription for a session, e.g. after a failed job."""
    row = db.get_conn().execute("SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone()
    if not row:
        raise HTTPException(404, "Session not found")
    session = _row_to_dict(row)
    video_path = Path(session["video_path"])
    campaign_id = session.get("campaign_id")
    if campaign_id:
        audio_path = config.campaign_audio_dir(campaign_id) / f"{session_id}.wav"
        transcript_path = config.campaign_transcript_dir(campaign_id) / f"{session_id}.json"
    else:
        audio_path = config.AUDIO_DIR / f"{session_id}.wav"
        transcript_path = config.TRANSCRIPT_DIR / f"{session_id}.json"
    job_id = _run_transcription(session_id, video_path, audio_path, transcript_path, campaign_id)
    return {"job_id": job_id}


@router.post("/upload-chunk")
async def upload_chunk(
    file: UploadFile = File(...),
    chunk_index: int = Form(...),
    total_chunks: int = Form(...),
    filename: str = Form(...),
    upload_id: str = Form(...),
):
    """Assemble a large upload sent in chunks (bypasses nginx's client_max_body_size
    for files bigger than that). Returns a filepath under UPLOAD_DIR that can later
    be referenced when creating a session, for very large recordings."""
    safe_filename = os.path.basename(filename)
    temp_dir = config.UPLOAD_DIR / f"tmp_{upload_id}"
    temp_dir.mkdir(parents=True, exist_ok=True)

    chunk_path = temp_dir / f"chunk_{chunk_index}"
    with open(chunk_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    if len(os.listdir(temp_dir)) == total_chunks:
        final_file_path = _dedup_path(config.UPLOAD_DIR / safe_filename)
        with open(final_file_path, "wb") as final_file:
            for i in range(total_chunks):
                with open(temp_dir / f"chunk_{i}", "rb") as chunk_f:
                    final_file.write(chunk_f.read())
        shutil.rmtree(temp_dir)
        return {"status": "completed", "filepath": str(final_file_path), "filename": safe_filename}

    return {"status": "chunk_received", "chunk_index": chunk_index}