AIRegulation-DocAnalysis/backend/app/application/compliance/pipeline.py

"""Compliance analysis pipeline helpers.

All functions are synchronous — call them via asyncio.to_thread() in async SSE generators.
"""

from __future__ import annotations

import asyncio
import json
import os
import re
import tempfile
from typing import TYPE_CHECKING

from loguru import logger
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential

# Shared retry policy for LLM calls: 3 attempts, exponential back-off 1–4 s.
_llm_retry = retry(
    stop=stop_after_attempt(3),
    wait=wait_exponential(multiplier=1, min=1, max=4),
    retry=retry_if_exception_type((ValueError, TimeoutError, ConnectionError)),
    reraise=True,
)

if TYPE_CHECKING:
    from app.application.knowledge import KnowledgeRetrievalService
    from app.domain.retrieval import RetrievedChunk
    from app.domain.compliance.ports import AnalysisRecord, FindingRecord
    from app.services.llm.base_client import BaseLLMClient


def _extract_json(text: str):
    """Extract JSON from LLM response, tolerating markdown wrappers."""
    stripped = text.strip()
    match = re.search(r"```(?:json)?\s*([\s\S]*?)```", stripped)
    if match:
        stripped = match.group(1).strip()
    try:
        return json.loads(stripped)
    except json.JSONDecodeError:
        pass
    for pattern in (r"(\[[\s\S]*\])", r"(\{[\s\S]*\})"):
        m = re.search(pattern, stripped)
        if m:
            try:
                return json.loads(m.group(1))
            except json.JSONDecodeError:
                continue
    raise ValueError(f"No valid JSON found in LLM response: {text[:300]}")


def extract_text_from_doc_id(doc_id: str) -> str:
    """Fetch the full text of a document by retrieving its chunks filtered by doc_id.

    Uses a high top_k and doc_id filter to reconstruct the document in chunk order,
    avoiding the previous approach of semantic search by doc_name which could return
    chunks from unrelated documents.
    """
    from app.shared.bootstrap import get_document_query_service, get_retrieval_service
    doc = get_document_query_service().get(doc_id)
    if not doc:
        raise ValueError(f"Document '{doc_id}' not found")
    service = get_retrieval_service()
    # Use doc_name as a broad query, filter strictly by doc_id so we only get
    # this document's chunks; top_k=100 covers most real-world documents.
    chunks = service.retrieve(query=doc.doc_name, top_k=100, filters=doc_id)
    doc_chunks = [c for c in chunks if getattr(c, "doc_id", None) == doc_id]
    if not doc_chunks:
        # Fallback: use top results even without doc_id match (e.g., legacy store)
        doc_chunks = chunks[:30]
    # Sort by chunk_index to preserve document reading order
    doc_chunks.sort(key=lambda c: getattr(c, "chunk_index", 0))
    return "\n\n".join(c.text for c in doc_chunks[:40])


def extract_text_from_file(content: bytes, filename: str) -> str:
    """Parse an uploaded file and return its full text content.

    Removed previous 4000-char cap so large specifications and standards are
    fully analysed. The caller is responsible for splitting the text into
    clause-sized chunks before passing to the LLM.
    """
    from app.shared.bootstrap import get_document_command_service
    suffix = os.path.splitext(filename or "doc.pdf")[1] or ".pdf"
    tmp_path = ""
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            tmp.write(content)
            tmp_path = tmp.name
        service = get_document_command_service()
        parsed = service.parser.parse(file_path=tmp_path, doc_id="tmp_analysis", doc_name=filename)
        if parsed.raw_text:
            # Return full text — truncation happens in split_into_clauses()
            return parsed.raw_text
        return "\n".join(
            b.get("text", "") for b in parsed.semantic_blocks if b.get("text")
        )
    except Exception as exc:
        logger.warning("File text extraction failed: {}", exc)
        return ""
    finally:
        if tmp_path:
            try: os.unlink(tmp_path)
            except OSError: pass


