Blame for api_prompts.py

"""
DocuLens - Prompt Configuration API endpoints.

Provides CRUD for prompt versions: list, get active, update/create versions,
activate a specific version, seed defaults, and A/B prompt testing.
"""

import io
import os
import json
import time
import base64
import logging
import tempfile
from typing import Optional

from fastapi import APIRouter, HTTPException, Query, UploadFile, File, Form, Depends
from pydantic import BaseModel
from PIL import Image

from extraction import (
    PROMPT_LIBRARY,
    image_to_base64,
    _parse_vlm_json,
    _vlm_data_to_result,
)
from db.supabase import (
    is_enabled,
    list_prompt_versions,
    get_active_prompt,
    save_prompt_version,
    activate_prompt_version,
    get_all_active_prompts,
)
from middleware import check_rate_limit, validate_file_upload, sanitize_filename

logger = logging.getLogger(__name__)

prompts_router = APIRouter()

HF_TOKEN = os.environ.get("HF_TOKEN", "")


# ── Request / response models ────────────────────────────────────────

class PromptUpdateRequest(BaseModel):
    prompt_text: str
    created_by: Optional[str] = None


class ActivateVersionRequest(BaseModel):
    version: int


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

@prompts_router.get("/api/v1/prompts")
async def api_list_all_prompts(api_key: str = Depends(check_rate_limit)):
    """List all prompts — active DB versions merged with hardcoded defaults.

    For each (use_case, stage) in the hardcoded library, returns:
      - The active DB prompt if one exists, otherwise the hardcoded default.
      - source: "db" or "default" so the UI knows which is in effect.
    """
    # Build the full catalogue from the hardcoded library
    result = []
    for stage, stage_prompts in PROMPT_LIBRARY.items():
        for use_case, prompt_text in stage_prompts.items():
            if use_case == "_default":
                continue  # Skip _default in the UI listing
            entry = {
                "use_case": use_case,
                "stage": stage,
                "prompt_text": prompt_text,
                "source": "default",
                "version": 0,
                "is_active": True,
            }
            result.append(entry)

    # Also include _default prompts for each stage
    for stage, stage_prompts in PROMPT_LIBRARY.items():
        if "_default" in stage_prompts:
            result.append({
                "use_case": "_default",
                "stage": stage,
                "prompt_text": stage_prompts["_default"],
                "source": "default",
                "version": 0,
                "is_active": True,
            })

    # Overlay with DB versions if Supabase is configured
    if is_enabled():
        db_prompts = get_all_active_prompts()
        db_map = {(p["use_case"], p["stage"]): p for p in db_prompts}
        for entry in result:
            key = (entry["use_case"], entry["stage"])
            if key in db_map:
                db_row = db_map[key]
                entry["prompt_text"] = db_row["prompt_text"]
                entry["source"] = "db"
                entry["version"] = db_row["version"]
                entry["db_id"] = db_row["id"]

    return {"prompts": result}


@prompts_router.get("/api/v1/prompts/{use_case}/{stage}")
async def api_get_prompt(use_case: str, stage: str, api_key: str = Depends(check_rate_limit)):
    """Get prompt details for a specific (use_case, stage) pair.

    Returns the active prompt (DB or default) plus version history.
    """
    # Active prompt
    hardcoded = PROMPT_LIBRARY.get(stage, {}).get(
        use_case, PROMPT_LIBRARY.get(stage, {}).get("_default", "")
    )

    active_text = hardcoded
    source = "default"
    active_version = 0

    if is_enabled():
        db_text = get_active_prompt(use_case, stage)
        if db_text:
            active_text = db_text
            source = "db"

    # Version history
    versions = []
    if is_enabled():
        versions = list_prompt_versions(use_case, stage)
        for v in versions:
            if v.get("is_active"):
                active_version = v["version"]

    return {
        "use_case": use_case,
        "stage": stage,
        "active_prompt": active_text,
        "default_prompt": hardcoded,
        "source": source,
        "active_version": active_version,
        "versions": versions,
    }


@prompts_router.put("/api/v1/prompts/{use_case}/{stage}")
async def api_update_prompt(
    use_case: str,
    stage: str,
    body: PromptUpdateRequest,
    api_key: str = Depends(check_rate_limit),
):
    """Save a new version of the prompt for (use_case, stage).

    Creates a new version and makes it active. The previous version
    is preserved in history.
    """
    if not is_enabled():
        raise HTTPException(
            status_code=503,
            detail="Supabase not configured — prompt persistence unavailable",
        )

    row = save_prompt_version(
        use_case=use_case,
        stage=stage,
        prompt_text=body.prompt_text,
        created_by=body.created_by,
    )
    if not row:
        raise HTTPException(status_code=500, detail="Failed to save prompt version")

    return {"status": "ok", "version": row}


