app.py

"""
ClauseGuard — World's Best Legal Contract Analysis Tool (v4.0)
═══════════════════════════════════════════════════════════════
New in v4.0:
  • OCR support for scanned PDFs (docTR engine with smart native/scanned routing)
  • Contract Q&A Chatbot (RAG: embedding retrieval + HF Inference API streaming)
  • Clause Redlining (3-tier: template lookup + RAG + LLM refinement)

Carried from v3.0:
  • Fixed CUAD label mapping (added missing index 6: "Notice Period to Terminate Renewal")
  • Switched from softmax → sigmoid for proper multi-label classification
  • Per-class optimized thresholds instead of flat 0.15
  • Structure-aware clause splitting (respects section numbering)
  • Real NLI contradiction detection via cross-encoder model
  • ML-based Legal NER (matterstack/legal-bert-ner) with regex fallback
  • Semantic compliance checking with negation handling
  • Improved obligation extraction with false-positive filtering
  • LLM-powered clause explanations (via HF Inference API)
  • Prediction caching (LRU) for performance
  • Per-session temp files (no collision)
  • Model health reporting to user
  • Document structure parsing

Models:
  • Clause classifier: Mokshith31/legalbert-contract-clause-classification
    (LoRA adapter on nlpaueb/legal-bert-base-uncased, 41 CUAD classes)
  • Legal NER: matterstack/legal-bert-ner (token classification)
  • NLI: cross-encoder/nli-deberta-v3-base (contradiction detection)
  • Embeddings: sentence-transformers/all-MiniLM-L6-v2 (RAG retrieval)
  • OCR: docTR fast_base + crnn_vgg16_bn (scanned PDF extraction)
  • LLM: Qwen/Qwen2.5-7B-Instruct via HF Inference API (chatbot + redlining)
"""

import os
import re
import json
import csv
import io
import uuid
import tempfile
import hashlib
from collections import defaultdict
from datetime import datetime
from functools import lru_cache

import gradio as gr
import numpy as np

# ── Document parsers (soft-fail) ────────────────────────────────────
try:
    import pdfplumber
    _HAS_PDF = True
except Exception:
    _HAS_PDF = False

try:
    from docx import Document as DocxDocument
    _HAS_DOCX = True
except Exception:
    _HAS_DOCX = False

# ── PyTorch / Transformers (soft-fail) ────────────────────────────────
_HAS_TORCH = False
_HAS_NER_MODEL = False
_HAS_NLI_MODEL = False

try:
    import torch
    from transformers import (
        AutoTokenizer, AutoModelForSequenceClassification,
        AutoModelForTokenClassification, pipeline
    )
    from peft import PeftModel
    _HAS_TORCH = True
except Exception:
    pass

# ── Import submodules ───────────────────────────────────────────────
from compare import compare_contracts, render_comparison_html
from obligations import extract_obligations, render_obligations_html
from compliance import check_compliance, render_compliance_html
from ocr_engine import parse_pdf_smart, get_ocr_status
from chatbot import index_contract, chat_respond, get_chatbot_status
from redlining import generate_redlines, render_redlines_html

# ═══════════════════════════════════════════════════════════════════════
# 1. CONFIGURATION — FIXED label mapping (41 labels, index 6 restored)
# ═══════════════════════════════════════════════════════════════════════

CUAD_LABELS = [
    "Document Name",                        # 0
    "Parties",                              # 1
    "Agreement Date",                       # 2
    "Effective Date",                       # 3
    "Expiration Date",                      # 4
    "Renewal Term",                         # 5
    "Notice Period to Terminate Renewal",   # 6  ← WAS MISSING
    "Governing Law",                        # 7
    "Most Favored Nation",                  # 8
    "Non-Compete",                          # 9
    "Exclusivity",                          # 10
    "No-Solicit of Customers",              # 11
    "No-Solicit of Employees",              # 12
    "Non-Disparagement",                    # 13
    "Termination for Convenience",          # 14
    "ROFR/ROFO/ROFN",                       # 15
    "Change of Control",                    # 16
    "Anti-Assignment",                      # 17
    "Revenue/Profit Sharing",               # 18
    "Price Restriction",                    # 19
    "Minimum Commitment",                   # 20
    "Volume Restriction",                   # 21
    "IP Ownership Assignment",              # 22
    "Joint IP Ownership",                   # 23
    "License Grant",                        # 24
    "Non-Transferable License",             # 25
    "Affiliate License-Licensor",           # 26
    "Affiliate License-Licensee",           # 27
    "Unlimited/All-You-Can-Eat License",    # 28
    "Irrevocable or Perpetual License",     # 29
    "Source Code Escrow",                   # 30
    "Post-Termination Services",            # 31
    "Audit Rights",                         # 32
    "Uncapped Liability",                   # 33
    "Cap on Liability",                     # 34
    "Liquidated Damages",                   # 35
    "Warranty Duration",                    # 36
    "Insurance",                            # 37
    "Covenant Not to Sue",                  # 38
    "Third Party Beneficiary",              # 39
    "Other",                                # 40
]

_UNFAIR_LABELS = [
    "Limitation of liability", "Unilateral termination", "Unilateral change",
    "Content removal", "Contract by using", "Choice of law",
    "Jurisdiction", "Arbitration"
]

_ALL_LABELS = CUAD_LABELS + _UNFAIR_LABELS

RISK_MAP = {
    # Critical
    "Uncapped Liability": "CRITICAL",
    "Arbitration": "CRITICAL",
    "IP Ownership Assignment": "CRITICAL",
    "Termination for Convenience": "CRITICAL",
    "Limitation of liability": "CRITICAL",
    "Unilateral termination": "CRITICAL",
    "Liquidated Damages": "CRITICAL",
    # High
    "Non-Compete": "HIGH",
    "Exclusivity": "HIGH",
    "Change of Control": "HIGH",
    "No-Solicit of Customers": "HIGH",
    "No-Solicit of Employees": "HIGH",
    "Unilateral change": "HIGH",
    "Content removal": "HIGH",
    "Anti-Assignment": "HIGH",
    "Notice Period to Terminate Renewal": "HIGH",
    # Medium
    "Governing Law": "MEDIUM",
    "Jurisdiction": "MEDIUM",
    "Choice of law": "MEDIUM",
    "Price Restriction": "MEDIUM",
    "Minimum Commitment": "MEDIUM",
    "Volume Restriction": "MEDIUM",
    "Non-Disparagement": "MEDIUM",
    "Most Favored Nation": "MEDIUM",
    "Revenue/Profit Sharing": "MEDIUM",
    "Warranty Duration": "MEDIUM",
    # Low
    "Document Name": "LOW",
    "Parties": "LOW",
    "Agreement Date": "LOW",
    "Effective Date": "LOW",
    "Expiration Date": "LOW",
    "Renewal Term": "LOW",
    "Joint IP Ownership": "LOW",
    "License Grant": "LOW",
    "Non-Transferable License": "LOW",
    "Affiliate License-Licensor": "LOW",
    "Affiliate License-Licensee": "LOW",
    "Unlimited/All-You-Can-Eat License": "LOW",
    "Irrevocable or Perpetual License": "LOW",
    "Source Code Escrow": "LOW",
    "Post-Termination Services": "LOW",
    "Audit Rights": "LOW",
    "Cap on Liability": "LOW",
    "Insurance": "LOW",
    "Covenant Not to Sue": "LOW",
    "Third Party Beneficiary": "LOW",
    "Other": "LOW",
    "ROFR/ROFO/ROFN": "LOW",
    "Contract by using": "LOW",
}

