"""
Self-contained contract drafting system + evaluation.
Downloads all code from the repo and runs everything.
"""

import os
import sys
import json

# Install deps
os.system("pip install -q datasets rank-bm25 sentence-transformers numpy")

# Download source files from the Hub
from huggingface_hub import hf_hub_download

repo = "narcolepticchicken/contract-drafting-assistant"
for fname in ["playbook.py", "clause_retriever.py", "drafting_engine.py", "eval_runner.py"]:
    path = hf_hub_download(repo_id=repo, filename=fname)
    os.makedirs("/app", exist_ok=True)
    os.system(f"cp {path} /app/{fname}")

sys.path.insert(0, "/app")

from clause_retriever import ClauseRetriever
from drafting_engine import ContractDraftingEngine, DraftingContext
from eval_runner import EvalRunner, GOLD_TASKS

print("=" * 60)
print("CONTRACT DRAFTING ASSISTANT - FULL EVALUATION RUN")
print("=" * 60)

# Build retriever with seed clauses
print("\n[1] Building clause retriever...")
retriever = ClauseRetriever(use_bm25=True, use_embeddings=False)
try:
    from datasets import load_dataset
    ds = load_dataset("asapworks/Contract_Clause_SampleDataset", split="train")
    for row in ds:
        retriever.add_clauses([{
            "clause_text": row["clause_text"],
            "clause_type": row.get("clause_type", "unknown"),
            "source": row.get("file", "seed"),
        }])
    print(f"  Loaded {len(retriever.corpus)} seed clauses from asapworks/Contract_Clause_SampleDataset")
except Exception as e:
    print(f"  Warning: could not load seed clauses: {e}")

# Also try loading ContractNLI for additional context
try:
    ds = load_dataset("kiddothe2b/contract-nli", "contractnli_a", split="train")
    for row in ds:
        retriever.add_clauses([{
            "clause_text": row["premise"],
            "clause_type": "nda_clause",
            "source": "contract-nli",
        }])
    print(f"  Loaded contract-nli premises")
except Exception as e:
    print(f"  Warning: could not load contract-nli: {e}")

# Build engine
print("\n[2] Initializing drafting engine...")
engine = ContractDraftingEngine(retriever=retriever)

# Run eval suite
print("\n[3] Running gold task evaluation suite...")
runner = EvalRunner(engine)
results = runner.run_suite(GOLD_TASKS)

# Print report
report = runner.report(results)
print(report)

# Save report
with open("/app/eval_report.md", "w") as f:
    f.write(report)

# Save detailed JSON
detailed = []
for r in results:
    detailed.append({
        "task_id": r.task_id,
        "contract_type": r.contract_type,
        "total_score": r.total_score,
        "scores": r.scores,
    })
with open("/app/eval_results.json", "w") as f:
    json.dump(detailed, f, indent=2)

# Generate sample drafted agreements
print("\n[4] Generating sample drafted agreements...")
samples = [
    DraftingContext(
        contract_type="saas_agreement",
        party_position="pro_company",
        deal_context="Enterprise SaaS for financial analytics. Customer is a mid-size bank.",
        business_constraints=["SOC 2 Type II", "annual billing", "99.9% uptime"],
        governing_law="Delaware",
        company_name="FinAnalytics Inc",
        counterparty_name="MidSize Bank",
    ),
    DraftingContext(
        contract_type="nda",
        party_position="balanced",
        deal_context="Mutual NDA for M&A discussions between two tech companies.",
        business_constraints=["3 year term", "mutual obligations", "return of information"],
        governing_law="California",
        company_name="TechCorp A",
        counterparty_name="TechCorp B",
    ),
    DraftingContext(
        contract_type="dpa",
        party_position="balanced",
        deal_context="GDPR DPA for SaaS provider processing EU personal data.",
        business_constraints=["GDPR compliant", "subprocessor list", "audit rights"],
        governing_law="Ireland",
        company_name="CloudProvider",
        counterparty_name="EU Controller",
    ),
]

for ctx in samples:
    contract = engine.draft(ctx)
    md = engine.export(contract, fmt="markdown")
    fname = f"/app/sample_{ctx.contract_type}_{ctx.party_position}.md"
    with open(fname, "w") as f:
        f.write(md)
    print(f"  Saved {fname}")

# Generate failure report
print("\n[5] Generating failure report...")
failure_lines = ["# Failure Report", ""]
for r in results:
    failures = []
    for dim, score in r.scores.items():
        if score < 0.5:
            failures.append(f"  - {dim}: {score:.3f} (BELOW THRESHOLD)")
    if failures:
        failure_lines.append(f"## {r.task_id} ({r.contract_type})")
        failure_lines.extend(failures)
        failure_lines.append("")
        failure_lines.append("Root causes:")
        if r.scores.get("clause_completeness", 1) < 0.5:
            failure_lines.append("  - Missing required clauses - template coverage insufficient")
        if r.scores.get("business_usefulness", 1) < 0.5:
            failure_lines.append("  - Business constraints not reflected - need constraint-aware generation")
        if r.scores.get("risk_flag_accuracy", 1) < 0.5:
            failure_lines.append("  - Risk flag detection heuristics too simplistic")
        if r.scores.get("citation_support", 1) < 0.5:
            failure_lines.append("  - No retrieved precedent clauses for this clause type")
        failure_lines.append("")

if len(failure_lines) <= 2:
    failure_lines.append("No critical failures detected. All scores above 0.5 threshold.")
    failure_lines.append("")
    failure_lines.append("Known limitations:")
    failure_lines.append("  - Template-based generation produces short, generic clauses")
    failure_lines.append("  - Risk flag detection uses keyword matching, not semantic understanding")
    failure_lines.append("  - Playbook compliance scoring is heuristic-based")
    failure_lines.append("  - No LLM-based generation when running without GPU/transformers")
    failure_lines.append("  - Business constraints checked via string matching, not semantic embedding")

with open("/app/failure_report.md", "w") as f:
    f.write("\n".join(failure_lines))

print("Done! All outputs saved to /app/")