@prompts_router.post("/api/v1/prompts/{use_case}/{stage}/activate")
async def api_activate_version(
    use_case: str,
    stage: str,
    body: ActivateVersionRequest,
    api_key: str = Depends(check_rate_limit),
):
    """Activate a specific historical version of a prompt."""
    if not is_enabled():
        raise HTTPException(
            status_code=503,
            detail="Supabase not configured — prompt persistence unavailable",
        )

    ok = activate_prompt_version(use_case, stage, body.version)
    if not ok:
        raise HTTPException(status_code=500, detail="Failed to activate version")

    return {"status": "ok", "activated_version": body.version}


@prompts_router.post("/api/v1/prompts/{use_case}/{stage}/reset")
async def api_reset_prompt(use_case: str, stage: str, api_key: str = Depends(check_rate_limit)):
    """Reset a prompt to the hardcoded default by creating a new version
    with the default text. This preserves history — the user can always
    revert to any previous version.
    """
    if not is_enabled():
        raise HTTPException(
            status_code=503,
            detail="Supabase not configured — prompt persistence unavailable",
        )

    hardcoded = PROMPT_LIBRARY.get(stage, {}).get(
        use_case, PROMPT_LIBRARY.get(stage, {}).get("_default", "")
    )
    if not hardcoded:
        raise HTTPException(status_code=404, detail="No default prompt for this use case/stage")

    row = save_prompt_version(
        use_case=use_case,
        stage=stage,
        prompt_text=hardcoded,
        created_by="system_reset",
    )
    if not row:
        raise HTTPException(status_code=500, detail="Failed to reset prompt")

    return {"status": "ok", "version": row}


@prompts_router.post("/api/v1/prompts/seed")
async def api_seed_defaults(api_key: str = Depends(check_rate_limit)):
    """Seed all hardcoded prompts into the DB as version 1.

    Skips any (use_case, stage) that already has a version in the DB.
    Useful for initial setup.
    """
    if not is_enabled():
        raise HTTPException(
            status_code=503,
            detail="Supabase not configured — prompt persistence unavailable",
        )

    seeded = []
    skipped = []
    for stage, stage_prompts in PROMPT_LIBRARY.items():
        for use_case, prompt_text in stage_prompts.items():
            existing = list_prompt_versions(use_case, stage)
            if existing:
                skipped.append(f"{stage}/{use_case}")
                continue
            row = save_prompt_version(
                use_case=use_case,
                stage=stage,
                prompt_text=prompt_text,
                created_by="seed",
            )
            if row:
                seeded.append(f"{stage}/{use_case}")

    return {"seeded": seeded, "skipped": skipped}


# ---------------------------------------------------------------------------
# A/B Prompt Testing
# ---------------------------------------------------------------------------

def _run_extraction_with_prompt(
    image: Image.Image,
    prompt_text: str,
    model_id: str,
) -> dict:
    """Run extraction on an image using a specific prompt and model.

    Returns the parsed extraction result as a dict.
    """
    from hf_client import create_inference_client, robust_chat_completion

    img_b64 = image_to_base64(image)
    data_url = f"data:image/jpeg;base64,{img_b64}"

    client = create_inference_client(api_key=HF_TOKEN)

    response = robust_chat_completion(
        client,
        model=model_id,
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": data_url}},
                    {"type": "text", "text": prompt_text},
                ],
            }
        ],
        max_tokens=4096,
    )

    result_text = response.choices[0].message.content
    data = _parse_vlm_json(result_text)
    result = _vlm_data_to_result(data, model_id, "auto")

    # Return a clean dict without large base64 images
    d = result.to_dict()
    d.pop("page_image", None)
    d.pop("page_images", None)
    d.pop("raw_text", None)
    d.pop("bounding_boxes", None)
    return d


