main.py

"""
NoteAbstract — Video Course to Structured Notes.
FastAPI application entry point.
"""

import asyncio
import json
import os
import shutil
import subprocess
import threading
import uuid
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

# Fix SSL certificate verification on Windows
try:
    import certifi
    os.environ["SSL_CERT_FILE"] = certifi.where()
    os.environ["REQUESTS_CA_BUNDLE"] = certifi.where()
except ImportError:
    pass
# Also disable HF Hub SSL verification as fallback
os.environ.setdefault("HF_HUB_DISABLE_SSL_VERIFY", "1")

from fastapi import FastAPI, File, Form, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

from config import PROVIDERS, settings
from models.job import Job, get_session, init_db

# ---------------------------------------------------------------------------
# Lifespan — startup / shutdown
# ---------------------------------------------------------------------------


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    Path(settings.storage_dir).mkdir(parents=True, exist_ok=True)
    init_db()
    app.state.ffmpeg_available = _check_ffmpeg()
    app.state.loop = asyncio.get_running_loop()
    app.state.cancelled_jobs: set[str] = set()
    app.state.processing_lock = asyncio.Lock()
    app.state.ws_manager = ConnectionManager()
    if app.state.ffmpeg_available:
        print("[OK] FFmpeg found")
    else:
        print("[WARN] FFmpeg not found — audio extraction will fail. Install from https://ffmpeg.org/download.html")
    print(f"NoteAbstract running at http://{settings.host}:{settings.port}")
    yield
    # Shutdown
    print("Shutting down...")


def _check_ffmpeg() -> bool:
    try:
        subprocess.run(["ffmpeg", "-version"], capture_output=True, timeout=5)
        return True
    except Exception:
        return False


# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------

app = FastAPI(title="NoteAbstract", lifespan=lifespan)

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


# ---------------------------------------------------------------------------
# WebSocket Connection Manager
# ---------------------------------------------------------------------------


class ConnectionManager:
    """Manages WebSocket connections keyed by job_id."""

    def __init__(self):
        self._connections: dict[str, list[WebSocket]] = {}

    async def connect(self, ws: WebSocket, job_id: str):
        await ws.accept()
        self._connections.setdefault(job_id, []).append(ws)

    async def disconnect(self, ws: WebSocket, job_id: str):
        if job_id in self._connections:
            self._connections[job_id] = [c for c in self._connections[job_id] if c is not ws]
            if not self._connections[job_id]:
                del self._connections[job_id]

    async def broadcast(self, job_id: str, data: dict):
        stale = []
        for ws in self._connections.get(job_id, []):
            try:
                await ws.send_json(data)
            except Exception:
                stale.append(ws)
        for ws in stale:
            await self.disconnect(ws, job_id)

    async def broadcast_log(self, job_id: str, level: str, message: str):
        await self.broadcast(job_id, {"type": "log", "level": level, "message": message})


def ws_manager() -> ConnectionManager:
    """Dependency to get the ConnectionManager from app state."""
    from fastapi import Request
    # We access via app.state internally; routes use the module-level reference.
    return app.state.ws_manager


# ---------------------------------------------------------------------------
# Static
# ---------------------------------------------------------------------------


@app.get("/")
async def index():
    return FileResponse(Path(__file__).parent / "static" / "index.html")


# ---------------------------------------------------------------------------
# REST — Jobs
# ---------------------------------------------------------------------------


# Valid LLM providers
VALID_PROVIDERS = ("claude", "openai", "deepseek", "qwen", "zhipu", "moonshot", "custom")