def split_into_clauses(text: str, client: "BaseLLMClient") -> list[str]:
    """Split a compliance document into semantically independent clauses.

    For long texts (> 2 000 chars) the document is processed in overlapping
    2 000-char windows so no content is missed. Each window produces up to 4
    clauses; results are deduplicated and capped at 12 total to keep analysis
    latency reasonable.
    """
    # Window size and step for sliding-window clause extraction
    _WINDOW = 2000
    _STEP = 1800  # 200-char overlap to avoid cutting clauses at boundaries
    _MAX_CLAUSES = 12

    windows = []
    if len(text) <= _WINDOW:
        windows = [text]
    else:
        pos = 0
        while pos < len(text):
            windows.append(text[pos: pos + _WINDOW])
            pos += _STEP

    all_clauses: list[str] = []
    for window in windows:
        prompt = (
            "You are a compliance analysis expert. Split the following text into "
            "3-4 semantically complete compliance clauses. Each clause must be an "
            "independent requirement or technical statement. Omit section headings, "
            "definitions, and non-normative text.\n"
            "Return as JSON array of strings, e.g.:\n"
            '["Clause one...", "Clause two..."]\n'
            "Return ONLY the JSON array.\n\n"
            f"Text:\n{window}"
        )
        response = client.chat([{"role": "user", "content": prompt}], max_tokens=800)
        if response.is_success:
            try:
                result = _extract_json(response.content)
                if isinstance(result, list):
                    clauses = [str(c).strip() for c in result if str(c).strip()]
                    all_clauses.extend(clauses[:4])
            except (ValueError, TypeError):
                logger.warning("Clause split JSON parse failed for window, using sentence fallback")
                sentences = re.split(r"[.?!;\n]+", window)
                all_clauses.extend(s.strip() for s in sentences if len(s.strip()) > 20)
        else:
            # LLM unavailable — fall back to sentence splitting for this window
            sentences = re.split(r"[.?!;\n]+", window)
            all_clauses.extend(s.strip() for s in sentences if len(s.strip()) > 20)

        if len(all_clauses) >= _MAX_CLAUSES:
            break

    # Deduplicate near-duplicates (same first 80 chars) that span window boundaries
    seen: set[str] = set()
    deduped: list[str] = []
    for c in all_clauses:
        key = c[:80].lower()
        if key not in seen:
            seen.add(key)
            deduped.append(c)

    return deduped[:_MAX_CLAUSES]


def retrieve_for_clause(
    clause: str,
    retrieval_service: "KnowledgeRetrievalService",
    top_k: int = 5,
    domains: str | None = None,
) -> list["RetrievedChunk"]:
    """Retrieve regulation chunks relevant to a clause.

    If the best retrieval score is below 0.55, rewrite the clause into a more
    technical query and retry once to improve coverage.
    """
    chunks = retrieval_service.retrieve(query=clause, top_k=top_k, filters=domains)
    if not chunks:
        return chunks

    best_score = max((getattr(c, "score", 0) for c in chunks), default=0)
    if best_score < 0.55:
        # Rewrite clause as technical keyword query and retry
        keywords = " ".join(
            w for w in re.split(r"\W+", clause) if len(w) > 3
        )[:200]
        retry_chunks = retrieval_service.retrieve(query=keywords, top_k=top_k, filters=domains)
        if retry_chunks:
            # Merge: keep unique chunks, prefer higher-score version
            seen_ids: set[str] = {getattr(c, "chunk_id", str(i)) for i, c in enumerate(chunks)}
            for rc in retry_chunks:
                rid = getattr(rc, "chunk_id", "")
                if rid not in seen_ids:
                    chunks.append(rc)
                    seen_ids.add(rid)
            chunks.sort(key=lambda c: getattr(c, "score", 0), reverse=True)
            chunks = chunks[:top_k]
    return chunks


def process_single_clause(
    clause: str,
    index: int,
    retrieval_service: "KnowledgeRetrievalService",
    client: "BaseLLMClient",
    top_k: int = 5,
    domains: str | None = None,
) -> dict:
    """Process one clause: retrieve relevant regulations then check compliance.

    Returns a dict with keys:
      - index: clause position (for ordering)
      - chunks: list of RetrievedChunk (for source events)
      - finding: dict with title/desc/status/clause_ref/confidence (may be None on LLM failure)

    Designed to run inside asyncio.to_thread() for parallel execution.
    The finding now includes a 'source_refs' list linking back to the chunks
    that informed the verdict, enabling the frontend to correlate sources with findings.
    """
    chunks = retrieve_for_clause(clause, retrieval_service, top_k, domains)
    finding = check_clause_compliance(clause, chunks, client)
    if finding is not None:
        # Attach source references so the frontend can link finding ↔ sources
        finding["source_refs"] = [
            {
                "standard": getattr(c, "doc_title", "") or getattr(c, "doc_name", ""),
                "clause": getattr(c, "section_title", "") or "",
                "score": round(float(getattr(c, "score", 0)), 3),
            }
            for c in chunks[:3]
        ]
    return {"index": index, "chunks": chunks, "finding": finding}


