main.py

"""Backend: serves signed URL for ElevenLabs agent + PDF upload to Databricks."""
import logging
import re
import sys
from pathlib import Path
import os
import io
import uuid
import json
from datetime import datetime
from typing import Any

import numpy as np
from dotenv import load_dotenv

from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.prompts import ChatPromptTemplate
from langchain_classic.chains.retrieval import create_retrieval_chain
from langchain_classic.chains.combine_documents import create_stuff_documents_chain
from langchain_google_genai import ChatGoogleGenerativeAI

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    stream=sys.stdout,
)
log = logging.getLogger(__name__)
import httpx
from fastapi import FastAPI, Query, HTTPException, UploadFile, File, Body, Request
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
from pypdf import PdfReader
from google import genai
from google.genai import types
from databricks.sdk import WorkspaceClient
from databricks import sql

load_dotenv(Path(__file__).resolve().parent / ".env")

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

VOLUME_PATH = "/Volumes/nyusw/buildathon/rag_pdfs"
CHUNKS_TABLE = "nyusw.buildathon.rag_chunks"
CONVERSATION_RAW_TABLE = "nyusw.buildathon.conversations_raw"
TABLE_NAME = "nyusw.buildathon.rag_chunks"
CHUNK_SIZE = 8000
CHUNK_OVERLAP = 200


class ConversationRequest(BaseModel):
    case_id: str
    thread_id: str
    Defence: str = ""
    Plaintiff: str = ""


class QueryRequest(BaseModel):
    case_id: str
    query: str


class JudgementSchema(BaseModel):
    """Structured output for judgement generation. No defaults - Gemini API does not allow them in response_schema."""
    verdict: str  # "Plaintiff" or "Defendant" - who won
    plaintiff_won: bool
    rating: int  # 1-10 strength of winning case
    detailed_statement: str


@app.get("/health")
def health():
    return {"status": "ok"}


@app.get("/signed-url")
def get_signed_url(
    agent_id: str = Query(default=None),
    include_conversation_id: bool = Query(default=False),
    branch_id: str = Query(default=None),
):
    agent = agent_id or os.getenv("ELEVENLABS_AGENT_ID")
    api_key = os.getenv("ELEVENLABS_API_KEY")
    if not agent:
        raise HTTPException(status_code=400, detail="agent_id required (or set ELEVENLABS_AGENT_ID)")
    if not api_key:
        raise HTTPException(status_code=500, detail="ELEVENLABS_API_KEY not configured")

    params = {"agent_id": agent, "include_conversation_id": include_conversation_id}
    if branch_id:
        params["branch_id"] = branch_id

    resp = httpx.get(
        "https://api.elevenlabs.io/v1/convai/conversation/get-signed-url",
        params=params,
        headers={"xi-api-key": api_key},
    )
    if resp.status_code != 200:
        raise HTTPException(status_code=resp.status_code, detail=resp.text)
    data = resp.json()
    return {"signed_url": data.get("signed_url")}


@app.get("/api/cases/{case_id}/files/{filename}")
def get_case_file(case_id: str, filename: str):
    """Return file bytes from the Databricks volume for preview/download. Path params are sanitized; filename is URL-decoded by FastAPI."""
    try:
        cid = _sanitize_case_id(case_id)
        # filename: basename only, no path traversal (sanitized for vol path and Content-Disposition)
        fname = _sanitize_filename(filename)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

    host = os.getenv("DATABRICKS_HOST")
    token = os.getenv("DATABRICKS_TOKEN")
    if not host or not token:
        raise HTTPException(status_code=500, detail="Databricks not configured")

    vol_path = f"{VOLUME_PATH}/{cid}/{fname}"

    try:
        w = WorkspaceClient(host=host, token=token)
        content = w.files.download(vol_path)
        if hasattr(content, "contents") and hasattr(content.contents, "read"):
            body = content.contents.read()
        elif hasattr(content, "read"):
            body = content.read()
        else:
            body = content if isinstance(content, bytes) else b""
        if not body:
            raise HTTPException(status_code=404, detail="File not found or empty")
    except Exception as e:
        log.exception("Volume download failed for %s", vol_path)
        if "not found" in str(e).lower() or "404" in str(e):
            raise HTTPException(status_code=404, detail="File not found")
        raise HTTPException(status_code=500, detail="Failed to download file")

    ext = fname.rsplit(".", 1)[-1].lower() if "." in fname else ""
    media_type = "application/pdf" if ext == "pdf" else "application/octet-stream"
    return Response(
        content=body,
        media_type=media_type,
        headers={
            "Content-Disposition": f'inline; filename="{fname}"',
        },
    )


