app.py

"""
ClauseGuard — World's Best Legal Contract Analysis Tool
════════════════════════════════════════════════════════
Features:
  • 41 CUAD clause categories via fine-tuned Legal-BERT
  • 4-tier risk scoring (Critical / High / Medium / Low)
  • Legal NER: parties, dates, monetary values, jurisdictions, defined terms
  • NLI contradiction & missing-clause detection
  • PDF / DOCX / TXT parsing
  • Professional 3-panel Gradio UI
  • JSON & CSV export

Models:
  • Clause classifier: Mokshith31/legalbert-contract-clause-classification
    (LoRA adapter on nlpaueb/legal-bert-base-uncased, 41 CUAD classes)
"""

import os
import re
import json
import csv
import io
import textwrap
from collections import defaultdict
from datetime import datetime

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) ────────────────────────────────
try:
    import torch
    from transformers import AutoTokenizer, AutoModelForSequenceClassification
    from peft import PeftModel
    _HAS_TORCH = True
except Exception:
    _HAS_TORCH = False

# ═══════════════════════════════════════════════════════════════════════
# 1. CONFIGURATION
# ═══════════════════════════════════════════════════════════════════════

CUAD_LABELS = [
    "Document Name", "Parties", "Agreement Date", "Effective Date",
    "Expiration Date", "Renewal Term", "Governing Law", "Most Favored Nation",
    "Non-Compete", "Exclusivity", "No-Solicit of Customers",
    "No-Solicit of Employees", "Non-Disparagement",
    "Termination for Convenience", "ROFR/ROFO/ROFN", "Change of Control",
    "Anti-Assignment", "Revenue/Profit Sharing", "Price Restriction",
    "Minimum Commitment", "Volume Restriction", "IP Ownership Assignment",
    "Joint IP Ownership", "License Grant", "Non-Transferable License",
    "Affiliate License-Licensor", "Affiliate License-Licensee",
    "Unlimited/All-You-Can-Eat License", "Irrevocable or Perpetual License",
    "Source Code Escrow", "Post-Termination Services", "Audit Rights",
    "Uncapped Liability", "Cap on Liability", "Liquidated Damages",
    "Warranty Duration", "Insurance", "Covenant Not to Sue",
    "Third Party Beneficiary", "Other"
]

# Original 8 unfair-clause labels (backward-compat + consumer contracts)
_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",
    # 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",
    "Liquidated Damages": "HIGH",
    "Anti-Assignment": "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",
    # 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",
    "Warranty Duration": "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.",
})

# Risk weights for scoring
RISK_WEIGHTS = {"CRITICAL": 40, "HIGH": 20, "MEDIUM": 10, "LOW": 3}

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

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

cuad_tokenizer = None
cuad_model = None

def _load_cuad_model():
    global cuad_tokenizer, cuad_model
    if not _HAS_TORCH:
        print("[ClauseGuard] PyTorch not available — using regex fallback")
        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()
        print("[ClauseGuard] CUAD model loaded successfully")
    except Exception as e:
        print(f"[ClauseGuard] CUAD model load failed: {e}")
        cuad_tokenizer = None
        cuad_model = None

_load_cuad_model()

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

def parse_pdf(file_path):
    if not _HAS_PDF:
        return None, "PDF parsing not available (pdfplumber not installed)"
    try:
        text = ""
        with pdfplumber.open(file_path) as pdf:
            for page in pdf.pages:
                page_text = page.extract_text()
                if page_text:
                    text += page_text + "\n\n"
        return text.strip(), None
    except Exception as e:
        return None, f"PDF parse error: {e}"

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. CLAUSE DETECTION
# ═══════════════════════════════════════════════════════════════════════

def split_clauses(text):
    """Split contract text into individual clauses."""
    text = re.sub(r'\n{3,}', '\n\n', text.strip())
    parts = re.split(
        r'(?<=[.!?])\s+(?=[A-Z0-9(])|(?:\n\n)(?=\d+[.)]\s|\([a-z]\)\s|[A-Z][A-Z\s]{2,})',
        text
    )
    clauses = []
    for p in parts:
        p = p.strip()
        if len(p) > 30:
            clauses.append(p)
    return clauses

