routes/Memory_Upload.py

import socket
from ipaddress import ip_address, ip_network

from fastapi import APIRouter, Depends, UploadFile, File, Form, HTTPException, Request
from fastapi.responses import StreamingResponse
from starlette.concurrency import run_in_threadpool
from routes.limiter import limiter
from dotenv import load_dotenv
import json
import os
import uuid
from pathlib import Path
from urllib.parse import urlparse

from .auth import verify_token

_BLOCKED_NETS = [
    ip_network("127.0.0.0/8"),    # Loopback
    ip_network("10.0.0.0/8"),     # RFC 1918 private
    ip_network("172.16.0.0/12"),  # RFC 1918 private
    ip_network("192.168.0.0/16"), # RFC 1918 private
    ip_network("169.254.0.0/16"), # Link-local / cloud metadata (AWS, GCP, Azure)
    ip_network("100.64.0.0/10"),  # Shared address space
    ip_network("0.0.0.0/8"),      # This network
    ip_network("::1/128"),        # IPv6 loopback
    ip_network("fc00::/7"),       # IPv6 private (ULA)
    ip_network("fe80::/10"),      # IPv6 link-local
]


def _assert_ssrf_safe(url: str) -> None:
    parsed = urlparse(url)
    hostname = parsed.hostname
    if not hostname:
        raise HTTPException(status_code=400, detail="Invalid URL: missing hostname")
    try:
        infos = socket.getaddrinfo(hostname, None)
        if not infos:
            raise HTTPException(status_code=400, detail="Could not resolve hostname")
        ip = ip_address(infos[0][4][0])
    except HTTPException:
        raise
    except Exception:
        raise HTTPException(status_code=400, detail="Could not resolve hostname")
    if any(ip in net for net in _BLOCKED_NETS):
        raise HTTPException(status_code=400, detail="URL resolves to a blocked/private address")

from langchain_core.documents import Document
from typing import Any

from .rag_chunking import (
    CHROMA_COLLECTION_NAME,
    CHROMA_PATH,
    build_embedding_function,
    enrich_documents,
    scrape_page,
    split_documents,
)

load_dotenv()

try:
    from supabase import create_client
except Exception:
    create_client = None

SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_SERVICE_ROLE_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY")
if not SUPABASE_URL or not SUPABASE_SERVICE_ROLE_KEY:
    raise RuntimeError("SUPABASE_URL and SUPABASE_SERVICE_ROLE_KEY must be set in environment")

# File upload hard limits and allowed MIME types (can be overridden via env)
MAX_UPLOAD_SIZE_BYTES = int(os.getenv("MAX_UPLOAD_SIZE_BYTES", str(10 * 1024 * 1024)))  # 10 MB default
ALLOWED_MIME_TYPES = ["application/pdf"]

if create_client is None:
    raise RuntimeError("supabase-py is not installed. Please install the 'supabase' package to use storage features")

supabase = create_client(SUPABASE_URL, SUPABASE_SERVICE_ROLE_KEY)

router = APIRouter(tags=["Embeddings and RAG"])

# Number of chunks to embed+insert per batch. Kept small so streamed progress
# updates arrive frequently on the client.
EMBED_BATCH_SIZE = 32


def _ndjson(obj: dict) -> str:
    """Serialize a progress/result event as a single newline-delimited JSON line."""
    return json.dumps(obj) + "\n"


def _extract_doc_id(inserted, db_res):
    """Best-effort extraction of the documents.id for the freshly-inserted row."""
    if inserted:
        if isinstance(inserted, list) and inserted and isinstance(inserted[0], dict):
            return inserted[0].get("id")
        if isinstance(inserted, dict):
            return inserted.get("id")
    try:
        db_data = getattr(db_res, "data", None)
        if isinstance(db_data, list) and db_data and isinstance(db_data[0], dict):
            return db_data[0].get("id")
        if isinstance(db_data, dict):
            return db_data.get("id")
    except Exception:
        pass
    return None