def _compute_field_diff(result_a: dict, result_b: dict) -> list:
    """Compare two extraction results field by field.

    Returns a list of diff entries:
      {field, value_a, value_b, match: bool, type: 'added'|'removed'|'changed'|'same'}
    """
    fields_a = result_a.get("extracted_fields", {})
    fields_b = result_b.get("extracted_fields", {})

    all_keys = sorted(set(list(fields_a.keys()) + list(fields_b.keys())))
    diff = []

    for key in all_keys:
        val_a = fields_a.get(key)
        val_b = fields_b.get(key)

        if key not in fields_a:
            diff.append({"field": key, "value_a": None, "value_b": val_b, "match": False, "type": "added"})
        elif key not in fields_b:
            diff.append({"field": key, "value_a": val_a, "value_b": None, "match": False, "type": "removed"})
        elif str(val_a).strip() == str(val_b).strip():
            diff.append({"field": key, "value_a": val_a, "value_b": val_b, "match": True, "type": "same"})
        else:
            diff.append({"field": key, "value_a": val_a, "value_b": val_b, "match": False, "type": "changed"})

    # Also compare line_items count
    li_a = result_a.get("line_items", [])
    li_b = result_b.get("line_items", [])
    diff.append({
        "field": "_line_items_count",
        "value_a": len(li_a),
        "value_b": len(li_b),
        "match": len(li_a) == len(li_b),
        "type": "same" if len(li_a) == len(li_b) else "changed",
    })

    return diff


@prompts_router.post("/api/v1/prompts/ab-test")
async def api_ab_test(
    file: UploadFile = File(...),
    prompt_a: str = Form(...),
    prompt_b: str = Form(...),
    model: str = Form("Qwen/Qwen2.5-VL-72B-Instruct"),
    golden_fields: str = Form(""),
    api_key: str = Depends(check_rate_limit),
):
    """Run A/B test: extract from the same document using two different prompts.

    Returns both results plus a field-by-field diff, and optionally compares
    against golden-copy fields if provided.
    """
    validate_file_upload(file.filename, file.size or 0, file.content_type)
    safe_filename = sanitize_filename(file.filename)

    if not HF_TOKEN:
        raise HTTPException(status_code=503, detail="HF_TOKEN not configured")

    contents = await file.read()
    ext = os.path.splitext(safe_filename)[1].lower()

    # Load image
    if ext == ".pdf":
        try:
            import fitz
            doc = fitz.open(stream=contents, filetype="pdf")
            page = doc[0]
            pix = page.get_pixmap(dpi=200)
            image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            doc.close()
        except Exception as e:
            raise HTTPException(status_code=422, detail=f"Failed to render PDF: {e}")
    else:
        try:
            image = Image.open(io.BytesIO(contents))
        except Exception as e:
            raise HTTPException(status_code=422, detail=f"Failed to open image: {e}")

    # Run both extractions
    errors = {}
    result_a = None
    result_b = None

    start_a = time.time()
    try:
        result_a = _run_extraction_with_prompt(image, prompt_a, model)
        result_a["_processing_ms"] = int((time.time() - start_a) * 1000)
    except Exception as e:
        errors["prompt_a"] = str(e)[:200]
        logger.warning(f"A/B test prompt_a failed: {e}")

    start_b = time.time()
    try:
        result_b = _run_extraction_with_prompt(image, prompt_b, model)
        result_b["_processing_ms"] = int((time.time() - start_b) * 1000)
    except Exception as e:
        errors["prompt_b"] = str(e)[:200]
        logger.warning(f"A/B test prompt_b failed: {e}")

    if not result_a and not result_b:
        raise HTTPException(
            status_code=500,
            detail=f"Both prompts failed: {errors}",
        )

    # Compute diff between A and B
    diff_ab = _compute_field_diff(result_a or {}, result_b or {}) if result_a and result_b else []

    # Compare against golden copy if provided
    golden_diff_a = []
    golden_diff_b = []
    golden_score_a = None
    golden_score_b = None

    if golden_fields:
        try:
            golden = json.loads(golden_fields)
            golden_wrapper = {"extracted_fields": golden}

            if result_a:
                golden_diff_a = _compute_field_diff(golden_wrapper, result_a)
                matches = sum(1 for d in golden_diff_a if d["match"])
                golden_score_a = round(matches / max(len(golden_diff_a), 1), 2)

            if result_b:
                golden_diff_b = _compute_field_diff(golden_wrapper, result_b)
                matches = sum(1 for d in golden_diff_b if d["match"])
                golden_score_b = round(matches / max(len(golden_diff_b), 1), 2)
        except json.JSONDecodeError:
            pass  # Ignore invalid golden fields

    return {
        "result_a": result_a,
        "result_b": result_b,
        "diff": diff_ab,
        "golden_comparison": {
            "score_a": golden_score_a,
            "score_b": golden_score_b,
            "diff_a": golden_diff_a,
            "diff_b": golden_diff_b,
        } if golden_fields else None,
        "errors": errors if errors else None,
    }


