drafting_engine.py

"""
Contract Drafting Engine.
Orchestrates clause retrieval, playbook rules, fallback positions,
risk flags, drafting checklist, and verifier pass.
"""

import json
from typing import List, Dict, Optional, Any
from dataclasses import dataclass, asdict

from playbook import (
    get_required_clauses,
    get_fallback_position,
    get_risk_flags,
    get_checklist,
)
from clause_retriever import ClauseRetriever


@dataclass
class DraftingContext:
    contract_type: str
    party_position: str  # pro_company, pro_counterparty, balanced
    deal_context: str
    business_constraints: List[str]
    governing_law: Optional[str] = None
    counterparty_name: Optional[str] = None
    company_name: Optional[str] = None
    deal_value: Optional[str] = None
    term_length: Optional[str] = None


@dataclass
class DraftedClause:
    clause_name: str
    clause_text: str
    source: str
    fallback_applied: bool
    risk_flags: List[Dict[str, str]]
    checklist_items: List[str]
    retrieved_clauses: List[Dict]


@dataclass
class DraftedContract:
    contract_type: str
    context: DraftingContext
    clauses: List[DraftedClause]
    risk_flags: List[Dict[str, Any]]
    checklist: List[Dict[str, Any]]
    verifier_notes: List[str]


class ContractDraftingEngine:
    def __init__(
        self,
        retriever: Optional[ClauseRetriever] = None,
        generator_model_name: Optional[str] = None,
    ):
        self.retriever = retriever or ClauseRetriever()
        self.generator_model_name = generator_model_name
        self.generator = None
        if generator_model_name:
            try:
                from transformers import pipeline
                self.generator = pipeline(
                    "text-generation",
                    model=generator_model_name,
                    device_map="auto",
                )
            except Exception as e:
                print(f"Warning: could not load generator {generator_model_name}: {e}")

    def draft(self, context: DraftingContext) -> DraftedContract:
        required = get_required_clauses(context.contract_type)
        checklist = get_checklist(context.contract_type)
        drafted_clauses: List[DraftedClause] = []
        all_risk_flags: List[Dict[str, Any]] = []

        for clause_name in required:
            # Retrieve precedent clauses
            query = f"{clause_name.replace('_', ' ')} clause for {context.contract_type.replace('_', ' ')}"
            retrieved = self.retriever.retrieve(
                query=query,
                clause_type=clause_name,
                top_k=3,
            )

            # Get fallback position
            fallback = get_fallback_position(clause_name, context.party_position)

            # Generate clause text
            clause_text = self._generate_clause(
                clause_name=clause_name,
                context=context,
                retrieved_clauses=retrieved,
                fallback=fallback,
            )

            # Get risk flags
            flags = get_risk_flags(clause_name)
            # Apply contextual risk analysis
            active_flags = self._evaluate_risk_flags(clause_text, flags, context)
            all_risk_flags.extend([
                {"clause": clause_name, **f} for f in active_flags
            ])

            # Checklist items for this clause
            clause_checklist = [
                c["item"] for c in checklist
                if clause_name.replace("_", " ") in c["item"].lower()
                or c["category"] in clause_name
            ]

            drafted_clauses.append(DraftedClause(
                clause_name=clause_name,
                clause_text=clause_text,
                source="retrieved+generated",
                fallback_applied=fallback is not None,
                risk_flags=active_flags,
                checklist_items=clause_checklist,
                retrieved_clauses=retrieved,
            ))

        # Verifier pass
        verifier_notes = self._verifier_pass(drafted_clauses, context)

        return DraftedContract(
            contract_type=context.contract_type,
            context=context,
            clauses=drafted_clauses,
            risk_flags=all_risk_flags,
            checklist=[{"item": c["item"], "category": c["category"], "checked": False} for c in checklist],
            verifier_notes=verifier_notes,
        )

    def _generate_clause(
        self,
        clause_name: str,
        context: DraftingContext,
        retrieved_clauses: List[Dict],
        fallback: Optional[Dict[str, str]],
    ) -> str:
        # Build prompt from retrieved clauses + playbook
        prompt_parts = [
            f"Draft a {clause_name.replace('_', ' ')} clause for a {context.contract_type.replace('_', ' ')}.",
            f"Party position: {context.party_position}.",
            f"Deal context: {context.deal_context}",
        ]
        if fallback:
            prompt_parts.append(f"Fallback position: {json.dumps(fallback)}")
        if retrieved_clauses:
            prompt_parts.append("Precedent clauses:")
            for rc in retrieved_clauses:
                prompt_parts.append(f"- {rc['clause_text'][:500]}")
        prompt = "\n".join(prompt_parts)

        if self.generator:
            try:
                out = self.generator(
                    prompt,
                    max_new_tokens=512,
                    do_sample=True,
                    temperature=0.3,
                )
                return out[0]["generated_text"][len(prompt):].strip()
            except Exception as e:
                print(f"Generation failed for {clause_name}: {e}")

        # Fallback: template-based generation
        return self._template_clause(clause_name, context, fallback)

    def _template_clause(
        self,
        clause_name: str,
        context: DraftingContext,
        fallback: Optional[Dict[str, str]],
    ) -> str:
        """Simple template-based clause generation when no LLM is available."""
        templates = {
            "limitation_of_liability": (
                "LIMITATION OF LIABILITY. "
                "Except for breaches of confidentiality, IP infringement, or gross negligence, "
                "each party's aggregate liability arising out of this agreement shall not exceed "
                f"{fallback.get('cap', 'the fees paid in the 12 months preceding the claim') if fallback else 'the fees paid in the 12 months preceding the claim'}."
            ),
            "indemnification": (
                "INDEMNIFICATION. "
                f"{context.company_name or 'Company'} shall indemnify {context.counterparty_name or 'Counterparty'} against third-party claims arising from "
                f"{fallback.get('scope', 'IP infringement and breach of confidentiality') if fallback else 'IP infringement and breach of confidentiality'}."
            ),
            "data_protection": (
                "DATA PROTECTION. "
                "Each party shall process personal data in accordance with applicable data protection laws. "
                f"{fallback.get('role', 'The parties shall act as independent controllers') if fallback else 'The parties shall act as independent controllers'}."
            ),
            "termination": (
                "TERMINATION. "
                f"Either party may terminate this agreement {fallback.get('for_convenience', 'for convenience with 60 days notice') if fallback else 'for convenience with 60 days notice'}. "
                "Upon termination, all fees owed survive and data shall be returned within 30 days."
            ),
            "intellectual_property": (
                "INTELLECTUAL PROPERTY. "
                f"{fallback.get('ownership', 'Each party retains its pre-existing IP') if fallback else 'Each party retains its pre-existing IP'}. "
                "All custom deliverables shall be owned as specified in the applicable SOW."
            ),
            "confidentiality": (
                "CONFIDENTIALITY. "
                "Each party agrees to hold all Confidential Information in strict confidence and not disclose it to any third parties except as required by law."
            ),
            "governing_law": (
                f"GOVERNING LAW. This agreement shall be governed by the laws of {context.governing_law or 'the State of Delaware'}, without regard to conflict of laws principles."
            ),
        }
        return templates.get(
            clause_name,
            f"[{clause_name.replace('_', ' ').title()}]. [Placeholder clause for {context.contract_type}.]"
        )

    def _evaluate_risk_flags(
        self,
        clause_text: str,
        flags: List[Dict[str, str]],
        context: DraftingContext,
    ) -> List[Dict[str, str]]:
        active = []
        text_lower = clause_text.lower()
        for flag in flags:
            if flag["flag"] == "NO_CAP" and "cap" not in text_lower and "limited" not in text_lower:
                active.append(flag)
            elif flag["flag"] == "NO_IP_CARVEOUT" and "intellectual property" not in text_lower and "ip" not in text_lower:
                active.append(flag)
            elif flag["flag"] == "NO_DPA" and "data processing" not in text_lower and "dpa" not in text_lower:
                active.append(flag)
            elif flag["flag"] == "NO_CURE_PERIOD" and "cure" not in text_lower:
                active.append(flag)
            elif flag["flag"] == "NO_DATA_RETURN" and "return" not in text_lower and "delete" not in text_lower:
                active.append(flag)
            elif flag["flag"] == "NO_MUTUALITY" and "mutual" not in text_lower:
                active.append(flag)
        return active

    def _verifier_pass(
        self,
        clauses: List[DraftedClause],
        context: DraftingContext,
    ) -> List[str]:
        notes = []
        clause_names = {c.clause_name for c in clauses}
        required = set(get_required_clauses(context.contract_type))
        missing = required - clause_names
        if missing:
            notes.append(f"MISSING CLAUSES: {', '.join(missing)}")

        # Check internal consistency
        has_limitation = any(c.clause_name == "limitation_of_liability" for c in clauses)
        has_indemnity = any(c.clause_name == "indemnification" for c in clauses)
        if has_limitation and has_indemnity:
            notes.append("PASS: Limitation and indemnification both present.")
        if not has_limitation:
            notes.append("WARNING: No limitation of liability clause.")

        # Check for invented terms (basic heuristic)
        for c in clauses:
            if "[Placeholder" in c.clause_text:
                notes.append(f"WARNING: {c.clause_name} contains placeholder text.")

        return notes

    def export(self, contract: DraftedContract, fmt: str = "json") -> str:
        if fmt == "json":
            return json.dumps(asdict(contract), indent=2)
        elif fmt == "markdown":
            lines = [
                f"# {contract.contract_type.replace('_', ' ').title()}",
                "",
                "## Context",
                f"- Party position: {contract.context.party_position}",
                f"- Deal context: {contract.context.deal_context}",
                "",
                "## Clauses",
            ]
            for c in contract.clauses:
                lines.append(f"### {c.clause_name.replace('_', ' ').title()}")
                lines.append(c.clause_text)
                lines.append("")
            lines.append("## Risk Flags")
            for rf in contract.risk_flags:
                lines.append(f"- **{rf['severity']}** [{rf['clause']}]: {rf['description']}")
            lines.append("")
            lines.append("## Verifier Notes")
            for note in contract.verifier_notes:
                lines.append(f"- {note}")
            return "\n".join(lines)
        else:
            raise ValueError(f"Unknown format: {fmt}")