def _parse_pdf_to_docs(tmp_path: str, original_filename: str) -> list[Document]:
    """Extract page text (pdfplumber) and tables (camelot) into LangChain Documents.

    Falls back to PyMuPDFLoader when pdfplumber is unavailable. Runs synchronously;
    call via run_in_threadpool from async contexts.
    """
    pdfplumber: Any = None
    camelot: Any = None
    try:
        import pdfplumber as _pdfplumber
        pdfplumber = _pdfplumber
    except Exception:
        pdfplumber = None
    try:
        import camelot as _camelot
        camelot = _camelot
    except Exception:
        camelot = None

    docs: list[Document] = []
    if pdfplumber:
        with pdfplumber.open(tmp_path) as pdf:
            for i, page in enumerate(pdf.pages):
                text = page.extract_text() or ""
                docs.append(Document(page_content=text, metadata={"source": original_filename, "page": i}))

        if camelot:
            tables = []
            try:
                tables = camelot.read_pdf(tmp_path, pages="all", flavor="lattice")
                if not tables or len(tables) == 0:
                    tables = camelot.read_pdf(tmp_path, pages="all", flavor="stream")
            except Exception:
                tables = []

            for tbl in tables:
                page_idx = None
                page_attr = getattr(tbl, "page", None)
                if isinstance(page_attr, (str, int)):
                    try:
                        page_idx = int(page_attr) - 1
                    except Exception:
                        page_idx = None

                table_text = tbl.df.to_csv(sep="\t", index=False)
                meta: dict[str, Any] = {"source": original_filename, "chunk_type": "table"}
                if page_idx is not None:
                    meta["page"] = page_idx
                docs.append(Document(page_content=table_text, metadata=meta))
    else:
        try:
            from langchain_community.document_loaders import PyMuPDFLoader
            docs = PyMuPDFLoader(str(tmp_path)).load()
        except Exception as e:
            raise RuntimeError(f"PDF parsing requires either pdfplumber or PyMuPDFLoader: {e}")

    return docs


async def _embed_and_store(splits, *, user_uuid, doc_id_val, source_fields: dict):
    """Async generator: embed `splits` in batches, insert rows, yielding NDJSON
    progress events (55→90%). Yields the total inserted-row count as the last item
    via a final event with key `inserted_rows`.
    """
    total = len(splits)
    if total == 0:
        yield {"inserted_rows": 0}
        return

    done_count = 0
    embedder = await run_in_threadpool(build_embedding_function)

    for i in range(0, total, EMBED_BATCH_SIZE):
        batch_docs = splits[i : i + EMBED_BATCH_SIZE]
        batch_texts = [d.page_content for d in batch_docs]
        batch_embs = await run_in_threadpool(embedder.embed_documents, batch_texts)

        rows = []
        for doc_chunk, emb in zip(batch_docs, batch_embs):
            row_item = {
                "user_id": str(user_uuid),
                "content": doc_chunk.page_content,
                "embedding": emb,
                "chunk_type": doc_chunk.metadata.get("chunk_type", "text"),
                **source_fields,
            }
            if doc_id_val:
                row_item["doc_id"] = doc_id_val
            rows.append(row_item)

        def _insert(batch=rows):
            resp = supabase.table("rag_user_documents").insert(batch).execute()
            resp_error = getattr(resp, "error", None)
            resp_status = getattr(resp, "status_code", None)
            if resp_error or (resp_status is not None and int(resp_status) >= 400):
                raise RuntimeError(f"Supabase insert failed: {resp_error or resp_status}")

        await run_in_threadpool(_insert)

        done_count += len(batch_docs)
        progress = 55 + int(35 * done_count / total)  # 55 → 90
        yield {
            "stage": "embedding",
            "message": f"Embedding chunks ({done_count}/{total})…",
            "progress": progress,
        }

    yield {"inserted_rows": total}


