import httpx
from sqlmodel import Session, select
from app.config import settings
from app.models import Assessment, Control, Evidence, Policy, AILog, User

SYSTEM = "You are SecureFlow's compliance assistant. Use only supplied internal records. Do not claim certification. If evidence is absent, say it is missing. Include citations like [Control:AC-1] or [Evidence:12]."


def collect_sources(session: Session, tenant_id: int, prompt: str, related_control: str | None = None):
    controls = session.exec(select(Control).limit(30)).all()
    assessments = session.exec(select(Assessment).where(Assessment.tenant_id == tenant_id).limit(80)).all()
    evidence = session.exec(select(Evidence).where(Evidence.tenant_id == tenant_id).limit(30)).all()
    policies = session.exec(select(Policy).where(Policy.tenant_id == tenant_id).limit(20)).all()
    items = []
    for c in controls:
        if not related_control or related_control.lower() in c.code.lower() or related_control.lower() in c.title.lower() or related_control.lower() in c.requirement.lower():
            a = next((x for x in assessments if x.control_id == c.id), None)
            items.append({"type": "control", "id": c.id, "code": c.code, "title": c.title, "requirement": c.requirement, "assessment": a.answer.value if a else "Not assessed", "notes": a.notes if a else ""})
    for e in evidence:
        items.append({"type": "evidence", "id": e.id, "title": e.title, "tags": e.tags, "control_id": e.control_id, "review_date": str(e.review_date)})
    for p in policies:
        items.append({"type": "policy", "id": p.id, "title": p.title, "version": p.version, "category": p.category})
    return items[:60]


async def generate_answer(session: Session, user: User, prompt: str, related_control: str | None = None) -> AILog:
    sources = collect_sources(session, user.tenant_id, prompt, related_control)
    source_text = "\n".join(str(s) for s in sources)
    fallback = "Draft answer based on internal records only:\n"
    if not sources:
        fallback += "No matching internal records were found. Treat this as a gap requiring human review."
    else:
        citations = ", ".join(f"[{s['type'].title()}:{s['id']}]" for s in sources[:8])
        fallback += f"Relevant records reviewed: {citations}. "
        gaps = [s for s in sources if s.get("assessment") in ("No", "Partial", "Not assessed")]
        fallback += "Current posture is presented as readiness and gap analysis only. "
        fallback += "Potential gaps require remediation before using this answer externally." if gaps else "Records indicate alignment for the cited controls, subject to qualified review."
    response = fallback
    if settings.openai_api_key:
        async with httpx.AsyncClient(timeout=30) as client:
            r = await client.post(
                f"{settings.openai_base_url.rstrip('/')}/chat/completions",
                headers={"Authorization": f"Bearer {settings.openai_api_key}"},
                json={"model": settings.openai_model, "messages": [{"role": "system", "content": SYSTEM}, {"role": "user", "content": f"Prompt: {prompt}\nInternal records:\n{source_text}"}], "temperature": 0.2},
            )
            r.raise_for_status()
            response = r.json()["choices"][0]["message"]["content"]
    log = AILog(tenant_id=user.tenant_id, user_id=user.id, prompt=prompt, response=response, source_records=sources)
    session.add(log)
    session.commit()
    session.refresh(log)
    return log