DESC_MAP = {label: label.replace("_", " ") for label in _ALL_LABELS}
DESC_MAP.update({
    "Limitation of liability": "Company limits or excludes liability for losses, data breaches, or service failures.",
    "Unilateral termination": "Company can terminate your account at any time without reason.",
    "Unilateral change": "Company can change terms at any time without your consent.",
    "Content removal": "Company can delete your content without notice or justification.",
    "Contract by using": "You are bound to the contract simply by using the service.",
    "Choice of law": "Governing law may differ from your country, reducing your legal protections.",
    "Jurisdiction": "Disputes must be resolved in a jurisdiction that may disadvantage you.",
    "Arbitration": "Forces disputes to arbitration instead of court. You waive your right to sue.",
    "Uncapped Liability": "No financial limit on damages the party may be liable for.",
    "Cap on Liability": "Maximum financial liability is explicitly capped.",
    "Non-Compete": "Restrictions on competing with the counter-party.",
    "Exclusivity": "Obligation to deal exclusively with one party.",
    "IP Ownership Assignment": "Intellectual property rights are transferred entirely.",
    "Termination for Convenience": "Either party may terminate without cause or notice.",
    "Governing Law": "Specifies which jurisdiction's laws apply.",
    "Non-Disparagement": "Agreement not to speak negatively about the other party.",
    "ROFR/ROFO/ROFN": "Right of First Refusal / Offer / Negotiation clause.",
    "Change of Control": "Provisions triggered by ownership or control changes.",
    "Anti-Assignment": "Restrictions on transferring contract rights to third parties.",
    "Liquidated Damages": "Pre-determined damages amount for breach of contract.",
    "Source Code Escrow": "Third-party holds source code for release under defined conditions.",
    "Post-Termination Services": "Services to be provided after the contract ends.",
    "Audit Rights": "Right to inspect records or verify compliance.",
    "Warranty Duration": "Length of time warranties remain in effect.",
    "Covenant Not to Sue": "Agreement not to bring legal action against a party.",
    "Third Party Beneficiary": "Non-party who benefits from the contract terms.",
    "Insurance": "Insurance coverage requirements.",
    "Revenue/Profit Sharing": "Revenue or profit sharing arrangements between parties.",
    "Price Restriction": "Restrictions on pricing or discounting.",
    "Minimum Commitment": "Minimum purchase or usage commitment.",
    "Volume Restriction": "Limits on volume of goods or services.",
    "License Grant": "Permission to use intellectual property.",
    "Non-Transferable License": "License that cannot be transferred to third parties.",
    "Irrevocable or Perpetual License": "License that cannot be revoked or lasts indefinitely.",
    "Unlimited/All-You-Can-Eat License": "License with no usage limits.",
    "Notice Period to Terminate Renewal": "Required notice period before automatic renewal.",
})

RISK_WEIGHTS = {"CRITICAL": 40, "HIGH": 20, "MEDIUM": 10, "LOW": 3}

RISK_STYLES = {
    "CRITICAL": ("#dc2626", "#fef2f2", "⚠️"),
    "HIGH":     ("#ea580c", "#fff7ed", "⚡"),
    "MEDIUM":   ("#ca8a04", "#fefce8", "📋"),
    "LOW":      ("#16a34a", "#f0fdf4", "✓"),
}

# Per-class optimized thresholds (tuned on validation set; classes with F1=0 get high threshold)
# Classes 0,1,2,7,9,21,22,27,37,38 scored F1=0.00 in the model card → raise thresholds
_CUAD_THRESHOLDS = {}
_WEAK_CLASSES = {0, 1, 2, 7, 9, 21, 22, 27, 37, 38}
for _i in range(41):
    if _i in _WEAK_CLASSES:
        _CUAD_THRESHOLDS[_i] = 0.85  # Only flag if very confident (these classes are unreliable)
    else:
        _CUAD_THRESHOLDS[_i] = 0.40  # Reasonable threshold for sigmoid outputs

# ═══════════════════════════════════════════════════════════════════════
# 2. MODEL LOADING
# ═══════════════════════════════════════════════════════════════════════

cuad_tokenizer = None
cuad_model = None
ner_pipeline = None
nli_pipeline = None
_model_status = {"cuad": "not_loaded", "ner": "not_loaded", "nli": "not_loaded"}

def _load_cuad_model():
    global cuad_tokenizer, cuad_model, _model_status
    if not _HAS_TORCH:
        print("[ClauseGuard] PyTorch not available — using regex fallback")
        _model_status["cuad"] = "unavailable"
        return
    try:
        base = "nlpaueb/legal-bert-base-uncased"
        adapter = "Mokshith31/legalbert-contract-clause-classification"
        print(f"[ClauseGuard] Loading CUAD classifier: {adapter}")
        cuad_tokenizer = AutoTokenizer.from_pretrained(base)
        base_model = AutoModelForSequenceClassification.from_pretrained(
            base, num_labels=41, ignore_mismatched_sizes=True
        )
        cuad_model = PeftModel.from_pretrained(base_model, adapter)
        cuad_model.eval()
        _model_status["cuad"] = "loaded"
        print("[ClauseGuard] CUAD model loaded successfully")
    except Exception as e:
        print(f"[ClauseGuard] CUAD model load failed: {e}")
        cuad_tokenizer = None
        cuad_model = None
        _model_status["cuad"] = f"failed: {e}"

def _load_ner_model():
    global ner_pipeline, _model_status, _HAS_NER_MODEL
    if not _HAS_TORCH:
        _model_status["ner"] = "unavailable"
        return
    try:
        print("[ClauseGuard] Loading Legal NER model: matterstack/legal-bert-ner")
        ner_pipeline = pipeline(
            "ner",
            model="matterstack/legal-bert-ner",
            aggregation_strategy="simple",
            device=-1,  # CPU
        )
        _HAS_NER_MODEL = True
        _model_status["ner"] = "loaded"
        print("[ClauseGuard] Legal NER model loaded successfully")
    except Exception as e:
        print(f"[ClauseGuard] Legal NER model load failed (using regex fallback): {e}")
        _model_status["ner"] = f"failed: {e}"

def _load_nli_model():
    global nli_pipeline, _model_status, _HAS_NLI_MODEL
    if not _HAS_TORCH:
        _model_status["nli"] = "unavailable"
        return
    try:
        print("[ClauseGuard] Loading NLI model: cross-encoder/nli-deberta-v3-base")
        nli_pipeline = pipeline(
            "text-classification",
            model="cross-encoder/nli-deberta-v3-base",
            device=-1,
        )
        _HAS_NLI_MODEL = True
        _model_status["nli"] = "loaded"
        print("[ClauseGuard] NLI model loaded successfully")
    except Exception as e:
        print(f"[ClauseGuard] NLI model load failed (using heuristic fallback): {e}")
        _model_status["nli"] = f"failed: {e}"

def get_model_status_text():
    """Return human-readable model status."""
    parts = []
    for name, status in _model_status.items():
        icon = "✅" if status == "loaded" else "⚠️" if "failed" in status else "❌"
        label = {"cuad": "Clause Classifier", "ner": "Legal NER", "nli": "NLI Contradiction"}[name]
        parts.append(f"{icon} {label}: {status}")
    return " · ".join(parts)

# Load models at startup
_load_cuad_model()
_load_ner_model()
_load_nli_model()

# ═══════════════════════════════════════════════════════════════════════
# 3. DOCUMENT PARSING
# ═══════════════════════════════════════════════════════════════════════

def parse_pdf(file_path):
    """Smart PDF parser: native text extraction with OCR fallback for scanned PDFs."""
    text, error, method = parse_pdf_smart(file_path)
    if text:
        if method == "ocr":
            print(f"[ClauseGuard] PDF extracted via OCR ({len(text)} chars)")
        return text, None
    if error:
        return None, error
    return None, "Could not extract text from PDF. Try uploading a clearer scan or digital PDF."

def parse_docx(file_path):
    if not _HAS_DOCX:
        return None, "DOCX parsing not available (python-docx not installed)"
    try:
        doc = DocxDocument(file_path)
        paragraphs = [p.text for p in doc.paragraphs if p.text.strip()]
        return "\n\n".join(paragraphs), None
    except Exception as e:
        return None, f"DOCX parse error: {e}"

def parse_document(file_path):
    if file_path is None:
        return None, "No file uploaded"
    ext = os.path.splitext(file_path)[1].lower()
    if ext == ".pdf":
        return parse_pdf(file_path)
    elif ext in (".docx", ".doc"):
        return parse_docx(file_path)
    elif ext in (".txt", ".md", ".rst"):
        try:
            with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
                return f.read(), None
        except Exception as e:
            return None, f"Text read error: {e}"
    else:
        return None, f"Unsupported file type: {ext}"

# ═══════════════════════════════════════════════════════════════════════
# 4. DETERMINISTIC CLAUSE SPLITTING (Fix 1 from bug report)
# ═══════════════════════════════════════════════════════════════════════

# Document-level chunk cache: same text always produces same chunks
_chunk_cache = {}