@app.post("/api/jobs")
async def create_job(
    file: Optional[UploadFile] = File(None),
    url: Optional[str] = Form(None),
    language: str = Form("auto"),
    asr_model: str = Form("base"),
    llm_provider: str = Form("claude"),
    llm_model: str = Form(""),
):
    """Create a new processing job. Provide either a file or a URL."""
    if not file and not url:
        raise HTTPException(400, "Either 'file' or 'url' must be provided.")
    if file and url:
        raise HTTPException(400, "Provide only one of 'file' or 'url', not both.")
    if asr_model not in ("tiny", "base", "small", "medium"):
        raise HTTPException(400, "asr_model must be one of: tiny, base, small, medium.")
    if language not in ("auto", "en", "zh"):
        raise HTTPException(400, "language must be one of: auto, en, zh.")
    if llm_provider not in VALID_PROVIDERS:
        raise HTTPException(400, f"llm_provider must be one of: {', '.join(VALID_PROVIDERS)}.")

    job_id = str(uuid.uuid4())
    storage_dir = Path(settings.storage_dir) / "jobs" / job_id
    storage_dir.mkdir(parents=True, exist_ok=True)

    source_type = "file" if file else "url"
    source_url = url if url else None
    source_filename = None

    # Save uploaded file
    if file:
        source_filename = file.filename or "uploaded_video"
        # Sanitize filename
        safe_name = Path(file.filename or "uploaded_video").name
        file_path = storage_dir / safe_name
        # Stream to disk
        with open(file_path, "wb") as f:
            while chunk := await file.read(8 * 1024 * 1024):
                f.write(chunk)
        source_filename = safe_name

    # Use provider's default model if not specified
    llm = llm_model or PROVIDERS.get(llm_provider, {}).get("default_model", "")

    with get_session() as sess:
        job = Job(
            id=job_id,
            status="pending",
            source_type=source_type,
            source_url=source_url,
            source_filename=source_filename,
            language=language,
            asr_model=asr_model,
            llm_provider=llm_provider,
            llm_model=llm,
            storage_dir=str(storage_dir),
        )
        sess.add(job)
        sess.commit()

    # Enqueue processing (fire-and-forget background task)
    asyncio.create_task(run_job_pipeline(job_id))

    return JSONResponse({"job_id": job_id, "status": "pending"})


@app.get("/api/jobs")
async def list_jobs(limit: int = 20, offset: int = 0):
    """List recent jobs."""
    with get_session() as sess:
        query = sess.query(Job).order_by(Job.created_at.desc()).offset(offset).limit(limit)
        jobs = [j.to_dict() for j in query.all()]
        total = sess.query(Job).count()
    return {"jobs": jobs, "total": total}


@app.get("/api/jobs/{job_id}")
async def get_job(job_id: str):
    """Get a single job's status."""
    with get_session() as sess:
        job = sess.get(Job, job_id)
        if not job:
            raise HTTPException(404, "Job not found.")
        return job.to_dict()


@app.delete("/api/jobs/{job_id}")
async def delete_job(job_id: str):
    """Delete a job and its storage directory."""
    with get_session() as sess:
        job = sess.get(Job, job_id)
        if not job:
            raise HTTPException(404, "Job not found.")
        storage = Path(job.storage_dir)
        sess.delete(job)
        sess.commit()
    if storage.exists():
        shutil.rmtree(storage, ignore_errors=True)
    return {"deleted": True}


@app.get("/api/jobs/{job_id}/notes")
async def get_job_notes(job_id: str):
    """Get the generated notes as JSON (structured + raw markdown)."""
    with get_session() as sess:
        job = sess.get(Job, job_id)
        if not job:
            raise HTTPException(404, "Job not found.")
        if job.status != "completed":
            raise HTTPException(400, "Job is not completed yet.")

    notes_path = Path(job.storage_dir) / "notes.md"
    if not notes_path.exists():
        raise HTTPException(404, "Notes file not found.")

    raw_md = notes_path.read_text(encoding="utf-8")
    structured = parse_notes_markdown(raw_md)

    return {
        "job_id": job_id,
        "title": job.title or structured.get("title", "Untitled"),
        **structured,
        "full_markdown": raw_md,
    }


@app.get("/api/jobs/{job_id}/download")
async def download_notes(job_id: str):
    """Download the notes as a .md file."""
    with get_session() as sess:
        job = sess.get(Job, job_id)
        if not job:
            raise HTTPException(404, "Job not found.")

    notes_path = Path(job.storage_dir) / "notes.md"
    if not notes_path.exists():
        raise HTTPException(404, "Notes file not found.")

    filename = (job.title or "notes")[:80].replace("/", "_").replace("\\", "_")
    return FileResponse(
        notes_path,
        media_type="text/markdown",
        filename=f"{filename}.md",
    )