def extract_text_from_pdf(pdf_bytes: bytes) -> str:
    log.info("Extracting text from PDF (%d bytes)", len(pdf_bytes))
    reader = PdfReader(io.BytesIO(pdf_bytes))
    text = "\n".join(p.extract_text() or "" for p in reader.pages)
    log.info("Extracted %d chars from %d pages", len(text), len(reader.pages))
    return text


def chunk_text(text: str) -> list[str]:
    chunks = []
    start = 0
    while start < len(text):
        end = start + CHUNK_SIZE
        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        start = end - CHUNK_OVERLAP
    log.info("Chunked into %d chunks", len(chunks))
    return chunks


def get_embeddings(texts: list[str], task_type: str = "RETRIEVAL_DOCUMENT") -> list[list[float]]:
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("GEMINI_API_KEY not configured")
    log.info("Getting Gemini embeddings for %d texts (task=%s)", len(texts), task_type)
    client = genai.Client(api_key=api_key)
    result = client.models.embed_content(
        model="gemini-embedding-001",
        contents=texts,
        config=types.EmbedContentConfig(
            output_dimensionality=768,
            task_type=task_type,
        ),
    )
    log.info("Got %d embeddings", len(result.embeddings))
    return [list(e.values) for e in result.embeddings]


def fetch_chunks_for_case(case_id: str) -> list[dict]:
    """Fetch all chunks for a case_id from rag_chunks."""
    conn = get_db_conn()
    cursor = conn.cursor()
    cursor.execute(
        f"""
        SELECT chunk_id, doc_id, source_type, source_uri, chunk_index, chunk_text
        FROM {CHUNKS_TABLE}
        WHERE case_id = ?
        """,
        (case_id,),
    )
    rows = cursor.fetchall()
    conn.close()
    columns = ["chunk_id", "doc_id", "source_type", "source_uri", "chunk_index", "chunk_text"]
    return [dict(zip(columns, row)) for row in rows]


def fetch_chunks_with_vectors(case_id: str) -> list[dict]:
    """Fetch chunks for a case_id including vector_embeddings for similarity search."""
    conn = get_db_conn()
    cursor = conn.cursor()
    cursor.execute(
        f"""
        SELECT chunk_id, doc_id, source_type, source_uri, chunk_index, chunk_text, vector_embeddings
        FROM {CHUNKS_TABLE}
        WHERE case_id = ?
        """,
        (case_id,),
    )
    rows = cursor.fetchall()
    conn.close()
    columns = ["chunk_id", "doc_id", "source_type", "source_uri", "chunk_index", "chunk_text", "vector_embeddings"]
    result = []
    for row in rows:
        d = dict(zip(columns, row))
        vec = d.get("vector_embeddings")
        if vec is not None:
            if isinstance(vec, str):
                try:
                    vec = json.loads(vec)
                except (json.JSONDecodeError, TypeError):
                    vec = []
            elif isinstance(vec, np.ndarray):
                vec = vec.tolist()
            elif not isinstance(vec, list):
                vec = list(vec) if hasattr(vec, "__iter__") and not isinstance(vec, (str, bytes)) else []
        # Avoid "truth value of array is ambiguous" - don't use `if vec` with numpy arrays
        d["vector_embeddings"] = vec if (vec is not None and len(vec) > 0) else []
        result.append(d)
    return result