@router.post("/chunk_pdf")
@limiter.limit("5/minute")
async def chunk_pdf(request: Request, file: UploadFile = File(...), user_id: str = Form(...), token_user_id: str = Depends(verify_token)):
    # Basic MIME check (client-provided); perform stronger validation after reading bytes
    if file.content_type not in ALLOWED_MIME_TYPES:
        raise HTTPException(status_code=400, detail="Only PDF files are allowed")

    # validate user_id is a proper UUID (Supabase `auth.users.id`)
    try:
        user_uuid = uuid.UUID(user_id)
    except Exception:
        raise HTTPException(status_code=400, detail="user_id must be a valid UUID corresponding to auth.users.id")

    if str(user_uuid) != token_user_id:
        raise HTTPException(status_code=403, detail="Forbidden: user_id does not match authenticated user")

    contents = await file.read()

    # Size check
    if len(contents) > MAX_UPLOAD_SIZE_BYTES:
        raise HTTPException(status_code=413, detail=f"File too large. Max size: {MAX_UPLOAD_SIZE_BYTES} bytes")

    # Quick PDF header/magic check
    if not isinstance(contents, (bytes, bytearray)) or not contents.startswith(b"%PDF"):
        raise HTTPException(status_code=400, detail="Uploaded file does not appear to be a valid PDF")

    # Sanitize filename
    original_filename = Path(file.filename or "upload.pdf").name
    unique_name = f"{uuid.uuid4().hex}_{original_filename}"
    # store under user folder to match RLS storage policies
    storage_path = f"{user_uuid}/{unique_name}"

    async def event_stream():
        tmp_path = None
        try:
            yield _ndjson({"stage": "uploading", "message": "Uploading file to storage…", "progress": 10})

            # upload bytes to Supabase storage (bucket: documents)
            try:
                await run_in_threadpool(lambda: supabase.storage.from_('documents').upload(
                    storage_path,
                    contents,
                    {"content-type": "application/pdf"},
                ))
            except Exception as e:
                err_text = str(e)
                if 'Bucket not found' in err_text or "statusCode': 404" in err_text:
                    err_text = (
                        "Supabase bucket 'documents' was not found. Please create a storage bucket "
                        "named 'documents' in your Supabase project. Original error: " + err_text
                    )
                yield _ndjson({"stage": "error", "error": f"Storage upload failed: {err_text}"})
                return

            yield _ndjson({"stage": "saving", "message": "Saving document record…", "progress": 20})

            row = {"user_id": str(user_uuid), "filename": original_filename, "storage_path": storage_path}
            try:
                db_res = await run_in_threadpool(lambda: supabase.table("documents").insert(row).execute())
            except Exception as e:
                yield _ndjson({"stage": "error", "error": f"DB insert failed: {e}"})
                return

            try:
                sel = await run_in_threadpool(
                    lambda: supabase.table("documents").select("*").eq("storage_path", storage_path).execute()
                )
                sel_data = getattr(sel, "data", None)
                inserted = sel_data[0] if isinstance(sel_data, list) and len(sel_data) > 0 else sel_data
            except Exception:
                inserted = None

            # Embedding phase. Failures here are non-fatal: the document is already stored.
            inserted_rows = 0
            embed_error = None
            try:
                yield _ndjson({"stage": "parsing", "message": "Extracting text & tables…", "progress": 30})

                import tempfile
                suffix = Path(original_filename).suffix or ".pdf"
                with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
                    tmp.write(contents)
                    tmp_path = tmp.name

                docs = await run_in_threadpool(_parse_pdf_to_docs, tmp_path, original_filename)
                docs = await run_in_threadpool(enrich_documents, docs)

                yield _ndjson({"stage": "chunking", "message": "Splitting into chunks…", "progress": 45})
                embedder = await run_in_threadpool(build_embedding_function)
                splits = await run_in_threadpool(lambda: split_documents(docs, embedder=embedder))

                doc_id_val = _extract_doc_id(inserted, db_res)
                async for evt in _embed_and_store(
                    splits, user_uuid=user_uuid, doc_id_val=doc_id_val, source_fields={"source_type": "pdf"}
                ):
                    if "inserted_rows" in evt:
                        inserted_rows = evt["inserted_rows"]
                    else:
                        yield _ndjson(evt)
            except Exception as e:
                embed_error = str(e)
            finally:
                if tmp_path:
                    try:
                        os.remove(tmp_path)
                    except Exception:
                        pass

            result = {
                "status": "ok",
                "storage_path": storage_path,
                "document": inserted,
                "db_result": getattr(db_res, "data", db_res),
                "embeddings": {"inserted_rows": inserted_rows, "error": embed_error},
            }
            yield _ndjson({"stage": "done", "message": "Done", "progress": 100, "result": result})
        except Exception as e:
            yield _ndjson({"stage": "error", "error": str(e)})

    return StreamingResponse(event_stream(), media_type="application/x-ndjson")


