admin / Apex

"""Background jobs: document extraction and AI generation.

Each task opens its own DB session. Tasks degrade gracefully when AI is
disabled (placeholder output) so the pipeline always completes.
"""
from __future__ import annotations

from app.ai import load_prompt
from app.core.database import SessionLocal
from app.models import (
    Bid,
    BidRequirement,
    BidResponse,
    Capability,
    CapabilityResponse,
    ComplianceFramework,
    Document,
)
from app.services import ai
from app.workers.celery_app import celery


@celery.task(name="extract_capabilities")
def extract_capabilities(document_ids: list[int], framework_ids: list[int]) -> dict:
    db = SessionLocal()
    try:
        docs = [db.get(Document, d) for d in document_ids]
        docs = [d for d in docs if d]
        corpus = "\n\n".join((d.text_content or "")[:40000] for d in docs)
        system = load_prompt("capability_extraction")
        items, meta = ai.extract_json(system, f"SOURCE DOCUMENTS:\n{corpus}")
        if not items and not ai.ai_available():
            items = [{
                "name": "Sample capability (AI disabled)",
                "category": "Service Delivery",
                "description": "Placeholder — configure ANTHROPIC_API_KEY to extract real capabilities from the uploaded documents.",
                "evidence": ["(no AI)"],
            }]
        created = []
        for it in items if isinstance(items, list) else []:
            cap = Capability(
                name=it.get("name", "Capability")[:255],
                category=it.get("category", ""),
                description=it.get("description", ""),
                evidence=it.get("evidence", []),
                status="extracted",
                source_document_ids=document_ids,
            )
            db.add(cap)
            db.flush()
            for fw_id in framework_ids:
                _generate_capability_response(db, cap, fw_id)
            created.append(cap.id)
        db.commit()
        return {"created": created, "meta": meta}
    finally:
        db.close()


def _generate_capability_response(db, cap: Capability, framework_id: int) -> None:
    fw = db.get(ComplianceFramework, framework_id)
    system = load_prompt("capability_response")
    prompt = (
        f"CAPABILITY: {cap.name}\nCATEGORY: {cap.category}\n"
        f"DESCRIPTION: {cap.description}\nEVIDENCE: {cap.evidence}\n"
        f"FRAMEWORK: {fw.name if fw else 'general'} — {fw.description if fw else ''}"
    )
    text, meta = ai.generate_text(system, prompt, max_tokens=1200)
    db.add(CapabilityResponse(
        capability_id=cap.id,
        framework_id=framework_id,
        response_text=text,
        model=meta.get("model", ""),
        prompt_version=meta.get("prompt_version", ""),
    ))


@celery.task(name="regenerate_capability_response")
def regenerate_capability_response(response_id: int) -> dict:
    db = SessionLocal()
    try:
        resp = db.get(CapabilityResponse, response_id)
        if not resp:
            return {"error": "not found"}
        cap = db.get(Capability, resp.capability_id)
        fw = db.get(ComplianceFramework, resp.framework_id) if resp.framework_id else None
        system = load_prompt("capability_response")
        prompt = (
            f"CAPABILITY: {cap.name}\nDESCRIPTION: {cap.description}\n"
            f"EVIDENCE: {cap.evidence}\nFRAMEWORK: {fw.name if fw else 'general'}"
        )
        text, meta = ai.generate_text(system, prompt, max_tokens=1200)
        resp.response_text = text
        resp.model = meta.get("model", "")
        db.commit()
        return {"ok": True}
    finally:
        db.close()


@celery.task(name="extract_requirements")
def extract_requirements(bid_id: int, document_ids: list[int]) -> dict:
    db = SessionLocal()
    try:
        docs = [db.get(Document, d) for d in document_ids]
        docs = [d for d in docs if d]
        corpus = "\n\n".join((d.text_content or "")[:40000] for d in docs)
        system = load_prompt("requirement_extraction")
        items, meta = ai.extract_json(system, f"TENDER DOCUMENTS:\n{corpus}")
        if not items and not ai.ai_available():
            items = [{
                "ref": "Q1", "type": "question", "mandatory": True,
                "question_text": "Describe your service management approach (AI disabled — placeholder).",
                "limits": "", "weighting": "", "source": "(no AI)",
            }]
        created = []
        for it in items if isinstance(items, list) else []:
            req = BidRequirement(
                bid_id=bid_id,
                ref=str(it.get("ref", ""))[:80],
                type=it.get("type", "question"),
                question_text=it.get("question_text", ""),
                limits=str(it.get("limits", "")),
                weighting=str(it.get("weighting", "")),
                mandatory=bool(it.get("mandatory", False)),
                source=str(it.get("source", "")),
                status="confirmed",
            )
            db.add(req)
            db.flush()
            created.append(req.id)
        db.commit()
        return {"created": created, "meta": meta}
    finally:
        db.close()


@celery.task(name="generate_response")
def generate_response(requirement_id: int) -> dict:
    db = SessionLocal()
    try:
        req = db.get(BidRequirement, requirement_id)
        if not req:
            return {"error": "not found"}
        bid = db.get(Bid, req.bid_id)
        # ground in approved capabilities (simple keyword match on category/name)
        caps = db.query(Capability).filter(Capability.status == "approved").all()
        evidence_blocks = []
        used = []
        for cap in caps[:12]:
            evidence_blocks.append(f"- {cap.name} ({cap.category}): {cap.description}")
            used.append(cap.id)
        fw_names = []
        for fw_id in (bid.frameworks or []):
            fw = db.get(ComplianceFramework, fw_id)
            if fw:
                fw_names.append(fw.name)
        system = load_prompt("response_generation")
        limit_txt = req.limits
        if req.word_limit:
            limit_txt = f"{req.limits} (hard limit: {req.word_limit} words — do not exceed)".strip()
        prompt = (
            f"TENDER REQUIREMENT ({req.ref}): {req.question_text}\n"
            f"LIMITS: {limit_txt}\nFRAMEWORKS: {', '.join(fw_names) or 'general'}\n"
            f"WIN THEMES: {bid.win_themes}\n\n"
            f"APPROVED CAPABILITY EVIDENCE:\n" + ("\n".join(evidence_blocks) or "(none approved yet)")
        )
        text, meta = ai.generate_text(system, prompt, max_tokens=2500)
        gaps = ""
        if "GAPS:" in text:
            body, _, gap_part = text.partition("GAPS:")
            text, gaps = body.strip(), gap_part.strip()
        resp = db.query(BidResponse).filter(BidResponse.requirement_id == requirement_id).first()
        if not resp:
            resp = BidResponse(requirement_id=requirement_id)
            db.add(resp)
        resp.draft_text = text
        resp.capabilities_used = used
        resp.gaps = gaps
        resp.model_meta = meta
        resp.status = "draft"
        db.commit()
        return {"ok": True, "gaps": bool(gaps)}
    finally:
        db.close()