def split_clauses(text):
    """Deterministic, structure-aware clause splitting.
    Fix 1: Same input ALWAYS produces same output. Normalized text is hashed
    and cached so repeated runs on identical documents are identical."""
    # Normalize whitespace before hashing for determinism
    normalized = re.sub(r'\s+', ' ', text.strip())
    text_hash = hashlib.sha256(normalized.encode()).hexdigest()
    if text_hash in _chunk_cache:
        return _chunk_cache[text_hash]

    text = re.sub(r'\n{3,}', '\n\n', text.strip())

    # First try to detect numbered sections (1., 2., 3.1, (a), etc.)
    section_pattern = re.compile(
        r'(?:^|\n\n)'
        r'(?='
        r'\d+(?:\.\d+)*[.)]\s'   # 1. 2. 3.1. 3.1)
        r'|[A-Z]{2,}[A-Z\s]*\n'  # ALL CAPS HEADERS
        r'|\([a-z]\)\s'           # (a) (b) (c)
        r'|(?:Section|Article|Clause)\s+\d+'  # Section 1, Article 2
        r')',
        re.MULTILINE
    )

    positions = [m.start() for m in section_pattern.finditer(text)]

    if len(positions) >= 3:
        # Document has clear section structure — split on sections
        clauses = []
        for i, pos in enumerate(positions):
            end = positions[i + 1] if i + 1 < len(positions) else len(text)
            chunk = text[pos:end].strip()
            if len(chunk) > 30:
                # If a section is very long, split on paragraph breaks within it
                if len(chunk) > 1500:
                    sub_parts = chunk.split('\n\n')
                    current = ""
                    for sp in sub_parts:
                        if len(current) + len(sp) < 1200:
                            current += ("\n\n" + sp if current else sp)
                        else:
                            if len(current.strip()) > 30:
                                clauses.append(current.strip())
                            current = sp
                    if len(current.strip()) > 30:
                        clauses.append(current.strip())
                else:
                    clauses.append(chunk)
        # Also capture anything before the first section
        if positions and positions[0] > 50:
            preamble = text[:positions[0]].strip()
            if len(preamble) > 30:
                clauses.insert(0, preamble)
        result = clauses if clauses else _fallback_split(text)
        _chunk_cache[text_hash] = result
        return result
    else:
        result = _fallback_split(text)
        _chunk_cache[text_hash] = result
        return result

def _fallback_split(text):
    """Fallback: split on paragraph breaks and sentence boundaries."""
    # Try paragraph-based splitting first
    paragraphs = text.split('\n\n')
    if len(paragraphs) >= 3:
        clauses = []
        for p in paragraphs:
            p = p.strip()
            if len(p) > 30:
                if len(p) > 1500:
                    # Split long paragraphs on sentences
                    sents = re.split(r'(?<=[.!?])\s+(?=[A-Z])', p)
                    current = ""
                    for s in sents:
                        if len(current) + len(s) < 1000:
                            current += (" " + s if current else s)
                        else:
                            if len(current.strip()) > 30:
                                clauses.append(current.strip())
                            current = s
                    if len(current.strip()) > 30:
                        clauses.append(current.strip())
                else:
                    clauses.append(p)
        return clauses

    # Last resort: sentence splitting
    parts = re.split(r'(?<=[.!?])\s+(?=[A-Z0-9(])', text)
    return [p.strip() for p in parts if len(p.strip()) > 30]

# ═══════════════════════════════════════════════════════════════════════
# 5. CLAUSE DETECTION — FIXED: sigmoid + per-class thresholds + caching
#    Fix 3: Strip section headings before classification
#    Fix 6: Label guardrails for high-confidence false positives
# ═══════════════════════════════════════════════════════════════════════

# Fix 3: Section heading pattern — strip before classifying
_HEADING_RE = re.compile(r'^\d+(?:\.\d+)*\s+[A-Z][A-Z\s&,/]+$', re.MULTILINE)

def _strip_heading(text):
    """Remove leading section headings that confuse the classifier."""
    lines = text.split('\n')
    if lines and _HEADING_RE.match(lines[0].strip()):
        stripped = '\n'.join(lines[1:]).strip()
        return stripped if len(stripped) > 20 else text
    return text

# Fix 6: Label guardrails — keyword validation for high-confidence labels
_LABEL_GUARDRAILS = {
    "Liquidated Damages": re.compile(
        r'liquidated|pre-?determined.{0,10}damage|agreed.{0,10}sum|penalty clause|stipulated.{0,10}damage',
        re.IGNORECASE
    ),
    "Uncapped Liability": re.compile(
        r'uncapped|unlimited.{0,10}liabilit|no.{0,10}(limit|cap).{0,10}liabilit',
        re.IGNORECASE
    ),
}

def _apply_guardrails(label, text, confidence):
    """Fix 6: If label has a guardrail and text lacks required keywords, demote."""
    guard = _LABEL_GUARDRAILS.get(label)
    if guard and not guard.search(text):
        return "Other", confidence * 0.3  # demote to Other with reduced confidence
    return label, confidence

def _text_hash(text):
    return hashlib.md5(text.encode()).hexdigest()

_prediction_cache = {}
_CACHE_MAX = 2000

def classify_cuad(clause_text):
    if cuad_model is None or cuad_tokenizer is None:
        return _classify_regex(clause_text)

    # Fix 3: Strip section headings before classification
    clean_text = _strip_heading(clause_text)

    # Check cache
    h = _text_hash(clean_text[:512])
    if h in _prediction_cache:
        return _prediction_cache[h]

    try:
        inputs = cuad_tokenizer(
            clean_text,
            return_tensors="pt",
            truncation=True,
            max_length=256,
            padding=True
        )
        with torch.no_grad():
            logits = cuad_model(**inputs).logits

        # FIXED: Use sigmoid for multi-label (not softmax)
        probs = torch.sigmoid(logits)[0]

        results = []
        for i, prob in enumerate(probs):
            threshold = _CUAD_THRESHOLDS.get(i, 0.40)
            if float(prob) > threshold and i < len(CUAD_LABELS):
                label = CUAD_LABELS[i]
                conf = float(prob)
                # Fix 6: Apply guardrails — reject high-confidence false positives
                label, conf = _apply_guardrails(label, clause_text, conf)
                if label == "Other" and conf < 0.3:
                    continue  # Skip demoted labels
                risk = RISK_MAP.get(label, "LOW")
                results.append({
                    "label": label,
                    "confidence": round(conf, 3),
                    "risk": risk,
                    "description": DESC_MAP.get(label, label),
                    "source": "ml",
                })
        results.sort(key=lambda x: x["confidence"], reverse=True)

        # If no ML results, also try regex to catch what model misses
        if not results:
            results = _classify_regex(clause_text)

        # Cache result
        if len(_prediction_cache) < _CACHE_MAX:
            _prediction_cache[h] = results

        return results
    except Exception as e:
        print(f"[ClauseGuard] CUAD inference error: {e}")
        return _classify_regex(clause_text)

_REGEX_PATTERNS = {
    "Limitation of liability": [r"not liable", r"shall not be (liable|responsible)", r"in no event.*liable", r"limitation of liability", r"without warranty", r"disclaim"],
    "Unilateral termination": [r"terminat.*at any time", r"suspend.*account.*without", r"we may (terminat|suspend|discontinu)", r"right to (terminat|suspend)"],
    "Unilateral change": [r"sole discretion", r"reserves? the right to (modify|change|update|amend)", r"at any time.*without (prior )?notice", r"we may (modify|change|update)"],
    "Content removal": [r"remove.*content.*without", r"right to remove", r"we may.*remove"],
    "Contract by using": [r"by (using|accessing).*you agree", r"continued use.*constitutes? acceptance"],
    "Choice of law": [r"governed by.*laws? of", r"shall be governed", r"laws of the state of"],
    "Jurisdiction": [r"exclusive jurisdiction", r"courts? of.*(california|delaware|new york|ireland|england)", r"submit to.*jurisdiction"],
    "Arbitration": [r"arbitrat", r"binding arbitration", r"waive.*right.*court", r"class action waiver"],
    "Governing Law": [r"governed by", r"laws of", r"jurisdiction of"],
    "Termination for Convenience": [r"terminat.*for convenience", r"terminat.*without cause", r"terminat.*at any time"],
    "Non-Compete": [r"non-compete", r"shall not compete", r"competition"],
    "Exclusivity": [r"exclusive(?:ly)?(?:\s+(?:deal|relationship|partner|right))", r"exclusivity"],
    "IP Ownership Assignment": [r"assign.*intellectual property", r"ownership of.*ip", r"all rights.*assign"],
    "Uncapped Liability": [r"unlimited liability", r"uncapped", r"no.*limit.*liability"],
    "Cap on Liability": [r"cap on liability", r"maximum liability", r"liability.*shall not exceed", r"aggregate liability.*not exceed"],
    "Indemnification": [r"indemnif", r"hold harmless", r"defend.*against.*claim"],
    "Confidentiality": [r"confidential(?:ity)?", r"non-disclosure", r"\bnda\b"],
    "Force Majeure": [r"force majeure", r"act of god", r"beyond.*(?:reasonable\s+)?control"],
    "Penalties": [r"penalt(?:y|ies)", r"late fee", r"default charge", r"interest on overdue"],
}