@router.post("/chunk_url")
@limiter.limit("5/minute")
async def chunk_url(request: Request, url: str = Form(...), user_id: str = Form(...), token_user_id: str = Depends(verify_token)):
    # Validate URL scheme and guard against SSRF
    parsed = urlparse(url)
    if parsed.scheme not in ("http", "https"):
        raise HTTPException(status_code=400, detail="URL must start with http:// or https://")
    _assert_ssrf_safe(url)

    try:
        user_uuid = uuid.UUID(user_id)
    except Exception:
        raise HTTPException(status_code=400, detail="user_id must be a valid UUID corresponding to auth.users.id")

    if str(user_uuid) != token_user_id:
        raise HTTPException(status_code=403, detail="Forbidden: user_id does not match authenticated user")

    original_filename = Path(parsed.path).name or parsed.netloc
    storage_path = url  # store the original URL in the documents row for reference

    async def event_stream():
        try:
            yield _ndjson({"stage": "fetching", "message": "Fetching page…", "progress": 10})
            try:
                content = await run_in_threadpool(scrape_page, url)
            except Exception as e:
                yield _ndjson({"stage": "error", "error": f"Failed to fetch/parse URL: {e}"})
                return

            yield _ndjson({"stage": "saving", "message": "Saving document record…", "progress": 25})
            row = {"user_id": str(user_uuid), "filename": original_filename, "storage_path": storage_path}
            try:
                db_res = await run_in_threadpool(lambda: supabase.table("documents").insert(row).execute())
            except Exception as e:
                yield _ndjson({"stage": "error", "error": f"DB insert failed: {e}"})
                return

            try:
                sel = await run_in_threadpool(
                    lambda: supabase.table("documents").select("*").eq("storage_path", storage_path).execute()
                )
                sel_data = getattr(sel, "data", None)
                inserted = sel_data[0] if isinstance(sel_data, list) and len(sel_data) > 0 else sel_data
            except Exception:
                inserted = None

            inserted_rows = 0
            embed_error = None
            try:
                yield _ndjson({"stage": "chunking", "message": "Splitting into chunks…", "progress": 40})
                docs: list[Document] = [Document(page_content=content, metadata={"source": original_filename})]
                docs = await run_in_threadpool(enrich_documents, docs)
                embedder = await run_in_threadpool(build_embedding_function)
                splits = await run_in_threadpool(lambda: split_documents(docs, embedder=embedder))

                doc_id_val = _extract_doc_id(inserted, db_res)
                async for evt in _embed_and_store(
                    splits, user_uuid=user_uuid, doc_id_val=doc_id_val,
                    source_fields={"source_type": "url", "url": url},
                ):
                    if "inserted_rows" in evt:
                        inserted_rows = evt["inserted_rows"]
                    else:
                        yield _ndjson(evt)
            except Exception as e:
                embed_error = str(e)

            result = {
                "status": "ok",
                "url": storage_path,
                "document": inserted,
                "db_result": getattr(db_res, "data", db_res),
                "embeddings": {"inserted_rows": inserted_rows, "error": embed_error},
            }
            yield _ndjson({"stage": "done", "message": "Done", "progress": 100, "result": result})
        except Exception as e:
            yield _ndjson({"stage": "error", "error": str(e)})

    return StreamingResponse(event_stream(), media_type="application/x-ndjson")