# ---------------------------------------------------------------------------
# REST — Settings
# ---------------------------------------------------------------------------


@app.get("/api/settings")
async def get_settings():
    """Return current settings (API keys masked)."""
    return {
        "ffmpeg_available": app.state.ffmpeg_available,
        "asr_models": ["tiny", "base", "small", "medium"],
        "default_asr_model": settings.default_asr_model,
        "default_language": settings.default_language,
        "default_llm_provider": settings.default_llm_provider,
        "default_llm_model": settings.default_llm_model,
        "providers": list(PROVIDERS.keys()),
        "provider_names": {k: v["name"] for k, v in PROVIDERS.items()},
        "provider_types": {k: v["type"] for k, v in PROVIDERS.items()},
        "api_keys": {
            "claude": {"set": settings.has_claude_key(), "preview": mask_key_str(settings.claude_api_key)},
            "openai": {"set": bool(settings.openai_api_key), "preview": mask_key_str(settings.openai_api_key)},
            "deepseek": {"set": bool(settings.deepseek_api_key), "preview": mask_key_str(settings.deepseek_api_key)},
            "qwen": {"set": bool(settings.qwen_api_key), "preview": mask_key_str(settings.qwen_api_key)},
            "zhipu": {"set": bool(settings.zhipu_api_key), "preview": mask_key_str(settings.zhipu_api_key)},
            "moonshot": {"set": bool(settings.moonshot_api_key), "preview": mask_key_str(settings.moonshot_api_key)},
            "custom": {"set": bool(settings.custom_api_key), "preview": mask_key_str(settings.custom_api_key)},
        },
        "custom_base_url": mask_key_str(settings.custom_base_url) if settings.custom_base_url else "",
    }


@app.put("/api/settings")
async def update_settings(data: dict):
    """Update settings and persist to .env."""
    api_key_fields = [
        "claude_api_key", "openai_api_key", "deepseek_api_key",
        "qwen_api_key", "zhipu_api_key", "moonshot_api_key",
        "custom_api_key", "custom_base_url",
    ]
    for field in api_key_fields:
        if field in data:
            setattr(settings, field, data[field])
    for key in ("default_asr_model", "default_language", "default_llm_provider", "default_llm_model"):
        if key in data:
            setattr(settings, key, data[key])
    settings.save_to_env()
    return await get_settings()


# ---------------------------------------------------------------------------
# WebSocket
# ---------------------------------------------------------------------------


@app.websocket("/ws/{job_id}")
async def websocket_endpoint(ws: WebSocket, job_id: str):
    mgr: ConnectionManager = app.state.ws_manager
    await mgr.connect(ws, job_id)
    try:
        while True:
            raw = await ws.receive_text()
            msg = json.loads(raw)
            if msg.get("type") == "cancel":
                app.state.cancelled_jobs.add(job_id)
                await mgr.broadcast(job_id, {"type": "log", "level": "warn", "message": "Cancelling..."})
            elif msg.get("type") == "ping":
                await ws.send_json({"type": "pong"})
    except WebSocketDisconnect:
        pass
    except Exception:
        pass
    finally:
        await mgr.disconnect(ws, job_id)


# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------

# Import processors lazily to avoid circular issues
# We import at top of function to defer until everything is loaded


def _safe_schedule(coro):
    """Schedule a coroutine safely — works from both main loop and worker threads."""
    loop = app.state.loop
    if loop is None:
        return
    if threading.current_thread() is threading.main_thread():
        asyncio.create_task(coro)
    else:
        loop.call_soon_threadsafe(lambda c=coro: asyncio.ensure_future(c))