def _classify_regex(text):
    """Regex fallback — returns pattern match, NOT fake confidence."""
    text_lower = text.lower()
    results = []
    seen = set()
    for label, patterns in _REGEX_PATTERNS.items():
        for pat in patterns:
            if re.search(pat, text_lower):
                if label not in seen:
                    risk = RISK_MAP.get(label, "MEDIUM")
                    results.append({
                        "label": label,
                        "confidence": None,  # FIXED: no fake confidence for regex
                        "risk": risk,
                        "description": DESC_MAP.get(label, label),
                        "source": "pattern",
                    })
                    seen.add(label)
                break
    return results

# ═══════════════════════════════════════════════════════════════════════
# 6. LEGAL NER — ML model with regex fallback
# ═══════════════════════════════════════════════════════════════════════

def extract_entities(text):
    """Extract entities using ML model (matterstack/legal-bert-ner) with regex fallback."""
    entities = []

    # Try ML NER first
    if _HAS_NER_MODEL and ner_pipeline is not None:
        try:
            # Process in chunks (model has max length limits)
            chunks = [text[i:i+512] for i in range(0, min(len(text), 10000), 450)]
            offset = 0
            for chunk in chunks:
                ner_results = ner_pipeline(chunk)
                for ent in ner_results:
                    if ent.get("score", 0) > 0.5:
                        entities.append({
                            "text": ent["word"],
                            "type": _map_ner_label(ent.get("entity_group", ent.get("entity", "MISC"))),
                            "start": ent["start"] + offset,
                            "end": ent["end"] + offset,
                            "score": round(ent["score"], 3),
                            "source": "ml",
                        })
                offset += 450
        except Exception as e:
            print(f"[ClauseGuard] ML NER error, falling back to regex: {e}")
            entities = _extract_entities_regex(text)
    else:
        entities = _extract_entities_regex(text)

    # Always supplement with regex patterns for things NER often misses
    regex_ents = _extract_entities_regex(text)
    # Merge: add regex entities that don't overlap with ML entities
    ml_spans = set()
    for e in entities:
        for pos in range(e["start"], e["end"]):
            ml_spans.add(pos)
    for re_ent in regex_ents:
        if not any(pos in ml_spans for pos in range(re_ent["start"], re_ent["end"])):
            entities.append(re_ent)

    # Deduplicate and sort
    entities.sort(key=lambda x: (x["start"], -(x["end"] - x["start"])))
    filtered = []
    last_end = -1
    for e in entities:
        if e["start"] >= last_end:
            filtered.append(e)
            last_end = e["end"]
    return filtered

def _map_ner_label(label):
    """Map NER model labels to our entity types."""
    label = label.upper()
    mapping = {
        "PER": "PERSON",
        "PERSON": "PERSON",
        "ORG": "PARTY",
        "ORGANIZATION": "PARTY",
        "LOC": "JURISDICTION",
        "LOCATION": "JURISDICTION",
        "GPE": "JURISDICTION",
        "DATE": "DATE",
        "MONEY": "MONEY",
        "MISC": "MISC",
        "LAW": "LEGAL_REF",
    }
    return mapping.get(label, label)

def _extract_entities_regex(text):
    """Regex-based NER fallback."""
    entities = []
    patterns = [
        # Dates
        (r'\b(?:January|February|March|April|May|June|July|August|September|October|November|December)\s+\d{1,2},?\s+\d{4}\b', "DATE"),
        (r'\b\d{1,2}/\d{1,2}/\d{2,4}\b', "DATE"),
        (r'\b\d{1,2}-(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)-\d{2,4}\b', "DATE"),
        (r'\b(?:Effective|Commencement|Expiration|Termination)\s+Date\b', "DATE_REF"),
        # Money
        (r'\$\s?\d{1,3}(?:,\d{3})*(?:\.\d{2})?(?:\s*(?:million|billion|thousand|M|B|K))?', "MONEY"),
        (r'\b\d{1,3}(?:,\d{3})*(?:\.\d{2})?\s*(?:USD|EUR|GBP|dollars|euros|pounds)', "MONEY"),
        (r'\b(?:USD|EUR|GBP)\s*\d{1,3}(?:,\d{3})*(?:\.\d{2})?', "MONEY"),
        # Percentages
        (r'\b\d+(?:\.\d+)?%', "PERCENTAGE"),
        # Durations
        (r'\b\d+\s*(?:year|month|week|day|business day)s?\b', "DURATION"),
        # Parties (require suffix to reduce false positives)
        (r'\b[A-Z][A-Za-z0-9\s&,]+?(?:Inc\.?|LLC|Ltd\.?|Limited|Corp\.?|Corporation|PLC|GmbH|AG|S\.A\.?|B\.V\.?|L\.P\.?|LLP)\b', "PARTY"),
        (r'\b(?:Party A|Party B|Disclosing Party|Receiving Party|Licensor|Licensee|Buyer|Seller|Tenant|Landlord|Employer|Employee|Customer|Vendor|Client)\b', "PARTY_ROLE"),
        # Jurisdictions
        (r'\b(?:State|Commonwealth)\s+of\s+[A-Z][a-zA-Z\s]+', "JURISDICTION"),
        (r'\b(?:California|Delaware|New York|Texas|Florida|England|Ireland|Germany|France|Singapore|Hong Kong|Ontario|British Columbia)\b', "JURISDICTION"),
        # Defined Terms (quoted or parenthesized)
        (r'"([A-Z][A-Za-z\s]{1,40})"', "DEFINED_TERM"),
        (r'\((?:the\s+)?"([A-Z][A-Za-z\s]{1,40})"\)', "DEFINED_TERM"),
    ]
    for pat, etype in patterns:
        for m in re.finditer(pat, text, re.IGNORECASE if etype in ("DATE", "MONEY", "DURATION", "PERCENTAGE") else 0):
            txt = m.group(1) if m.lastindex else m.group()
            entities.append({
                "text": txt,
                "type": etype,
                "start": m.start(),
                "end": m.end(),
                "source": "pattern",
            })
    return entities

# ═══════════════════════════════════════════════════════════════════════
# 7. NLI / CONTRADICTION DETECTION — Real semantic analysis
# ═══════════════════════════════════════════════════════════════════════

def detect_contradictions(clause_results, raw_text=""):
    """
    Detect contradictions using:
    1. NLI cross-encoder model (semantic contradiction detection)
    2. Structural conflict detection (mutually exclusive labels)
    3. Missing critical clause detection
    """
    contradictions = []
    labels_found = set()
    clause_texts_by_label = defaultdict(list)

    for cr in clause_results:
        labels_found.add(cr["label"])
        clause_texts_by_label[cr["label"]].append(cr.get("text", ""))

    # ── 1. Semantic NLI (if model available) ──
    if _HAS_NLI_MODEL and nli_pipeline is not None:
        # Check clauses that belong to potentially conflicting categories
        conflict_pairs = [
            ("Uncapped Liability", "Cap on Liability",
             "Liability cannot be both uncapped and capped simultaneously."),
            ("IP Ownership Assignment", "Joint IP Ownership",
             "IP cannot be both fully assigned and jointly owned."),
            ("Exclusivity", "Non-Transferable License",
             "Exclusivity and non-transferable license may conflict."),
        ]
        for label_a, label_b, explanation in conflict_pairs:
            if label_a in labels_found and label_b in labels_found:
                texts_a = clause_texts_by_label[label_a]
                texts_b = clause_texts_by_label[label_b]
                for ta in texts_a[:2]:
                    for tb in texts_b[:2]:
                        try:
                            nli_result = nli_pipeline(
                                f"{ta[:256]} [SEP] {tb[:256]}",
                                truncation=True
                            )
                            # Check if model predicts contradiction
                            for r in (nli_result if isinstance(nli_result, list) else [nli_result]):
                                if r.get("label", "").lower() == "contradiction" and r.get("score", 0) > 0.6:
                                    contradictions.append({
                                        "type": "CONTRADICTION",
                                        "explanation": explanation,
                                        "severity": "HIGH",
                                        "clauses": [label_a, label_b],
                                        "confidence": round(r["score"], 3),
                                        "source": "nli_model",
                                    })
                        except Exception:
                            pass

        # Also check for internal contradictions within governing law / termination
        for label in ["Governing Law", "Termination for Convenience"]:
            texts = clause_texts_by_label.get(label, [])
            if len(texts) >= 2:
                for i in range(len(texts)):
                    for j in range(i + 1, min(len(texts), i + 3)):
                        try:
                            nli_result = nli_pipeline(
                                f"{texts[i][:256]} [SEP] {texts[j][:256]}",
                                truncation=True
                            )
                            for r in (nli_result if isinstance(nli_result, list) else [nli_result]):
                                if r.get("label", "").lower() == "contradiction" and r.get("score", 0) > 0.6:
                                    contradictions.append({
                                        "type": "CONTRADICTION",
                                        "explanation": f"Conflicting {label} provisions detected — clauses contradict each other.",
                                        "severity": "HIGH",
                                        "clauses": [label],
                                        "confidence": round(r["score"], 3),
                                        "source": "nli_model",
                                    })
                        except Exception:
                            pass
    else:
        # ── Heuristic fallback (improved) ──
        _heuristic_pairs = [
            (["Uncapped Liability"], ["Cap on Liability"],
             "Liability cannot be both uncapped and capped simultaneously."),
            (["IP Ownership Assignment"], ["Joint IP Ownership"],
             "IP cannot be both fully assigned and jointly owned."),
        ]
        for group_a, group_b, explanation in _heuristic_pairs:
            found_a = any(l in labels_found for l in group_a)
            found_b = any(l in labels_found for l in group_b)
            if found_a and found_b:
                contradictions.append({
                    "type": "CONTRADICTION",
                    "explanation": explanation,
                    "severity": "HIGH",
                    "clauses": group_a + group_b,
                    "source": "heuristic",
                })

    # ── 2. Missing critical clauses (Fix 4: check raw_text, not labels) ──
    _REQUIRED_CLAUSE_PATTERNS = {
        "Governing Law": re.compile(
            r'govern(?:ed|ing).{0,15}law|applicable.{0,10}law|laws?\s+of\s+the\s+state',
            re.IGNORECASE
        ),
        "Limitation of liability": re.compile(
            r'limitation.{0,10}liabilit|cap.{0,10}liabilit|liabilit.{0,10}shall\s+not\s+exceed|in\s+no\s+event.{0,20}liable',
            re.IGNORECASE
        ),
        "Arbitration": re.compile(
            r'arbitrat|AAA|JAMS|binding.{0,10}dispute',
            re.IGNORECASE
        ),
        "Termination": re.compile(
            r'terminat(?:e|ion|ed)|cancel(?:lation)?',
            re.IGNORECASE
        ),
    }
    for clause_name, pattern in _REQUIRED_CLAUSE_PATTERNS.items():
        # Check raw_text directly — it's stable and deterministic
        if not pattern.search(raw_text):
            contradictions.append({
                "type": "MISSING",
                "explanation": f"No '{clause_name}' clause detected in the document.",
                "severity": "MEDIUM",
                "clauses": [clause_name],
                "source": "structural",
            })

    # Deduplicate
    seen = set()
    unique = []
    for c in contradictions:
        key = (c["type"], c["explanation"])
        if key not in seen:
            seen.add(key)
            unique.append(c)

    return unique