def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Compute cosine similarity between two vectors."""
    a_norm = np.linalg.norm(a)
    b_norm = np.linalg.norm(b)
    if a_norm == 0 or b_norm == 0:
        return 0.0
    return float(np.dot(a, b) / (a_norm * b_norm))


RAG_TOP_K = 10
JUDGEMENT_RAG_TOP_K = 25

# Retrieval query for judgement - targets verdict, evidence, arguments
JUDGEMENT_RETRIEVAL_QUERY = """case verdict outcome plaintiff defendant evidence arguments legal reasoning key facts pleadings judgement decision"""


def retrieve_chunks_for_judgement(case_id: str) -> list[dict]:
    """Use RAG (vector similarity) to retrieve most relevant chunks for judgement."""
    chunks = fetch_chunks_with_vectors(case_id)
    if not chunks:
        return []
    valid = [c for c in chunks if c.get("vector_embeddings") is not None and len(c.get("vector_embeddings", [])) > 0]
    if not valid:
        return chunks[:JUDGEMENT_RAG_TOP_K]
    query_emb = np.array(get_embeddings([JUDGEMENT_RETRIEVAL_QUERY], task_type="RETRIEVAL_QUERY")[0], dtype=np.float32)
    scored = []
    for c in valid:
        raw = c["vector_embeddings"]
        if isinstance(raw, str):
            try:
                raw = json.loads(raw)
            except (json.JSONDecodeError, TypeError):
                continue
        vec = np.array(raw, dtype=np.float32)
        score = cosine_similarity(query_emb, vec)
        scored.append((score, c))
    scored.sort(key=lambda x: -x[0])
    return [c for _, c in scored[:JUDGEMENT_RAG_TOP_K]]


class DatabricksCaseRetriever(BaseRetriever):
    """LangChain retriever over Databricks rag_chunks with vector similarity."""

    case_id: str

    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun | None = None,
    ) -> list[Document]:
        chunks = fetch_chunks_with_vectors(self.case_id)
        if not chunks:
            return []
        valid = [c for c in chunks if c.get("vector_embeddings") is not None and len(c.get("vector_embeddings", [])) > 0]
        if not valid:
            return [Document(page_content=c["chunk_text"], metadata={"chunk_id": c["chunk_id"], "source_type": c["source_type"], "source_uri": c["source_uri"]}) for c in chunks[:RAG_TOP_K]]
        query_emb = np.array(get_embeddings([query], task_type="RETRIEVAL_QUERY")[0], dtype=np.float32)
        scored = []
        for c in valid:
            raw = c["vector_embeddings"]
            if isinstance(raw, str):
                try:
                    raw = json.loads(raw)
                except (json.JSONDecodeError, TypeError):
                    continue
            vec = np.array(raw, dtype=np.float32)
            score = cosine_similarity(query_emb, vec)
            scored.append((score, c))
        scored.sort(key=lambda x: -x[0])
        top = scored[:RAG_TOP_K]
        return [
            Document(
                page_content=c["chunk_text"],
                metadata={
                    "chunk_id": c["chunk_id"],
                    "source_type": c["source_type"],
                    "source_uri": c["source_uri"],
                    "score": float(s),
                },
            )
            for s, c in top
        ]


def build_rag_chain(case_id: str) -> Any:
    """Build LangChain RAG chain: retriever (DB) -> LLM."""
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("GEMINI_API_KEY not configured")
    retriever = DatabricksCaseRetriever(case_id=case_id)
    llm = ChatGoogleGenerativeAI(
        model="gemini-2.5-flash",
        google_api_key=api_key,
        temperature=0,
    )
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "Use ONLY the following case documents to answer the question. If the answer is not in the documents, say so.\n\nCASE DOCUMENTS:\n{context}"),
            ("human", "QUESTION: {input}\n\nANSWER:"),
        ]
    )
    combine_docs_chain = create_stuff_documents_chain(llm, prompt)
    return create_retrieval_chain(retriever, combine_docs_chain)


def ask_llm_with_context(query: str, context_chunks: list[str]) -> str:
    """Pass case chunks to LLM as context and get answer."""
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("GEMINI_API_KEY not configured")
    context = "\n\n---\n\n".join(context_chunks)
    prompt = f"""Use ONLY the following case documents to answer the question. If the answer is not in the documents, say so.

CASE DOCUMENTS:
{context}

QUESTION: {query}