async def run_clauses_streaming(
    clauses: list[str],
    retrieval_service: "KnowledgeRetrievalService",
    client: "BaseLLMClient",
    top_k: int = 5,
    domains: str | None = None,
):
    """Process all clauses concurrently and yield each result as it completes.

    Unlike the old gather()-based approach, this uses asyncio.Queue so that
    findings are emitted to the SSE stream immediately when each clause
    finishes — the user sees results progressively rather than waiting for
    the slowest clause before seeing any output.

    Yields dicts with keys: index, chunks, finding (same schema as
    process_single_clause, plus a sentinel {"_done": True} at the end).
    """
    queue: asyncio.Queue[dict] = asyncio.Queue()
    total = len(clauses)

    async def _worker(clause: str, i: int) -> None:
        """Run one clause in a thread and push the result into the queue."""
        try:
            result = await asyncio.to_thread(
                process_single_clause,
                clause, i, retrieval_service, client, top_k, domains,
            )
        except Exception as exc:
            logger.warning("Clause {} processing failed: {}", i, exc)
            result = {"index": i, "chunks": [], "finding": None}
        await queue.put(result)

    # Launch all workers concurrently
    tasks = [asyncio.create_task(_worker(clause, i)) for i, clause in enumerate(clauses)]

    received = 0
    while received < total:
        result = await queue.get()
        yield result
        received += 1

    # Wait for all tasks to complete (they should already be done by now)
    await asyncio.gather(*tasks, return_exceptions=True)


async def run_clauses_parallel(
    clauses: list[str],
    retrieval_service: "KnowledgeRetrievalService",
    client: "BaseLLMClient",
    top_k: int = 5,
    domains: str | None = None,
) -> list[dict]:
    """Legacy batch API kept for backward compatibility.

    Collects all streaming results and returns them sorted by clause index.
    New code should use run_clauses_streaming() directly.
    """
    results: list[dict] = []
    async for result in run_clauses_streaming(clauses, retrieval_service, client, top_k, domains):
        results.append(result)
    return sorted(results, key=lambda r: r["index"])


def check_clause_compliance(
    clause: str,
    chunks: list["RetrievedChunk"],
    client: "BaseLLMClient",
) -> dict | None:
    """Check whether a business clause complies with the retrieved regulations.

    The prompt explicitly instructs the LLM to:
    - extract clause_ref from the retrieved text (not invent it)
    - include a confidence score (0-1) reflecting how well the retrieved
      chunks cover the clause topic

    Returns None only when the LLM call fails after all retries.
    """
    reg_context = "\n".join(
        f"[{i+1}] {c.doc_title} {c.section_title or ''}: {c.text[:300]}"
        for i, c in enumerate(chunks[:5])
    ) if chunks else "(no regulatory context retrieved)"
    prompt = (
        "You are a compliance expert. Judge whether the following business clause "
        "complies with the retrieved regulations.\n\n"
        f"Business clause:\n{clause}\n\n"
        f"Retrieved regulations:\n{reg_context}\n\n"
        "Return JSON with these exact fields:\n"
        "{\n"
        '  "status": "ok" | "warn" | "risk",\n'
        '  "title": "Short finding title (max 30 chars)",\n'
        '  "desc": "Description (50-120 chars)",\n'
        '  "clause_ref": "Exact clause/article reference copied from the retrieved text above, '
        'e.g. Art.9.1 or Sec.3.1. Use null if no specific clause number appears in the retrieved text.",\n'
        '  "confidence": 0.0-1.0  // how well the retrieved context covers this clause topic\n'
        "}\n"
        "status: ok=compliant, warn=gap exists, risk=critical/missing\n"
        "IMPORTANT: copy clause_ref verbatim from the retrieved text; do NOT invent references.\n"
        "Return ONLY the JSON object."
    )

    def _do_check():
        resp = client.chat([{"role": "user", "content": prompt}], max_tokens=500)
        if not resp.is_success:
            raise ValueError("LLM returned non-success for gap check")
        return resp

    try:
        response = _llm_retry(_do_check)()
    except Exception as exc:
        logger.warning("check_clause_compliance LLM call failed after retries: {}", exc)
        return None

    try:
        result = _extract_json(response.content)
        if isinstance(result, dict) and "status" in result:
            return {
                "title": str(result.get("title", "Compliance finding")),
                "desc": str(result.get("desc", "")),
                "status": result.get("status", "info"),
                # None if LLM correctly found no clause number in retrieved text
                "clause_ref": result.get("clause_ref") or None,
                # Confidence score helps frontend show retrieval quality indicator
                "confidence": float(result.get("confidence", 0.5)),
            }
    except (ValueError, TypeError) as exc:
        logger.warning("Gap check JSON parse failed: {}", exc)
    return None