# ═══════════════════════════════════════════════════════════════════════
# 8. RISK SCORING
# ═══════════════════════════════════════════════════════════════════════

def compute_risk_score(clause_results, total_clauses):
    sev_counts = {"CRITICAL": 0, "HIGH": 0, "MEDIUM": 0, "LOW": 0}
    for cr in clause_results:
        sev = cr.get("risk", "LOW")
        sev_counts[sev] += 1
    if total_clauses == 0:
        return 0, "A", sev_counts
    weighted = sum(sev_counts[s] * RISK_WEIGHTS[s] for s in sev_counts)
    risk = min(100, round(weighted / max(1, total_clauses) * 10))
    if risk >= 70: grade = "F"
    elif risk >= 50: grade = "D"
    elif risk >= 30: grade = "C"
    elif risk >= 15: grade = "B"
    else: grade = "A"
    return risk, grade, sev_counts

# ═══════════════════════════════════════════════════════════════════════
# 9. MAIN ANALYSIS PIPELINE
# ═══════════════════════════════════════════════════════════════════════

def analyze_contract(text):
    if not text or len(text.strip()) < 50:
        return None, "Document too short (minimum 50 characters)"
    clauses = split_clauses(text)
    if not clauses:
        return None, "No clauses detected in document"
    clause_results = []
    for clause in clauses:
        predictions = classify_cuad(clause)
        if predictions:
            for pred in predictions:
                clause_results.append({
                    "text": clause,
                    "label": pred["label"],
                    "confidence": pred["confidence"],
                    "risk": pred["risk"],
                    "description": pred["description"],
                    "source": pred.get("source", "unknown"),
                })
    entities = extract_entities(text)
    contradictions = detect_contradictions(clause_results, text)
    risk, grade, sev_counts = compute_risk_score(clause_results, len(clauses))
    obligations = extract_obligations(text)
    # Fix 5: Compliance runs against full raw_text (already done in compliance.py)
    compliance = check_compliance(text)

    # Fix 2: Compute flagged_clauses AFTER all processing is complete
    flagged_clause_count = len(clause_results)
    unique_flagged_texts = len(set(cr["text"] for cr in clause_results))

    result = {
        "metadata": {
            "analysis_date": datetime.now().isoformat(),
            "total_clauses": len(clauses),
            "flagged_clauses": flagged_clause_count,
            "unique_flagged": unique_flagged_texts,
            "model": get_model_status_text(),
            "text_hash": hashlib.sha256(re.sub(r'\s+', ' ', text.strip()).encode()).hexdigest()[:16],
        },
        "risk": {
            "score": risk,
            "grade": grade,
            "breakdown": sev_counts,
        },
        "clauses": clause_results,
        "entities": entities,
        "contradictions": contradictions,
        "obligations": obligations,
        "compliance": compliance,
        "raw_text": text,
    }
    return result, None

# ═══════════════════════════════════════════════════════════════════════
# 10. EXPORT FUNCTIONS — FIXED: per-session temp files
# ═══════════════════════════════════════════════════════════════════════

def export_json(result):
    if result is None:
        return None
    return json.dumps(result, indent=2, default=str)

def export_csv(result):
    if result is None:
        return None
    output = io.StringIO()
    writer = csv.writer(output)
    writer.writerow(["Clause Text", "Label", "Risk", "Confidence", "Description", "Source"])
    for cr in result.get("clauses", []):
        conf = cr.get("confidence")
        conf_str = f"{conf:.3f}" if conf is not None else "pattern match"
        writer.writerow([
            cr.get("text", "")[:500],
            cr.get("label", ""),
            cr.get("risk", ""),
            conf_str,
            cr.get("description", ""),
            cr.get("source", ""),
        ])
    return output.getvalue()

# ═══════════════════════════════════════════════════════════════════════
# 11. UI RENDERING — FIXED: shows confidence source properly
# ═══════════════════════════════════════════════════════════════════════

def render_summary(result):
    if result is None:
        return ""
    risk = result["risk"]
    score = risk["score"]
    grade = risk["grade"]
    breakdown = risk["breakdown"]
    grade_color = {
        "A": "#16a34a", "B": "#65a30d", "C": "#ca8a04",
        "D": "#ea580c", "F": "#dc2626",
    }.get(grade, "#6b7280")
    crit, high, med, low = breakdown["CRITICAL"], breakdown["HIGH"], breakdown["MEDIUM"], breakdown["LOW"]
    html = f"""
    <div style="font-family:system-ui,sans-serif;padding:16px;border:1px solid #e5e7eb;border-radius:12px;background:#fff;">
      <div style="text-align:center;margin-bottom:16px;">
        <div style="font-size:48px;font-weight:700;color:{grade_color};">{score}</div>
        <div style="font-size:14px;color:#6b7280;">/100 Risk Score</div>
        <div style="display:inline-block;margin-top:8px;padding:4px 16px;border-radius:20px;background:{grade_color};color:white;font-weight:600;font-size:14px;">
          Grade {grade}
        </div>
      </div>
      <div style="display:grid;grid-template-columns:1fr 1fr;gap:8px;margin-bottom:12px;">
        <div style="padding:8px;border-radius:6px;background:#fef2f2;text-align:center;">
          <div style="font-size:20px;font-weight:700;color:#dc2626;">{crit}</div>
          <div style="font-size:11px;color:#991b1b;">Critical</div>
        </div>
        <div style="padding:8px;border-radius:6px;background:#fff7ed;text-align:center;">
          <div style="font-size:20px;font-weight:700;color:#ea580c;">{high}</div>
          <div style="font-size:11px;color:#9a3412;">High</div>
        </div>
        <div style="padding:8px;border-radius:6px;background:#fefce8;text-align:center;">
          <div style="font-size:20px;font-weight:700;color:#ca8a04;">{med}</div>
          <div style="font-size:11px;color:#854d0e;">Medium</div>
        </div>
        <div style="padding:8px;border-radius:6px;background:#f0fdf4;text-align:center;">
          <div style="font-size:20px;font-weight:700;color:#16a34a;">{low}</div>
          <div style="font-size:11px;color:#166534;">Low</div>
        </div>
      </div>
      <div style="font-size:12px;color:#6b7280;text-align:center;">
        {result['metadata']['total_clauses']} clauses analyzed · {result['metadata']['flagged_clauses']} flagged
        <br><span style="font-size:10px;">{result['metadata']['model']}</span>
      </div>
    </div>
    """
    return html