ANSWER:"""
    client = genai.Client(api_key=api_key)
    response = client.models.generate_content(
        model="gemini-1.5-flash",
        contents=prompt,
    )
    return (response.text or "").strip()


@app.post("/api/query")
async def rag_query(body: QueryRequest = Body(...)):
    """
    LangChain RAG: retriever (Databricks DB) -> LLM. Uses vector similarity for top-k retrieval.
    """
    log.info("RAG query: case_id=%s, query=%s", body.case_id, body.query[:50] + "..." if len(body.query) > 50 else body.query)

    if not body.query or not body.query.strip():
        raise HTTPException(status_code=400, detail="query required")

    try:
        chain = build_rag_chain(body.case_id)
    except ValueError as e:
        raise HTTPException(status_code=500, detail=str(e))

    try:
        result = chain.invoke({"input": body.query.strip()})
    except Exception as e:
        log.exception("RAG chain failed")
        raise HTTPException(status_code=500, detail=f"RAG failed: {e}")

    answer = (result.get("answer") or "").strip()
    context_docs = result.get("context", [])
    if isinstance(context_docs, str):
        context_docs = []
    chunks_used = [
        {
            "chunk_id": d.metadata.get("chunk_id", ""),
            "case_id": body.case_id,
            "source_type": d.metadata.get("source_type", ""),
            "source_uri": d.metadata.get("source_uri", ""),
            "chunk_text": d.page_content,
        }
        for d in context_docs
    ]
    log.info("RAG query done: case_id=%s, chunks=%d", body.case_id, len(chunks_used))
    return {"case_id": body.case_id, "query": body.query, "answer": answer, "chunks_used": chunks_used}


JUDGEMENT_PROMPT = """You are a legal expert and judge. Analyze the following case documents and produce a final judgement.

CASE DOCUMENTS (all RAG chunks for this case):
{context}

Based solely on the above case documents (PDFs, conversations, pleadings, etc.), provide:
1. **verdict**: Who won the case? Return exactly "Plaintiff" or "Defendant".
2. **plaintiff_won**: true if Plaintiff won, false if Defendant won.
3. **rating**: An integer 1-10 indicating the strength/confidence of the winning party's position (1=weak, 10=very strong).
4. **detailed_statement**: A comprehensive judgement statement (2-4 paragraphs) summarizing the case, key evidence, legal reasoning, and the final decision with justification. Write in formal legal language suitable for an official judgement document.