def synthesize_conclusion(
    para_text: str,
    findings: list[dict],
    client: "BaseLLMClient",
) -> dict:
    if not findings:
        return {
            "conclusion": "No significant compliance gaps found. Continue monitoring regulation updates.",
            "actions": [{"label": "Next action", "value": "Monitor regulation updates"}],
            "risk_score": 10,
            "highlight_terms": [],
            "para_text": para_text[:800],
        }
    findings_text = "\n".join(
        f"- [{f['status'].upper()}] {f['title']}: {f['desc']}"
        for f in findings
    )
    prompt = (
        "You are a compliance analysis expert. Generate a summary report "
        "based on the following compliance findings.\n\n"
        f"Original text (first 600 chars):\n{para_text[:600]}\n\n"
        f"Findings:\n{findings_text}\n\n"
        "Return JSON:\n"
        "{\n"
        '  "conclusion": "Overall compliance conclusion (100-200 chars)",\n'
        '  "actions": [\n'
        '    {"label": "Action label", "value": "Description"},\n'
        '    {"label": "Priority", "value": "High/Medium/Low", "risk": true}\n'
        '  ],\n'
        '  "risk_score": 0-100 (integer, higher=riskier),\n'
        '  "highlight_terms": ["term1", "term2"],  // up to 10 key technical/legal terms actually present in the text\n'
        '  "para_text": "Original text or summary (max 600 chars)"\n'
        "}\n"
        "Return ONLY the JSON object."
    )
    fallback = {
        "conclusion": "Compliance analysis complete. Review findings and create remediation plan.",
        "actions": [
            {"label": "Next action", "value": "Review critical findings"},
            {"label": "Escalation", "value": "Legal review required", "risk": True},
        ],
        "risk_score": 60,
        "highlight_terms": [],
        "para_text": para_text[:800],
    }

    def _do_synthesize():
        resp = client.chat([{"role": "user", "content": prompt}], max_tokens=1200)
        if not resp.is_success:
            raise ValueError("LLM returned non-success for synthesis")
        return resp

    try:
        response = _llm_retry(_do_synthesize)()
    except Exception as exc:
        logger.warning("synthesize_conclusion LLM call failed after retries: {}", exc)
        return fallback

    try:
        result = _extract_json(response.content)
        if isinstance(result, dict):
            return {
                "conclusion": str(result.get("conclusion", fallback["conclusion"])),
                "actions": result.get("actions", fallback["actions"]),
                "risk_score": int(result.get("risk_score", 60)),
                "highlight_terms": result.get("highlight_terms", []),
                "para_text": str(result.get("para_text", para_text[:800])),
            }
    except (ValueError, TypeError) as exc:
        logger.warning("Conclusion synthesis JSON parse failed: {}", exc)
    return fallback


_SUGGESTION_FOCUS = {
    "risk": "Focus on remediation steps, required certifications, and timeline to resolve.",
    "warn": "Focus on identifying the specific compliance gap and how to close it.",
    "ok": "Focus on maintaining compliance evidence and monitoring future changes.",
}