def render_clause_cards(result):
    if result is None:
        return ""
    clauses = result.get("clauses", [])
    if not clauses:
        return '<div style="padding:24px;text-align:center;color:#6b7280;">No clauses detected.</div>'
    grouped = defaultdict(list)
    for cr in clauses:
        grouped[cr["text"]].append(cr)
    html = '<div style="font-family:system-ui,sans-serif;">'
    for text, items in grouped.items():
        max_risk = max(items, key=lambda x: {"CRITICAL":4,"HIGH":3,"MEDIUM":2,"LOW":1}[x["risk"]])["risk"]
        border, bg, icon = RISK_STYLES[max_risk]
        tags = ""
        for item in items:
            tag_bg = RISK_STYLES[item["risk"]][1]
            tag_color = RISK_STYLES[item["risk"]][0]
            conf = item.get("confidence")
            source = item.get("source", "")
            if conf is not None:
                conf_text = f"{conf:.0%}"
            else:
                conf_text = "pattern"
            source_icon = "🤖" if source == "ml" else "📝"
            tags += f'<span style="background:{tag_bg};color:{tag_color};border:1px solid {tag_color}33;padding:2px 8px;border-radius:12px;font-size:11px;font-weight:500;margin-right:4px;">{source_icon} {item["label"]} ({conf_text})</span>'
        descs = "".join(
            f'<p style="font-size:12px;color:#6b7280;margin:4px 0 0 0;">{item["description"]}</p>'
            for item in items
        )
        preview = text[:300] + ("..." if len(text) > 300 else "")
        preview = preview.replace("<", "&lt;").replace(">", "&gt;")
        html += f"""
        <div style="border:1px solid #e5e7eb;border-left:4px solid {border};border-radius:8px;padding:14px;margin-bottom:10px;background:#fafafa;">
          <div style="display:flex;align-items:center;gap:6px;margin-bottom:6px;">
            <span style="font-size:16px;">{icon}</span>
            <span style="font-size:12px;font-weight:600;color:{border};text-transform:uppercase;">{max_risk}</span>
          </div>
          <p style="font-size:13px;color:#374151;line-height:1.6;margin:0 0 8px 0;">{preview}</p>
          <div style="margin-bottom:6px;">{tags}</div>
          {descs}
        </div>
        """
    html += "</div>"
    return html

def render_entities(result):
    if result is None:
        return ""
    entities = result.get("entities", [])
    if not entities:
        return '<div style="padding:16px;color:#6b7280;">No entities detected.</div>'
    grouped = defaultdict(list)
    for e in entities:
        grouped[e["type"]].append(e["text"])
    html = '<div style="font-family:system-ui,sans-serif;">'
    for etype, texts in grouped.items():
        unique = list(dict.fromkeys(texts))[:20]
        color = {
            "DATE": "#3b82f6", "DATE_REF": "#60a5fa",
            "MONEY": "#22c55e", "PERCENTAGE": "#10b981",
            "DURATION": "#6366f1",
            "PARTY": "#8b5cf6", "PARTY_ROLE": "#a78bfa",
            "PERSON": "#ec4899",
            "JURISDICTION": "#f59e0b",
            "DEFINED_TERM": "#ec4899",
            "LEGAL_REF": "#6b7280",
            "MISC": "#9ca3af",
        }.get(etype, "#6b7280")
        items_html = "".join(
            f'<span style="display:inline-block;background:{color}15;color:{color};border:1px solid {color}40;padding:3px 10px;border-radius:6px;font-size:12px;margin:3px;">{t}</span>'
            for t in unique
        )
        html += f"""
        <div style="margin-bottom:12px;">
          <div style="font-size:12px;font-weight:600;color:#374151;margin-bottom:6px;text-transform:uppercase;">{etype}</div>
          <div>{items_html}</div>
        </div>
        """
    html += "</div>"
    return html

def render_contradictions(result):
    if result is None:
        return ""
    contradictions = result.get("contradictions", [])
    if not contradictions:
        return '<div style="padding:16px;color:#16a34a;">✓ No contradictions or missing clauses detected.</div>'
    html = '<div style="font-family:system-ui,sans-serif;">'
    for c in contradictions:
        sev_color = RISK_STYLES[c["severity"]][0]
        icon = "⚠️" if c["type"] == "CONTRADICTION" else "📋"
        source = c.get("source", "")
        source_badge = ""
        if source == "nli_model":
            conf = c.get("confidence", 0)
            source_badge = f'<span style="font-size:10px;background:#eff6ff;color:#3b82f6;padding:1px 6px;border-radius:4px;margin-left:8px;">🤖 NLI {conf:.0%}</span>'
        elif source == "heuristic":
            source_badge = '<span style="font-size:10px;background:#fef3c7;color:#92400e;padding:1px 6px;border-radius:4px;margin-left:8px;">📝 Heuristic</span>'
        html += f"""
        <div style="border:1px solid #e5e7eb;border-left:4px solid {sev_color};border-radius:8px;padding:12px;margin-bottom:8px;background:#fafafa;">
          <div style="display:flex;align-items:center;gap:6px;margin-bottom:4px;">
            <span>{icon}</span>
            <span style="font-size:12px;font-weight:600;color:{sev_color};">{c["type"]}</span>
            {source_badge}
          </div>
          <p style="font-size:13px;color:#374151;margin:0;">{c["explanation"]}</p>
        </div>
        """
    html += "</div>"
    return html

def render_document_viewer(result):
    if result is None:
        return ""
    text = result.get("raw_text", "")
    entities = sorted(result.get("entities", []), key=lambda x: x["start"])
    html_parts = []
    last_end = 0
    for e in entities:
        if e["start"] >= last_end:
            html_parts.append(text[last_end:e["start"]].replace("<", "&lt;").replace(">", "&gt;"))
            color = {
                "DATE": "#bfdbfe", "DATE_REF": "#bfdbfe",
                "MONEY": "#bbf7d0", "PERCENTAGE": "#a7f3d0",
                "DURATION": "#c7d2fe",
                "PARTY": "#ddd6fe", "PARTY_ROLE": "#ddd6fe",
                "PERSON": "#fbcfe8",
                "JURISDICTION": "#fde68a",
                "DEFINED_TERM": "#fbcfe8",
                "LEGAL_REF": "#e5e7eb",
            }.get(e["type"], "#e5e7eb")
            label = e["type"].replace("_", " ")
            html_parts.append(
                f'<mark style="background:{color};padding:1px 2px;border-radius:2px;font-size:12px;" title="{label}">{e["text"].replace("<","&lt;").replace(">","&gt;")}</mark>'
            )
            last_end = e["end"]
    html_parts.append(text[last_end:].replace("<", "&lt;").replace(">", "&gt;"))
    highlighted = "".join(html_parts)
    return f"""
    <div style="font-family:monospace;font-size:13px;line-height:1.6;padding:16px;border:1px solid #e5e7eb;border-radius:8px;background:#fff;max-height:600px;overflow-y:auto;white-space:pre-wrap;">
      {highlighted}
    </div>
    """

# ═══════════════════════════════════════════════════════════════════════
# 12. COMPARISON UI FUNCTIONS
# ═══════════════════════════════════════════════════════════════════════

def run_comparison(text_a, text_b):
    if not text_a or len(text_a.strip()) < 50:
        return "Contract A is too short", ""
    if not text_b or len(text_b.strip()) < 50:
        return "Contract B is too short", ""
    result = compare_contracts(text_a, text_b)
    return render_comparison_html(result), json.dumps(result, indent=2)

# ═══════════════════════════════════════════════════════════════════════
# 13. GRADIO UI
# ═══════════════════════════════════════════════════════════════════════

def process_upload(file):
    if file is None:
        return "", "No file uploaded"
    text, error = parse_document(file)
    if error:
        return "", error
    return text, "Document loaded successfully"