Respond with valid JSON only, no markdown or extra text."""


def generate_judgement(case_id: str) -> dict:
    """Use RAG (vector similarity) to retrieve relevant chunks, then generate judgement via Gemini 2.5 Flash."""
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("GEMINI_API_KEY not configured")
    chunks = retrieve_chunks_for_judgement(case_id)
    if not chunks:
        raise ValueError("No chunks found for this case")
    context = "\n\n---\n\n".join(
        f"[{c.get('source_type', 'unknown')}] {c['chunk_text']}" for c in chunks
    )
    log.info("Generating judgement for case_id=%s with %d RAG-retrieved chunks", case_id, len(chunks))
    client = genai.Client(api_key=api_key)
    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=JUDGEMENT_PROMPT.format(context=context),
        config=types.GenerateContentConfig(
            response_mime_type="application/json",
            response_schema=JudgementSchema,
        ),
    )
    text = (response.text or "").strip()
    if not text:
        raise ValueError("Empty response from model")
    data = json.loads(text)
    data["case_id"] = case_id
    data["won"] = data.get("plaintiff_won", False)
    return data


def judgement_to_pdf(data: dict) -> bytes:
    """Convert judgement data to PDF bytes."""
    from reportlab.lib.pagesizes import letter
    from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
    from reportlab.lib.units import inch
    from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle

    buffer = io.BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=letter, rightMargin=inch, leftMargin=inch, topMargin=inch, bottomMargin=inch)
    styles = getSampleStyleSheet()
    title_style = ParagraphStyle(name="Title", parent=styles["Heading1"], fontSize=16, spaceAfter=12)
    heading_style = ParagraphStyle(name="Heading", parent=styles["Heading2"], fontSize=12, spaceAfter=6)
    body_style = styles["Normal"]

    story = []
    story.append(Paragraph("JUDGEMENT", title_style))
    story.append(Paragraph(f"Case ID: {data.get('case_id', 'N/A')}", body_style))
    story.append(Spacer(1, 12))

    verdict = data.get("verdict", "N/A")
    plaintiff_won = data.get("plaintiff_won", False)
    story.append(Paragraph(f"<b>Verdict:</b> {verdict} won", heading_style))
    story.append(Paragraph(f"<b>Plaintiff won:</b> {'Yes' if plaintiff_won else 'No'}", body_style))
    if "won" in data:
        story.append(Paragraph(f"<b>You won:</b> {'Yes' if data.get('won') else 'No'}", body_style))
    story.append(Paragraph(f"<b>Strength rating:</b> {data.get('rating', 5)}/10", body_style))
    story.append(Spacer(1, 12))

    story.append(Paragraph("<b>Detailed Statement</b>", heading_style))
    for para in (data.get("detailed_statement", "") or "").split("\n\n"):
        if para.strip():
            story.append(Paragraph(para.strip().replace("\n", " "), body_style))
            story.append(Spacer(1, 6))

    doc.build(story)
    return buffer.getvalue()


def _judgement_response(case_id: str, format: str, perspective: str | None):
    """Shared logic for GET/POST judgement: generate judgement and return JSON or PDF."""
    try:
        cid = _sanitize_case_id(case_id)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

    try:
        data = generate_judgement(cid)
    except ValueError as e:
        if "No chunks" in str(e):
            raise HTTPException(status_code=404, detail=str(e))
        raise HTTPException(status_code=500, detail=str(e))
    except Exception as e:
        log.exception("Judgement generation failed")
        raise HTTPException(status_code=500, detail=f"Judgement failed: {e}")

    if perspective and perspective.lower() in ("plaintiff", "defendant"):
        data["won"] = data.get("plaintiff_won", False) if perspective.lower() == "plaintiff" else not data.get("plaintiff_won", False)
    else:
        data["won"] = data.get("plaintiff_won", False)

    if format.lower() == "pdf":
        pdf_bytes = judgement_to_pdf(data)
        return Response(
            content=pdf_bytes,
            media_type="application/pdf",
            headers={"Content-Disposition": f'attachment; filename="judgement_{cid}.pdf"'},
        )
    return data


@app.get("/api/cases/{case_id}/judgement")
@app.post("/api/cases/{case_id}/judgement")
async def get_judgement(
    request: Request,
    case_id: str,
    format: str | None = Query(default=None, description="Response format: json or pdf (GET defaults to pdf, POST to json)"),
    perspective: str | None = Query(default=None, description="Your role: plaintiff or defendant - sets 'won' from your perspective"),
):
    """
    Generate a judgement for the case by analyzing all RAG chunks.
    Returns verdict (who won), plaintiff_won, rating (1-10), and detailed statement.
    GET and POST supported. GET defaults to PDF for download; use ?format=json for JSON. POST defaults to JSON; use ?format=pdf for PDF.
    Use ?perspective=plaintiff or ?perspective=defendant to get 'won' from your perspective.
    """
    if format is None:
        format = "pdf" if request.method == "GET" else "json"
    return _judgement_response(case_id, format, perspective)


def upload_to_volume(case_id: str, filename: str, pdf_bytes: bytes) -> str:
    host = os.getenv("DATABRICKS_HOST")
    token = os.getenv("DATABRICKS_TOKEN")
    if not host or not token:
        raise ValueError("DATABRICKS_HOST and DATABRICKS_TOKEN required")
    vol_path = f"{VOLUME_PATH}/{case_id}/{filename}"
    log.info("Uploading PDF to volume: %s", vol_path)
    w = WorkspaceClient(host=host, token=token)
    w.files.upload(vol_path, io.BytesIO(pdf_bytes), overwrite=True)
    log.info("Volume upload done")
    return vol_path


def _sanitize_filename(name: str) -> str:
    """Basename only, alphanumeric, dash, underscore, dot, space. No path traversal."""
    if not name or not name.strip():
        raise ValueError("filename required")
    base = name.strip().split("/")[-1].split("\\")[-1]
    if not re.match(r"^[a-zA-Z0-9._\- ]+$", base):
        raise ValueError("filename contains invalid characters")
    return base


def _sanitize_case_id(case_id: str) -> str:
    """Non-empty, safe for path (e.g. 6-digit numeric)."""
    if not case_id or not str(case_id).strip():
        raise ValueError("case_id required")
    cid = str(case_id).strip()
    if not re.match(r"^[a-zA-Z0-9-]+$", cid):
        raise ValueError("case_id contains invalid characters")
    return cid


def insert_chunks(chunks_data: list[dict]) -> None:
    log.info("Connecting to Databricks SQL")
    conn = get_db_conn()
    cursor = conn.cursor()
    log.info("Inserting %d chunks into %s", len(chunks_data), CHUNKS_TABLE)
    for i, row in enumerate(chunks_data):
        arr_sql = "array(" + ", ".join(str(f) for f in row["vector_embeddings"]) + ")"
        cursor.execute(
            f"""
            INSERT INTO {CHUNKS_TABLE}
            (chunk_id, doc_id, source_type, source_uri, chunk_index, chunk_text, updated_at, metadata_json, vector_embeddings, case_id)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, {arr_sql}, ?)
            """,
            (
                row["chunk_id"],
                row["doc_id"],
                row["source_type"],
                row["source_uri"],
                row["chunk_index"],
                row["chunk_text"],
                row["updated_at"],
                row["metadata_json"],
                row.get("case_id"),
            ),
        )
        if (i + 1) % 10 == 0:
            log.info("Inserted %d/%d chunks", i + 1, len(chunks_data))
    conn.close()
    log.info("All chunks inserted")


def get_db_conn():
    host = os.getenv("DATABRICKS_HOST")
    token = os.getenv("DATABRICKS_TOKEN")
    http_path = os.getenv("DATABRICKS_HTTP_PATH")
    if not all([host, token, http_path]):
        raise ValueError("DATABRICKS_HOST, DATABRICKS_TOKEN, DATABRICKS_HTTP_PATH required")
    server = host.replace("https://", "").replace("http://", "").rstrip("/")
    return sql.connect(
        server_hostname=server,
        http_path=http_path,
        access_token=token,
    )


def insert_conversation_raw(rows: list[dict]) -> None:
    log.info("Inserting %d rows into %s", len(rows), CONVERSATION_RAW_TABLE)
    conn = get_db_conn()
    cursor = conn.cursor()
    for row in rows:
        cursor.execute(
            f"""
            INSERT INTO {CONVERSATION_RAW_TABLE}
            (event_id, thread_id, tenant_id, role, text, created_at, metadata_json, case_id)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                row.get("event_id"),
                row["thread_id"],
                row.get("tenant_id"),
                row["role"],
                row["text"],
                row["created_at"],
                row.get("metadata_json"),
                row["case_id"],
            ),
        )
    conn.close()
    log.info("Conversation raw insert done")