# ---------------------------------------------------------------------------
# AI-Powered Prompt Improvement
# ---------------------------------------------------------------------------

class PromptImproveRequest(BaseModel):
    use_case: str
    stage: str = "extraction"
    feedback: str
    current_output: Optional[str] = None
    model: str = "Qwen/Qwen2.5-VL-72B-Instruct"


META_PROMPT_TEMPLATE = """You are an expert prompt engineer specializing in document extraction prompts for vision-language models.

## Current Extraction Prompt
```
{current_prompt}
```

## User Feedback
The user reports the following issues with the current extraction output:
```
{feedback}
```

{output_section}

## Your Task
Based on the user's feedback, generate an IMPROVED version of the extraction prompt that addresses the reported issues. The improved prompt must:
1. Keep the same general structure and JSON output format as the original
2. Address each specific issue mentioned in the feedback
3. Add clearer instructions or constraints where the original was ambiguous
4. Maintain any field names and output schema from the original prompt

Return your response in this exact JSON format:
{{
  "improved_prompt": "<the full improved extraction prompt>",
  "changes_summary": "<bullet-point summary of what you changed and why>"
}}

Return ONLY the JSON, no other text."""


@prompts_router.post("/api/v1/prompts/improve")
async def api_improve_prompt(body: PromptImproveRequest, api_key: str = Depends(check_rate_limit)):
    """Use AI to generate an improved prompt based on user feedback.

    Auto-loads the current active prompt for the given (use_case, stage),
    sends it along with user feedback to the VLM, and returns the
    improved prompt plus a summary of changes.
    """
    if not HF_TOKEN:
        raise HTTPException(status_code=503, detail="HF_TOKEN not configured")

    # 1. Auto-load the current active prompt
    current_prompt = None
    if is_enabled():
        current_prompt = get_active_prompt(body.use_case, body.stage)

    if not current_prompt:
        # Fall back to hardcoded PROMPT_LIBRARY
        current_prompt = PROMPT_LIBRARY.get(body.stage, {}).get(
            body.use_case,
            PROMPT_LIBRARY.get(body.stage, {}).get("_default", ""),
        )

    if not current_prompt:
        raise HTTPException(
            status_code=404,
            detail=f"No prompt found for use_case={body.use_case}, stage={body.stage}",
        )

    # 2. Build the meta-prompt
    output_section = ""
    if body.current_output:
        # Truncate if very long to stay within token limits
        output_text = body.current_output[:3000]
        output_section = f"""## Current Extraction Output (for reference)
```json
{output_text}
```"""

    meta_prompt = META_PROMPT_TEMPLATE.format(
        current_prompt=current_prompt,
        feedback=body.feedback,
        output_section=output_section,
    )

    # 3. Call the VLM (text-only, no image needed)
    from hf_client import create_inference_client, robust_chat_completion

    client = create_inference_client(api_key=HF_TOKEN)

    try:
        response = robust_chat_completion(
            client,
            model=body.model,
            messages=[{"role": "user", "content": meta_prompt}],
            max_tokens=4096,
        )
        result_text = response.choices[0].message.content
    except Exception as e:
        logger.error(f"Prompt improvement VLM call failed: {e}")
        raise HTTPException(
            status_code=502,
            detail=f"VLM call failed: {str(e)[:200]}",
        )

    # 4. Parse the response
    try:
        parsed = _parse_vlm_json(result_text)
        improved_prompt = parsed.get("improved_prompt", "")
        changes_summary = parsed.get("changes_summary", "")
    except Exception:
        # If parsing fails, treat the whole response as the improved prompt
        improved_prompt = result_text.strip()
        changes_summary = "AI generated an improved prompt but structured parsing failed."

    if not improved_prompt:
        raise HTTPException(
            status_code=500,
            detail="VLM returned empty improved prompt",
        )

    return {
        "current_prompt": current_prompt,
        "improved_prompt": improved_prompt,
        "changes_summary": changes_summary,
        "model": body.model,
        "use_case": body.use_case,
        "stage": body.stage,
    }