backend/api/upload.py

import shutil
import tempfile
import uuid
from pathlib import Path
from fastapi import APIRouter, UploadFile, File, HTTPException, BackgroundTasks
from pydantic import BaseModel
from backend.ingestion.pdf_loader     import ingest_pdf
from backend.ingestion.url_loader     import ingest_url
from backend.ingestion.youtube_loader import ingest_youtube
from backend.database.connection import get_connection

router = APIRouter()


# ── Request models ────────────────────────────────────────────────────────────

class UrlRequest(BaseModel):
    url: str
    language: str = "en"

class YoutubeRequest(BaseModel):
    url: str
    language: str = "en"


# ── Background Task Wrappers ──────────────────────────────────────────────────

def run_ingest_pdf_bg(tmp_path: Path, filename: str, source_id: str):
    """Safely runs PDF ingestion and always cleans up the temporary file."""
    try:
        ingest_pdf(tmp_path, filename, source_id)
    except Exception as e:
        print(f"[BG Ingest PDF] Critical failure: {e}")
    finally:
        if tmp_path.exists():
            tmp_path.unlink()
            print(f"[BG Ingest PDF] Cleaned up temporary file: {tmp_path}")


# ── Endpoints ─────────────────────────────────────────────────────────────────

@router.post("/upload-pdf")
async def upload_pdf(background_tasks: BackgroundTasks, file: UploadFile = File(...)):
    """
    POST /upload-pdf
    Accepts a PDF file upload, saves it temporarily,
    starts ingestion in a background task, and immediately returns source ID.
    """
    if not file.filename.endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Only PDF files are accepted.")

    # Save uploaded file to a temp location first
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
        shutil.copyfileobj(file.file, tmp)
        tmp_path = Path(tmp.name)

    source_id = str(uuid.uuid4())
    dest_path = tmp_path # temporary placeholder path until moved in background task

    try:
        # Create initial pending source entry in MySQL
        with get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                INSERT INTO sources (id, type, title, origin, language, status, progress_percentage, chunk_count)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            """, (source_id, "pdf", file.filename, str(dest_path), "en", "UPLOADING", 10, 0))
            conn.commit()

        # Dispatch background task
        background_tasks.add_task(run_ingest_pdf_bg, tmp_path, file.filename, source_id)
        
    except Exception as e:
        if tmp_path.exists():
            tmp_path.unlink()
        raise HTTPException(status_code=500, detail=f"Database initialization failed: {str(e)}")

    return {
        "message"    : "PDF upload initiated",
        "source_id"  : source_id,
        "title"      : file.filename,
        "chunk_count": 0,
        "status"     : "UPLOADING"
    }


@router.post("/add-url")
def add_url(req: UrlRequest, background_tasks: BackgroundTasks):
    """
    POST /add-url
    Accepts a website URL, scrapes and ingests it in the background.
    """
    if not req.url.startswith("http"):
        raise HTTPException(status_code=400, detail="Invalid URL. Must start with http:// or https://")

    source_id = str(uuid.uuid4())

    try:
        # Create initial pending source entry in MySQL
        with get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                INSERT INTO sources (id, type, title, origin, language, status, progress_percentage, chunk_count)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            """, (source_id, "url", req.url, req.url, "en", "EXTRACTING", 10, 0))
            conn.commit()

        # Dispatch background task
        background_tasks.add_task(ingest_url, req.url, source_id)

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Database initialization failed: {str(e)}")

    return {
        "message"    : "URL ingestion initiated",
        "source_id"  : source_id,
        "title"      : req.url,
        "chunk_count": 0,
        "status"     : "EXTRACTING"
    }


@router.post("/add-youtube")
def add_youtube(req: YoutubeRequest, background_tasks: BackgroundTasks):
    """
    POST /add-youtube
    Accepts a YouTube URL, fetches transcript and ingests it in the background.
    """
    if "youtube.com" not in req.url and "youtu.be" not in req.url:
        raise HTTPException(status_code=400, detail="Invalid YouTube URL.")

    source_id = str(uuid.uuid4())

    try:
        # Create initial pending source entry in MySQL
        with get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                INSERT INTO sources (id, type, title, origin, language, status, progress_percentage, chunk_count)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
            """, (source_id, "youtube", req.url, req.url, "en", "EXTRACTING", 10, 0))
            conn.commit()

        # Dispatch background task
        background_tasks.add_task(ingest_youtube, req.url, source_id)

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Database initialization failed: {str(e)}")

    return {
        "message"    : "YouTube video ingestion initiated",
        "source_id"  : source_id,
        "title"      : req.url,
        "chunk_count": 0,
        "language"   : req.language,
        "status"     : "EXTRACTING"
    }