async def run_job_pipeline(job_id: str):
    """
    Orchestrates the full processing pipeline for a job.
    Runs as a background asyncio task.
    """
    # Avoid concurrent processing — queue jobs
    lock: asyncio.Lock = app.state.processing_lock
    cancelled: set = app.state.cancelled_jobs
    mgr: ConnectionManager = app.state.ws_manager

    async with lock:
        # Load job
        with get_session() as sess:
            job = sess.get(Job, job_id)
            if not job:
                return

        storage_dir = Path(job.storage_dir)

        try:
            # ---- Stage 1: Download (URL jobs only) ----
            if job.source_type == "url" and not _has_video_file(storage_dir):
                if job_id in cancelled:
                    raise JobCancelled()
                _update_job(job_id, status="downloading", progress=0)
                await mgr.broadcast(job_id, {
                    "type": "stage_change", "stage": "downloading",
                    "label": "Downloading video...", "progress": 0
                })

                from utils.downloader import download_video

                def dl_progress(pct: float, msg: str = ""):
                    overall = pct * 0.20
                    _update_job(job_id, progress=overall)
                    _safe_schedule(mgr.broadcast(job_id, {
                        "type": "progress", "stage": "downloading",
                        "stage_progress": round(pct, 1),
                        "overall_progress": round(overall, 1),
                    }))
                    if msg:
                        _safe_schedule(mgr.broadcast_log(job_id, "info", msg))

                video_path = await asyncio.to_thread(
                    download_video, job.source_url, str(storage_dir), dl_progress
                )
                await mgr.broadcast(job_id, {
                    "type": "stage_complete", "stage": "downloading",
                    "result": {"filename": Path(video_path).name}
                })

            # ---- Stage 2: Extract audio ----
            if job_id in cancelled:
                raise JobCancelled()
            _update_job(job_id, status="extracting_audio", progress=20)
            await mgr.broadcast(job_id, {
                "type": "stage_change", "stage": "extracting_audio",
                "label": "Extracting audio...", "progress": 20
            })

            if not app.state.ffmpeg_available:
                raise FFmpegNotFound()

            video_file = _find_video_file(storage_dir)
            if not video_file:
                raise FileNotFoundError(f"No video file found in {storage_dir}")

            audio_path = storage_dir / "audio.wav"

            from processors.video import extract_audio

            def ex_progress(pct: float):
                overall = 20 + pct * 0.10
                _update_job(job_id, progress=overall)
                _safe_schedule(mgr.broadcast(job_id, {
                    "type": "progress", "stage": "extracting_audio",
                    "stage_progress": round(pct, 1),
                    "overall_progress": round(overall, 1),
                }))

            audio_info = await asyncio.to_thread(
                extract_audio, str(video_file), str(audio_path), ex_progress
            )

            await mgr.broadcast(job_id, {
                "type": "stage_complete", "stage": "extracting_audio",
                "result": audio_info
            })
            _update_job(job_id, progress=30)

            # ---- Stage 3: Transcribe ----
            if job_id in cancelled:
                raise JobCancelled()
            _update_job(job_id, status="transcribing", progress=30)
            await mgr.broadcast(job_id, {
                "type": "stage_change", "stage": "transcribing",
                "label": f"Transcribing with Whisper ({job.asr_model})...",
                "progress": 30
            })

            from processors.asr import transcribe_audio

            def asr_progress(pct: float):
                overall = 30 + pct * 0.40
                _update_job(job_id, progress=overall)
                _safe_schedule(mgr.broadcast(job_id, {
                    "type": "progress", "stage": "transcribing",
                    "stage_progress": round(pct, 1),
                    "overall_progress": round(overall, 1),
                }))

            transcript = await asyncio.to_thread(
                transcribe_audio, str(audio_path), job.asr_model,
                job.language, asr_progress, str(storage_dir)
            )

            await mgr.broadcast(job_id, {
                "type": "stage_complete", "stage": "transcribing",
                "result": {
                    "full_text_length": len(transcript.get("full_text", "")),
                    "segments_count": len(transcript.get("segments", [])),
                    "language_detected": transcript.get("language_detected", "unknown"),
                }
            })
            _update_job(job_id, progress=70)

            # ---- Stage 4: Generate Notes ----
            if job_id in cancelled:
                raise JobCancelled()

            # Truncate if needed
            full_text = transcript.get("full_text", "")
            if len(full_text) > settings.max_transcript_chars:
                await mgr.broadcast_log(job_id, "warn",
                    f"Transcript is {len(full_text):,} chars — truncating to {settings.max_transcript_chars:,} for LLM processing.")
                full_text = full_text[:settings.max_transcript_chars]

            _update_job(job_id, status="generating_notes", progress=70)
            await mgr.broadcast(job_id, {
                "type": "stage_change", "stage": "generating_notes",
                "label": "Generating notes with AI...", "progress": 70
            })

            from processors.llm import generate_notes
            from config import get_api_key_for_provider, get_base_url_for_provider

            api_key = get_api_key_for_provider(job.llm_provider)
            base_url = get_base_url_for_provider(job.llm_provider)
            if not api_key:
                raise NoAPIKey(job.llm_provider)

            def llm_progress(pct: float):
                overall = 70 + pct * 0.25
                _update_job(job_id, progress=overall)
                _safe_schedule(mgr.broadcast(job_id, {
                    "type": "progress", "stage": "generating_notes",
                    "stage_progress": round(pct, 1),
                    "overall_progress": round(overall, 1),
                }))

            notes = await asyncio.to_thread(
                generate_notes, full_text, job.llm_provider,
                job.llm_model or "", api_key, job.language,
                llm_progress, base_url
            )

            # Write notes.md
            full_md = notes.get("full_markdown", "")
            notes_path = storage_dir / "notes.md"
            notes_path.write_text(full_md, encoding="utf-8")

            title = notes.get("title", "Untitled Notes")[:200]

            _update_job(job_id, status="completed", progress=100, title=title)

            await mgr.broadcast(job_id, {
                "type": "complete",
                "job_id": job_id,
                "title": title,
                "notes_preview": notes.get("summary", "")[:500],
            })

        except JobCancelled:
            _update_job(job_id, status="failed", error_message="Cancelled by user.")
            await mgr.broadcast(job_id, {
                "type": "error", "stage": "cancelled",
                "message": "Job cancelled by user."
            })
        except FFmpegNotFound:
            _update_job(job_id, status="failed",
                error_message="FFmpeg is not installed. Please install FFmpeg and restart.")
            await mgr.broadcast(job_id, {
                "type": "error", "stage": "extracting_audio",
                "message": "FFmpeg not found. Install from https://ffmpeg.org/download.html"
            })
        except NoAPIKey as e:
            _update_job(job_id, status="failed",
                error_message=f"No {e.provider} API key configured.")
            await mgr.broadcast(job_id, {
                "type": "error", "stage": "generating_notes",
                "message": f"No {e.provider} API key. Add it in Settings."
            })
        except Exception as e:
            msg = str(e) if str(e) else type(e).__name__
            _update_job(job_id, status="failed", error_message=msg)
            await mgr.broadcast(job_id, {
                "type": "error", "stage": _current_stage(job_id),
                "message": msg
            })


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