@app.post("/conversation-upload")
async def store_conversation(body: ConversationRequest = Body(...)):
    """
    Store conversation text in conversation_raw, embed Defence and Plaintiff texts,
    and store embeddings in rag_chunks (one chunk per role).
    """
    log.info("Conversation request: case_id=%s, thread_id=%s", body.case_id, body.thread_id)

    now = datetime.utcnow().isoformat() + "Z"
    roles_data = [
        ("Defendant", body.Defence),
        ("Plaintiff", body.Plaintiff),
    ]

    # Filter empty texts
    non_empty = [(r, t) for r, t in roles_data if t and t.strip()]
    if not non_empty:
        raise HTTPException(status_code=400, detail="At least one of Defence or Plaintiff text required")

    # Insert into conversation_raw
    raw_rows = []
    for role, text in roles_data:
        if text and text.strip():
            raw_rows.append({
                "event_id": str(uuid.uuid4()),
                "thread_id": body.thread_id,
                "tenant_id": None,
                "role": role,
                "text": text.strip(),
                "created_at": now,
                "metadata_json": json.dumps({"case_id": body.case_id}),
                "case_id": body.case_id,
            })
    try:
        insert_conversation_raw(raw_rows)
    except Exception as e:
        log.exception("Conversation raw insert failed")
        raise HTTPException(status_code=500, detail=f"Conversation raw insert failed: {e}")

    # Embed and insert into rag_chunks (one chunk per role)
    texts_to_embed = [t for r, t in roles_data if t and t.strip()]
    try:
        embeddings = get_embeddings(texts_to_embed)
    except Exception as e:
        log.exception("Embedding failed")
        raise HTTPException(status_code=500, detail=f"Embedding failed: {e}")

    doc_id = str(uuid.uuid4())
    source_uri = f"conversation:{body.thread_id}"
    chunks_data = []
    for i, ((role, text), emb) in enumerate(zip(non_empty, embeddings)):
        chunks_data.append({
            "chunk_id": str(uuid.uuid4()),
            "doc_id": doc_id,
            "source_type": "conversation",
            "source_uri": source_uri,
            "chunk_index": i,
            "chunk_text": text.strip(),
            "updated_at": now,
            "metadata_json": json.dumps({"case_id": body.case_id, "thread_id": body.thread_id, "role": role}),
            "vector_embeddings": emb,
            "case_id": body.case_id,
        })
    try:
        insert_chunks(chunks_data)
    except Exception as e:
        log.exception("Rag chunks insert failed")
        raise HTTPException(status_code=500, detail=f"Rag chunks insert failed: {e}")

    log.info("Conversation stored: case_id=%s, thread_id=%s, chunks=%d", body.case_id, body.thread_id, len(chunks_data))
    return {
        "case_id": body.case_id,
        "thread_id": body.thread_id,
        "doc_id": doc_id,
        "chunks_count": len(chunks_data),
    }