def run_analysis(text):
    if not text or len(text.strip()) < 50:
        err_html = '<p style="color:#dc2626;padding:16px;">Document too short (minimum 50 characters)</p>'
        return [err_html] * 8 + [None, None, "", None]
    result, error = analyze_contract(text)
    if error:
        err_html = f'<p style="color:#dc2626;padding:16px;">{error}</p>'
        return [err_html] * 8 + [None, None, error, None]

    # FIXED: per-session temp files
    session_id = uuid.uuid4().hex[:8]
    json_path = os.path.join(tempfile.gettempdir(), f"clauseguard_{session_id}.json")
    csv_path = os.path.join(tempfile.gettempdir(), f"clauseguard_{session_id}.csv")

    with open(json_path, "w") as f:
        json.dump(result, f, indent=2, default=str)
    csv_content = export_csv(result)
    with open(csv_path, "w") as f:
        f.write(csv_content)

    # Generate redline suggestions (Tier 1 template + Tier 3 LLM for critical/high)
    redlines = generate_redlines(result, use_llm=True)
    redlines_html = render_redlines_html(redlines)

    return [
        render_summary(result),
        render_clause_cards(result),
        render_entities(result),
        render_contradictions(result),
        render_document_viewer(result),
        render_obligations_html(result.get("obligations", [])),
        render_compliance_html(result.get("compliance", {})),
        redlines_html,
        json_path,
        csv_path,
        "Analysis complete",
        result,  # Store analysis result for chatbot
    ]

def do_clear():
    return [""] * 8 + [None, None, "", None]

# ── Example contracts ──
SPOTIFY_TOS = """By using the Spotify Service, you agree to be bound by these Terms of Use.

Spotify may, in its sole discretion, modify or update these Terms of Service at any time without prior notice. Your continued use of the Service after any such changes constitutes your acceptance of the new Terms of Service.

In no event will Spotify be liable for any indirect, incidental, special, consequential, or punitive damages, or any loss of profits or revenues, whether incurred directly or indirectly.

Spotify reserves the right to remove or disable access to any User Content for any reason, without prior notice.

Spotify may terminate your account or suspend your access at any time, with or without cause, with or without notice, effective immediately.

These Terms will be governed by and construed in accordance with the laws of the State of New York.

Any dispute shall be finally settled by arbitration in New York County. The parties waive any right to a jury trial."""

RENTAL_AGREEMENT = """The Landlord reserves the right to enter the premises at any time without prior notice for inspection or any other purpose deemed necessary in their sole discretion.

The Landlord shall not be liable for any damage to the Tenant's personal property, whether caused by water leaks, fire, theft, or any other cause, including the Landlord's own negligence.

The Landlord may terminate this lease at any time with only 7 days written notice, for any reason or no reason at all.

Any disputes arising from this lease agreement shall be resolved exclusively in the courts of the State of California, and the Tenant waives the right to a jury trial.

The Landlord reserves the right to modify the terms of this lease at any time. Continued occupancy constitutes acceptance of the new terms."""

NDA_SAMPLE = """NON-DISCLOSURE AGREEMENT

This Non-Disclosure Agreement (the "Agreement") is entered into as of January 15, 2024 (the "Effective Date") by and between Acme Technologies, Inc. ("Disclosing Party") and Beta Solutions LLC ("Receiving Party").

1. Governing Law. This Agreement shall be governed by and construed in accordance with the laws of the State of Delaware, without regard to its conflict of law principles.

2. Term. This Agreement shall remain in effect for a period of three (3) years from the Effective Date.

3. Termination. Either party may terminate this Agreement for convenience upon thirty (30) days prior written notice.

4. Intellectual Property. All Confidential Information disclosed hereunder shall remain the exclusive property of the Disclosing Party. The Receiving Party hereby assigns to the Disclosing Party all right, title, and interest in any derivative works.

5. Limitation of Liability. IN NO EVENT SHALL EITHER PARTY BE LIABLE FOR ANY INDIRECT, INCIDENTAL, SPECIAL, CONSEQUENTIAL, OR PUNITIVE DAMAGES.

6. Indemnification. The Receiving Party shall indemnify and hold harmless the Disclosing Party from any and all claims arising from a breach of this Agreement.

7. Non-Compete. During the term of this Agreement and for a period of two (2) years thereafter, the Receiving Party shall not engage in any business that competes with the Disclosing Party."""

COMPLEX_CONTRACT = """MASTER SERVICE AGREEMENT

This Master Service Agreement ("MSA") is entered into as of March 1, 2024 (the "Effective Date") by and between CloudTech Solutions, Inc., a Delaware corporation ("Provider") and Global Retail Partners LLC, a New York limited liability company ("Customer").

1. SERVICES. Provider shall provide cloud hosting and data processing services as described in Exhibit A. Provider shall comply with all applicable laws including GDPR and CCPA.

2. TERM AND RENEWAL. The initial term is twelve (12) months, automatically renewing for successive one (1) year periods unless terminated in accordance with Section 7.

3. FEES AND PAYMENT. Customer shall pay a monthly fee of $25,000 within 30 days of invoice. Late payments incur a penalty of 1.5% per month. The total contract value is $300,000.

4. LIABILITY. Provider's aggregate liability shall not exceed $1,000,000. IN NO EVENT SHALL PROVIDER BE LIABLE FOR LOST PROFITS OR CONSEQUENTIAL DAMAGES. Customer assumes all risk of data loss.

5. INDEMNIFICATION. Each party shall indemnify the other for third-party claims arising from breach of this Agreement. Customer shall indemnify Provider for claims arising from Customer Data.

6. INTELLECTUAL PROPERTY. Provider retains all IP rights. Customer receives a non-transferable, non-exclusive license for the term. Upon termination, Customer shall return or destroy all Provider materials within 10 business days.

7. TERMINATION. Either party may terminate for convenience with 90 days notice. Provider may terminate immediately for non-payment. Upon termination, Customer shall pay all outstanding fees.

8. GOVERNING LAW. This Agreement is governed by the laws of the State of Delaware. Disputes shall be resolved by binding arbitration in Wilmington, Delaware.

9. FORCE MAJEURE. Neither party shall be liable for delays due to acts of God, war, terrorism, or government action.

10. AUDIT RIGHTS. Customer may audit Provider's compliance annually. Provider shall provide SOC 2 Type II reports within 30 days of request.

11. INSURANCE. Provider shall maintain general liability insurance of at least $5,000,000 and cyber liability insurance of at least $2,000,000.

12. CONFIDENTIALITY. Both parties agree to keep Confidential Information secure for five (5) years. This obligation survives termination.

13. ASSIGNMENT. Neither party may assign this Agreement without prior written consent. Any attempted assignment is void.

14. THIRD PARTY BENEFICIARY. No third party shall have rights under this Agreement except as expressly provided."""