class JobCancelled(Exception):
    pass


class FFmpegNotFound(Exception):
    pass


class NoAPIKey(Exception):
    def __init__(self, provider: str):
        self.provider = provider


def _current_stage(job_id: str) -> str:
    with get_session() as sess:
        job = sess.get(Job, job_id)
        return job.status if job else "unknown"


def _update_job(job_id: str, **kwargs):
    """Update job fields in DB (sync, thread-safe — each call creates a new session)."""
    with get_session() as sess:
        job = sess.get(Job, job_id)
        if job:
            for key, value in kwargs.items():
                if hasattr(job, key):
                    setattr(job, key, value)
            if kwargs.get("status") in ("completed", "failed"):
                job.completed_at = datetime.now(timezone.utc).isoformat()
            sess.commit()


def _find_video_file(storage_dir: Path) -> Path | None:
    """Find a video file in the storage directory."""
    for ext in (".mp4", ".mkv", ".webm", ".avi", ".mov", ".flv", ".wmv", ".m4v"):
        candidates = list(storage_dir.glob(f"*{ext}"))
        if candidates:
            return candidates[0]
    # Also look for any file that isn't audio.wav or notes.md or transcript.json
    for f in storage_dir.iterdir():
        if f.is_file() and f.name not in ("audio.wav", "notes.md", "transcript.json"):
            return f
    return None