def classify_cuad(clause_text):
    """Classify a single clause using the CUAD model."""
    if cuad_model is None or cuad_tokenizer is None:
        return _classify_regex(clause_text)

    try:
        inputs = cuad_tokenizer(
            clause_text,
            return_tensors="pt",
            truncation=True,
            max_length=256,
            padding=True
        )
        with torch.no_grad():
            logits = cuad_model(**inputs).logits
        probs = torch.softmax(logits, dim=-1)[0]

        # Multi-label: return all labels above threshold
        threshold = 0.15
        results = []
        for i, prob in enumerate(probs):
            if prob > threshold and i < len(CUAD_LABELS):
                label = CUAD_LABELS[i]
                risk = RISK_MAP.get(label, "LOW")
                results.append({
                    "label": label,
                    "confidence": round(float(prob), 3),
                    "risk": risk,
                    "description": DESC_MAP.get(label, label),
                })
        # Sort by confidence descending
        results.sort(key=lambda x: x["confidence"], reverse=True)
        # If no labels above threshold, take top-1
        if not results:
            top_idx = int(probs.argmax())
            label = CUAD_LABELS[top_idx] if top_idx < len(CUAD_LABELS) else "Other"
            results.append({
                "label": label,
                "confidence": round(float(probs[top_idx]), 3),
                "risk": RISK_MAP.get(label, "LOW"),
                "description": DESC_MAP.get(label, label),
            })
        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", 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"],
}

def _classify_regex(text):
    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": 0.7,
                        "risk": risk,
                        "description": DESC_MAP.get(label, label),
                    })
                    seen.add(label)
                break
    return results

# ═══════════════════════════════════════════════════════════════════════
# 5. LEGAL NER
# ═══════════════════════════════════════════════════════════════════════

def extract_entities(text):
    """Extract legal entities using regex patterns."""
    entities = []

    # Dates
    date_patterns = [
        (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}-\d{1,2}-\d{2,4}\b', "DATE"),
        (r'\b(?:Effective|Commencement|Expiration|Termination)\s+Date\b', "DATE_REF"),
    ]
    for pat, etype in date_patterns:
        for m in re.finditer(pat, text, re.IGNORECASE):
            entities.append({"text": m.group(), "type": etype, "start": m.start(), "end": m.end()})

    # Monetary values
    money_patterns = [
        (r'\$\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)', "MONEY"),
    ]
    for pat, etype in money_patterns:
        for m in re.finditer(pat, text, re.IGNORECASE):
            entities.append({"text": m.group(), "type": etype, "start": m.start(), "end": m.end()})

    # Party names
    party_patterns = [
        (r'\b[A-Z][A-Za-z0-9\s&]+(?:Inc\.|LLC|Ltd\.|Limited|Corp\.|Corporation|PLC|GmbH|AG|S\.A\.|B\.V\.)\b', "PARTY"),
        (r'\b(?:Party A|Party B|Licensor|Licensee|Buyer|Seller|Tenant|Landlord|Employer|Employee|Company|Customer)\b', "PARTY_ROLE"),
    ]
    for pat, etype in party_patterns:
        for m in re.finditer(pat, text):
            entities.append({"text": m.group(), "type": etype, "start": m.start(), "end": m.end()})

    # Jurisdictions
    jurisdiction_patterns = [
        (r'\b(?:State|Laws?) of [A-Z][a-zA-Z\s]+', "JURISDICTION"),
        (r'\b(?:California|Delaware|New York|Texas|Florida|England|Ireland|Germany|France|Singapore|Hong Kong)\b', "JURISDICTION"),
    ]
    for pat, etype in jurisdiction_patterns:
        for m in re.finditer(pat, text, re.IGNORECASE):
            entities.append({"text": m.group(), "type": etype, "start": m.start(), "end": m.end()})

    # Defined Terms
    defined_patterns = [
        (r'"([A-Z][A-Z\s]+)"', "DEFINED_TERM"),
        (r'\(([A-Z][A-Z\s]+)\)', "DEFINED_TERM"),
    ]
    for pat, etype in defined_patterns:
        for m in re.finditer(pat, text):
            entities.append({"text": m.group(1), "type": etype, "start": m.start(), "end": m.end()})

    # Deduplicate overlapping
    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

# ═══════════════════════════════════════════════════════════════════════
# 6. NLI / CONTRADICTION DETECTION
# ═══════════════════════════════════════════════════════════════════════

_CONTRADICTION_PAIRS = [
    (["Uncapped Liability", "unlimited liability"], ["Cap on Liability", "cap on liability"],
     "Liability cannot be both uncapped and capped simultaneously."),
    (["Governing Law"], ["Governing Law"],
     "Multiple governing law provisions detected — verify consistency."),
    (["Termination for Convenience", "terminat.*convenience"], ["Fixed Term", "fixed term"],
     "Contract has both fixed term and termination for convenience — review carefully."),
    (["IP Ownership Assignment", "assign.*ip"], ["Joint IP Ownership", "joint ownership"],
     "IP cannot be both fully assigned and jointly owned."),
]