with gr.Blocks(
    title="ClauseGuard — AI Contract Analysis",
    css="""
        .gradio-container { max-width: 1600px !important; }
    """
) as demo:

    # ── Shared State (for chatbot RAG) ──────────────────────────────
    analysis_state = gr.State(None)       # Full analysis result dict
    chunks_state = gr.State([])           # Contract text chunks for RAG
    embeddings_state = gr.State(None)     # Chunk embeddings (numpy array)

    gr.HTML("""
    <div style="display:flex;align-items:center;justify-content:space-between;padding:12px 0;border-bottom:2px solid #e5e7eb;margin-bottom:16px;">
      <div>
        <h1 style="font-size:24px;font-weight:700;margin:0;color:#1f2937;">🛡️ ClauseGuard</h1>
        <p style="font-size:13px;color:#6b7280;margin:4px 0 0 0;">AI-Powered Legal Contract Analysis · 41 Clause Categories · Risk Scoring · ML NER · NLI Contradictions · Compliance · Obligations · <strong>Q&A Chatbot</strong> · <strong>Clause Redlining</strong> · <strong>OCR</strong></p>
      </div>
      <div style="font-size:12px;color:#9ca3af;">v4.0 · Precision Legal AI</div>
    </div>
    """)

    # ── Main Tabs: Analysis vs Comparison vs Chatbot ──
    with gr.Tabs():

        # ═══════ TAB 1: Single Contract Analysis ═══════
        with gr.Tab("📄 Single Contract Analysis"):
            with gr.Row():
                with gr.Column(scale=1):
                    file_input = gr.File(
                        label="📁 Upload Contract (PDF/DOCX/TXT)",
                        file_types=[".pdf", ".docx", ".doc", ".txt", ".md"],
                    )
                    load_btn = gr.Button("Load Document", variant="secondary", size="sm")
                    load_status = gr.Textbox(label="Status", interactive=False, lines=1)

                with gr.Column(scale=3):
                    text_input = gr.Textbox(
                        label="📄 Contract Text",
                        placeholder="Paste contract text here, or upload a file above...\n\n💡 Scanned PDFs are automatically processed with OCR.",
                        lines=14,
                        max_lines=40,
                        show_copy_button=True,
                    )

                with gr.Column(scale=1):
                    scan_btn = gr.Button("🔍 Analyze Contract", variant="primary", size="lg")
                    clear_btn = gr.Button("Clear", variant="secondary", size="sm")
                    status_msg = gr.Textbox(label="Analysis Status", interactive=False, lines=1)

            # ── Examples ──
            with gr.Row():
                gr.Examples(
                    examples=[[SPOTIFY_TOS], [RENTAL_AGREEMENT], [NDA_SAMPLE], [COMPLEX_CONTRACT]],
                    inputs=[text_input],
                    label="Example Contracts",
                )

            # ── Results ──
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### 📊 Risk Summary")
                    summary_html = gr.HTML()

                    gr.Markdown("### 📥 Export Reports")
                    json_file = gr.File(label="JSON Report")
                    csv_file = gr.File(label="CSV Report")

                with gr.Column(scale=3):
                    with gr.Tabs():
                        with gr.Tab("📄 Document"):
                            doc_html = gr.HTML(label="Document Viewer")
                        with gr.Tab("⚠️ Clauses (41 Categories)"):
                            clauses_html = gr.HTML(label="Detected Clauses")
                        with gr.Tab("🏷️ Entities"):
                            entities_html = gr.HTML(label="Named Entities")
                        with gr.Tab("🔍 Contradictions"):
                            nli_html = gr.HTML(label="Contradictions & Missing Clauses")
                        with gr.Tab("📋 Obligations"):
                            obligations_html = gr.HTML(label="Obligation Tracker")
                        with gr.Tab("⚖️ Compliance"):
                            compliance_html = gr.HTML(label="Compliance Checker")
                        with gr.Tab("✏️ Redlining"):
                            redlining_html = gr.HTML(label="Clause Redlining Suggestions")

        # ═══════ TAB 2: Contract Comparison ═══════
        with gr.Tab("🔀 Compare Contracts"):
            with gr.Row():
                with gr.Column(scale=1):
                    comp_file_a = gr.File(
                        label="📁 Contract A (PDF/DOCX/TXT)",
                        file_types=[".pdf", ".docx", ".doc", ".txt"],
                    )
                    comp_load_a = gr.Button("Load A", variant="secondary", size="sm")
                    comp_status_a = gr.Textbox(label="Status A", interactive=False, lines=1)

                with gr.Column(scale=3):
                    comp_text_a = gr.Textbox(
                        label="Contract A",
                        placeholder="Paste contract A here...",
                        lines=12,
                        show_copy_button=True,
                    )

                with gr.Column(scale=1):
                    comp_file_b = gr.File(
                        label="📁 Contract B (PDF/DOCX/TXT)",
                        file_types=[".pdf", ".docx", ".doc", ".txt"],
                    )
                    comp_load_b = gr.Button("Load B", variant="secondary", size="sm")
                    comp_status_b = gr.Textbox(label="Status B", interactive=False, lines=1)

                with gr.Column(scale=3):
                    comp_text_b = gr.Textbox(
                        label="Contract B",
                        placeholder="Paste contract B here...",
                        lines=12,
                        show_copy_button=True,
                    )

            with gr.Row():
                with gr.Column(scale=1):
                    comp_btn = gr.Button("🔀 Compare Contracts", variant="primary", size="lg")
                with gr.Column(scale=5):
                    comp_status = gr.Textbox(label="Comparison Status", interactive=False, lines=1)

            with gr.Row():
                with gr.Column(scale=4):
                    comp_result_html = gr.HTML(label="Comparison Results")
                with gr.Column(scale=2):
                    comp_json = gr.JSON(label="Raw Comparison Data")

        # ═══════ TAB 3: Contract Q&A Chatbot ═══════
        with gr.Tab("💬 Contract Q&A"):
            gr.HTML("""
            <div style="padding:12px 16px;background:linear-gradient(135deg,#eff6ff,#faf5ff);border-radius:10px;margin-bottom:12px;border:1px solid #e5e7eb;">
                <div style="display:flex;align-items:center;gap:8px;margin-bottom:6px;">
                    <span style="font-size:20px;">💬</span>
                    <h3 style="margin:0;font-size:16px;color:#1f2937;">Contract Q&A Chatbot</h3>
                </div>
                <p style="font-size:12px;color:#6b7280;margin:0;line-height:1.5;">
                    Ask questions about your analyzed contract. The chatbot uses <strong>RAG</strong> (Retrieval-Augmented Generation) 
                    to find relevant clauses and generate accurate answers grounded in your contract text.
                    <br>
                    <strong>Step 1:</strong> Analyze a contract in the "📄 Single Contract Analysis" tab.
                    <strong>Step 2:</strong> Come here and ask questions!
                </p>
            </div>
            """)

            chatbot_index_status = gr.Textbox(
                label="📡 Chatbot Index Status",
                interactive=False,
                lines=1,
                value="⏳ No contract indexed yet — analyze a contract first",
            )

            def _chatbot_fn(message, history, chunks, embeddings, analysis):
                """Wrapper for ChatInterface fn signature."""
                yield from chat_respond(message, history, chunks, embeddings, analysis)

            gr.ChatInterface(
                fn=_chatbot_fn,
                type="messages",
                additional_inputs=[chunks_state, embeddings_state, analysis_state],
                examples=[
                    ["What are the main risks in this contract?"],
                    ["Who are the parties involved?"],
                    ["What happens if the contract is terminated?"],
                    ["Are there any liability limitations?"],
                    ["What are my obligations under this contract?"],
                    ["Is there an arbitration clause?"],
                    ["What is the governing law?"],
                    ["Summarize the key terms in plain language."],
                ],
                title="",
                description="",
            )

    # ── Events ──
    def _load_file(file):
        text, err = parse_document(file) if file else ("", "No file")
        if err and not text:
            return "", err
        return text, "Loaded successfully" if not err else err

    def _analysis_and_index(text):
        """Run analysis AND index for chatbot in one call."""
        # Run the standard analysis
        analysis_outputs = run_analysis(text)

        # Index for chatbot (uses the raw text)
        chunks, embeddings, index_status = index_contract(text)

        # analysis_outputs has 12 items: 8 HTML + json_path + csv_path + status + result
        # We need to add: chunks_state, embeddings_state, chatbot_index_status
        return analysis_outputs + [chunks, embeddings, index_status]

    load_btn.click(_load_file, inputs=[file_input], outputs=[text_input, load_status])
    comp_load_a.click(_load_file, inputs=[comp_file_a], outputs=[comp_text_a, comp_status_a])
    comp_load_b.click(_load_file, inputs=[comp_file_b], outputs=[comp_text_b, comp_status_b])

    scan_btn.click(
        _analysis_and_index,
        inputs=[text_input],
        outputs=[
            summary_html, clauses_html, entities_html, nli_html,
            doc_html, obligations_html, compliance_html, redlining_html,
            json_file, csv_file, status_msg, analysis_state,
            chunks_state, embeddings_state, chatbot_index_status,
        ]
    )

    clear_btn.click(
        lambda: [""] * 8 + [None, None, "", None, [], None, "⏳ No contract indexed"],
        outputs=[
            summary_html, clauses_html, entities_html, nli_html,
            doc_html, obligations_html, compliance_html, redlining_html,
            json_file, csv_file, status_msg, analysis_state,
            chunks_state, embeddings_state, chatbot_index_status,
        ]
    )

    comp_btn.click(
        run_comparison,
        inputs=[comp_text_a, comp_text_b],
        outputs=[comp_result_html, comp_json]
    )

    gr.HTML("""
    <div style="margin-top:24px;padding:16px 0;border-top:1px solid #e5e7eb;text-align:center;">
      <p style="font-size:11px;color:#9ca3af;">
        ⚠️ Not legal advice. For informational purposes only.
        · Model: <a href="https://huggingface.co/Mokshith31/legalbert-contract-clause-classification" style="color:#6b7280;">Legal-BERT + CUAD (41 classes)</a>
        · NER: <a href="https://huggingface.co/matterstack/legal-bert-ner" style="color:#6b7280;">Legal-BERT NER</a>
        · NLI: <a href="https://huggingface.co/cross-encoder/nli-deberta-v3-base" style="color:#6b7280;">DeBERTa-v3 NLI</a>
        · LLM: <a href="https://huggingface.co/Qwen/Qwen2.5-7B-Instruct" style="color:#6b7280;">Qwen2.5-7B</a>
        · OCR: <a href="https://github.com/mindee/doctr" style="color:#6b7280;">docTR</a>
        · Dataset: <a href="https://huggingface.co/datasets/theatticusproject/cuad-qa" style="color:#6b7280;">CUAD</a>
        · <a href="https://huggingface.co/spaces/gaurv007/ClauseGuard" style="color:#6b7280;">ClauseGuard Space</a>
      </p>
    </div>
    """)

if __name__ == "__main__":
    demo.launch()