_SUGGESTION_FALLBACK = {
    "risk": [
        "What specific certifications or documents are required to remediate this finding?",
        "What is the typical remediation timeline for this type of non-compliance?",
        "Which regulation clause defines the exact requirement?",
    ],
    "warn": [
        "What is the exact gap between the current state and the requirement?",
        "What evidence would demonstrate partial compliance?",
        "Which regulation clause applies to this warning?",
    ],
    "ok": [
        "What documentation should be maintained to evidence this compliance?",
        "How should this area be monitored as regulations evolve?",
        "Are there related clauses that may affect this compliant area?",
    ],
}


def build_finding_context(finding: "FindingRecord", analysis: "AnalysisRecord") -> str:
    """Build a grounded system context string for a finding chat thread.

    Combines finding details with analysis metadata so the LLM has full
    context without relying on the frontend to pass segment_context.
    """
    return (
        f"Document: {analysis.doc_name}\n"
        f"Standard: {analysis.standard_name}\n"
        f"Finding [{finding.seq + 1}]: {finding.title}\n"
        f"Status: {finding.status}\n"
        f"Clause reference: {finding.clause_ref or 'N/A'}\n"
        f"Description: {finding.description}\n"
        f"Overall conclusion: {analysis.conclusion}\n"
    )


def generate_suggestions(
    finding: "FindingRecord",
    analysis: "AnalysisRecord",
    client: "BaseLLMClient",
) -> list[str]:
    """Generate 3 context-aware follow-up questions for a finding chat thread.

    Returns exactly 3 question strings. Falls back to static templates on error.
    """
    fallback = _SUGGESTION_FALLBACK.get(finding.status, _SUGGESTION_FALLBACK["warn"])
    context = build_finding_context(finding, analysis)
    focus = _SUGGESTION_FOCUS.get(finding.status, _SUGGESTION_FOCUS["warn"])
    prompt = (
        f"{context}\n\n"
        f"Task: {focus}\n"
        "Generate exactly 3 concise follow-up questions a compliance analyst would ask.\n"
        'Return JSON: {"questions": ["question 1", "question 2", "question 3"]}\n'
        "Return ONLY the JSON object."
    )
    response = client.chat([{"role": "user", "content": prompt}], max_tokens=300)
    if not response.is_success:
        return fallback
    try:
        result = _extract_json(response.content)
        questions = result.get("questions", [])
        if isinstance(questions, list) and len(questions) >= 3:
            return [str(q) for q in questions[:3]]
    except (ValueError, TypeError) as exc:
        logger.warning("generate_suggestions JSON parse failed: {}", exc)
    return fallback


def detect_cross_clause_conflicts(
    findings: list[dict],
    client: "BaseLLMClient",
) -> list[dict]:
    """Detect contradictions and missing cross-references across all findings.

    Runs a single LLM call after all per-clause findings are collected.
    Returns a list of conflict dicts: {type, finding_a, finding_b, desc}.
    Returns an empty list on LLM failure so the caller can proceed without it.
    """
    if len(findings) < 2:
        # Need at least 2 findings to compare
        return []

    findings_text = "\n".join(
        f"[{i+1}] [{f['status'].upper()}] {f['title']}: {f['desc']}"
        + (f" (Ref: {f['clause_ref']})" if f.get("clause_ref") else "")
        for i, f in enumerate(findings)
    )
    prompt = (
        "You are a compliance expert. Review the following compliance findings from the same document "
        "and identify any cross-clause issues:\n\n"
        f"Findings:\n{findings_text}\n\n"
        "Return JSON array of conflicts (empty array [] if none found):\n"
        "[\n"
        "  {\n"
        '    "type": "contradiction" | "missing_ref" | "cumulative_risk",\n'
        '    "finding_a": <1-based index>,\n'
        '    "finding_b": <1-based index or null>,\n'
        '    "desc": "Brief description of the cross-clause issue (max 100 chars)"\n'
        "  }\n"
        "]\n"
        "Return ONLY the JSON array."
    )
    try:
        response = client.chat([{"role": "user", "content": prompt}], max_tokens=600)
        if not response.is_success:
            return []
        result = _extract_json(response.content)
        if isinstance(result, list):
            return [
                {
                    "type": str(c.get("type", "contradiction")),
                    "finding_a": int(c.get("finding_a", 0)),
                    "finding_b": c.get("finding_b"),
                    "desc": str(c.get("desc", "")),
                }
                for c in result
                if isinstance(c, dict)
            ]
    except Exception as exc:
        logger.warning("detect_cross_clause_conflicts failed: {}", exc)
    return []