def _has_video_file(storage_dir: Path) -> bool:
    return _find_video_file(storage_dir) is not None


def mask_key_str(key: str | None) -> str:
    if not key:
        return ""
    if len(key) <= 8:
        return key[:2] + "****"
    return key[:4] + "..." + key[-4:]


def parse_notes_markdown(raw: str) -> dict:
    """Parse structured notes from markdown. Fallback gracefully."""
    result = {
        "title": "Untitled",
        "summary": "",
        "chapters": [],
        "definitions": [],
        "action_items": [],
    }

    lines = raw.split("\n")

    # Title: first ## heading
    for line in lines:
        stripped = line.strip()
        if stripped.startswith("## ") and not stripped.startswith("### "):
            result["title"] = stripped[3:].strip()
            break

    # Summary: text between "Summary" heading and next --- or ###
    in_summary = False
    summary_lines = []
    for line in lines:
        s = line.strip()
        if s.lower().startswith("### summary") or s.lower().startswith("## summary"):
            in_summary = True
            continue
        if in_summary:
            if s.startswith("---") or s.startswith("### ") or s.startswith("## "):
                break
            if s:
                summary_lines.append(s)
    result["summary"] = " ".join(summary_lines).strip()

    # Chapters: split on "### Chapter" or numbered chapter headings
    in_chapter = False
    current_chapter = None
    for line in lines:
        s = line.strip()
        if s.startswith("### Chapter") or (s.startswith("### ") and any(
            kw in s.lower() for kw in ("chapter", "章节", "part", "部分")
        )):
            if current_chapter:
                result["chapters"].append(current_chapter)
            current_chapter = {"title": s.lstrip("#").strip(), "key_points": []}
            in_chapter = True
            continue
        if in_chapter and current_chapter:
            if s.startswith("---"):
                result["chapters"].append(current_chapter)
                current_chapter = None
                in_chapter = False
                continue
            if s.startswith("### ") and "chapter" not in s.lower():
                result["chapters"].append(current_chapter)
                current_chapter = None
                in_chapter = False
                continue
            if s.startswith("- ") or s.startswith("* "):
                current_chapter["key_points"].append(s[2:].strip())
    if current_chapter and current_chapter.get("key_points"):
        result["chapters"].append(current_chapter)

    # Definitions: look for "Key Definitions" or "定义" section
    in_defs = False
    for line in lines:
        s = line.strip()
        if s.lower().startswith("### key definition") or "关键定义" in s or "术语" in s:
            in_defs = True
            continue
        if in_defs:
            if s.startswith("### ") or s.startswith("## "):
                break
            if (s.startswith("- **") or s.startswith("* **")) and ":" in s:
                content = s.lstrip("-* ").strip()
                # Parse **Term**: Definition
                if content.startswith("**") and "**:" in content:
                    parts = content.split("**:", 1)
                    term = parts[0].strip("* ")
                    definition = parts[1].strip() if len(parts) > 1 else ""
                    result["definitions"].append({"term": term, "definition": definition})
                elif content.startswith("**") and "** " in content:
                    parts = content.split("** ", 1)
                    term = parts[0].strip("* ")
                    definition = parts[1].strip() if len(parts) > 1 else ""
                    result["definitions"].append({"term": term, "definition": definition})

    # Action items: look for "Action Items" or "Takeaways" or "行动"
    in_actions = False
    for line in lines:
        s = line.strip()
        if s.lower().startswith("### action") or s.lower().startswith("### takeaway") or "行动" in s or "收获" in s:
            in_actions = True
            continue
        if in_actions:
            if s.startswith("### ") or s.startswith("## "):
                break
            if s.startswith("- ") or s.startswith("* "):
                result["action_items"].append(s[2:].strip())

    return result


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host=settings.host, port=settings.port, reload=False)