def detect_contradictions(clause_results):
    """Detect contradictions and missing critical clauses."""
    contradictions = []
    labels_found = set()
    texts_found = {}

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

    # Contradiction pairs
    for group_a, group_b, explanation in _CONTRADICTION_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": list(set(group_a + group_b)),
            })

    # Missing critical clauses
    critical_clauses = ["Governing Law", "Termination for Convenience", "Limitation of liability", "Arbitration"]
    for cc in critical_clauses:
        if cc not in labels_found:
            contradictions.append({
                "type": "MISSING",
                "explanation": f"Critical clause '{cc}' not detected in the document.",
                "severity": "MEDIUM",
                "clauses": [cc],
            })

    return contradictions

# ═══════════════════════════════════════════════════════════════════════
# 7. 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

# ═══════════════════════════════════════════════════════════════════════
# 8. 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"

    # Analyze each clause
    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"],
                })

    # NER
    entities = extract_entities(text)

    # NLI / contradictions
    contradictions = detect_contradictions(clause_results)

    # Risk scoring
    risk, grade, sev_counts = compute_risk_score(clause_results, len(clauses))

    # Build result object
    result = {
        "metadata": {
            "analysis_date": datetime.now().isoformat(),
            "total_clauses": len(clauses),
            "flagged_clauses": len(set(cr["text"] for cr in clause_results)),
            "model": "Legal-BERT + CUAD (41 classes)" if cuad_model else "Regex fallback",
        },
        "risk": {
            "score": risk,
            "grade": grade,
            "breakdown": sev_counts,
        },
        "clauses": clause_results,
        "entities": entities,
        "contradictions": contradictions,
        "raw_text": text,
    }

    return result, None

# ═══════════════════════════════════════════════════════════════════════
# 9. EXPORT FUNCTIONS
# ═══════════════════════════════════════════════════════════════════════

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"])
    for cr in result.get("clauses", []):
        writer.writerow([
            cr.get("text", "")[:500],
            cr.get("label", ""),
            cr.get("risk", ""),
            cr.get("confidence", ""),
            cr.get("description", ""),
        ])
    return output.getvalue()

# ═══════════════════════════════════════════════════════════════════════
# 10. UI RENDERING
# ═══════════════════════════════════════════════════════════════════════

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>Engine: {result['metadata']['model']}
      </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>'

    # Group by clause text
    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]
            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;">{item["label"]} ({item["confidence"]})</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>'

    # Group by type
    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",
            "PARTY": "#8b5cf6", "PARTY_ROLE": "#a78bfa",
            "JURISDICTION": "#f59e0b",
            "DEFINED_TERM": "#ec4899",
        }.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 "📋"
        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>
          </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",
                "PARTY": "#ddd6fe", "PARTY_ROLE": "#ddd6fe",
                "JURISDICTION": "#fde68a",
                "DEFINED_TERM": "#fbcfe8",
            }.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>
    """

# ═══════════════════════════════════════════════════════════════════════
# 11. 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:
        return [""] * 5 + [None, None, "Document too short (minimum 50 characters)"]

    result, error = analyze_contract(text)
    if error:
        err_html = f'<p style="color:#dc2626;padding:16px;">{error}</p>'
        return [err_html] * 5 + [None, None, error]

    # Save export files
    json_path = "/tmp/clauseguard_report.json"
    with open(json_path, "w") as f:
        json.dump(result, f, indent=2, default=str)

    csv_content = export_csv(result)
    csv_path = "/tmp/clauseguard_report.csv"
    with open(csv_path, "w") as f:
        f.write(csv_content)

    return [
        render_summary(result),
        render_clause_cards(result),
        render_entities(result),
        render_contradictions(result),
        render_document_viewer(result),
        json_path,
        csv_path,
        "Analysis complete",
    ]

def do_clear():
    return [""] * 5 + [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."""

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

    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 · NER · NLI</p>
      </div>
      <div style="font-size:12px;color:#9ca3af;">v2.0 · World's Best Legal AI</div>
    </div>
    """)

    # ── Upload / Input ──
    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...",
                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]],
            inputs=[text_input],
            label="Example Contracts",
        )

    # ── Results ──
    with gr.Row():
        # Left: Summary + Export
        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")

        # Center: Main Content Tabs
        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")

    # ── 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

    load_btn.click(_load_file, inputs=[file_input], outputs=[text_input, load_status])

    scan_btn.click(
        run_analysis,
        inputs=[text_input],
        outputs=[summary_html, clauses_html, entities_html, nli_html,
                 doc_html, json_file, csv_file, status_msg]
    )

    clear_btn.click(
        do_clear,
        outputs=[summary_html, clauses_html, entities_html, nli_html,
                 doc_html, json_file, csv_file, status_msg]
    )

    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</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()