@app.post("/api/cases/{case_id}/pdf")
async def upload_case_pdf(
    case_id: str,
    file: UploadFile = File(...),
):
    """Upload PDF to Databricks volume, extract text, chunk, embed with Gemini (768d), store in rag_chunks."""
    log.info("PDF upload request: case_id=%s, filename=%s", case_id, file.filename or "(none)")
    if not file.filename or not file.filename.lower().endswith(".pdf"):
        raise HTTPException(status_code=400, detail="PDF file required")

    log.info("Reading file...")
    pdf_bytes = await file.read()
    log.info("Read %d bytes", len(pdf_bytes))
    if not pdf_bytes:
        raise HTTPException(status_code=400, detail="Empty file")

    try:
        text = extract_text_from_pdf(pdf_bytes)
    except Exception as e:
        log.exception("PDF extraction failed")
        raise HTTPException(status_code=400, detail=f"Failed to extract PDF text: {e}")

    if not text.strip():
        raise HTTPException(status_code=400, detail="No text extracted from PDF")

    chunks = chunk_text(text)
    if not chunks:
        raise HTTPException(status_code=400, detail="No chunks extracted")

    try:
        embeddings = get_embeddings(chunks)
    except Exception as e:
        log.exception("Embedding failed")
        raise HTTPException(status_code=500, detail=f"Embedding failed: {e}")

    doc_id = str(uuid.uuid4())
    source_uri = f"{VOLUME_PATH}/{case_id}/{file.filename}"
    now = datetime.utcnow().isoformat() + "Z"

    try:
        upload_to_volume(case_id, file.filename, pdf_bytes)
    except Exception as e:
        log.exception("Volume upload failed")
        raise HTTPException(status_code=500, detail=f"Volume upload failed: {e}")

    chunks_data = []
    for i, (chunk_text_val, emb) in enumerate(zip(chunks, embeddings)):
        chunks_data.append({
            "chunk_id": str(uuid.uuid4()),
            "doc_id": doc_id,
            "source_type": "pdf",
            "source_uri": source_uri,
            "chunk_index": i,
            "chunk_text": chunk_text_val,
            "updated_at": now,
            "case_id": case_id,
            "metadata_json": json.dumps({"case_id": case_id, "filename": file.filename}),
            "vector_embeddings": emb,
            "case_id": case_id,
        })

    try:
        insert_chunks(chunks_data)
    except Exception as e:
        log.exception("Insert failed")
        raise HTTPException(status_code=500, detail=f"Insert failed: {e}")

    log.info("PDF upload complete: case_id=%s, doc_id=%s, chunks=%d", case_id, doc_id, len(chunks))
    return {
        "case_id": case_id,
        "doc_id": doc_id,
        "source_uri": source_uri,
        "chunks_count": len(chunks),
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "8000")))