"""Test the classification agent with sample contract clauses."""

import json
import os
from agents.classification_agent import classify_clause, classification_node

OUTPUT_DIR = os.path.join(os.path.dirname(__file__), "..", "outputs")

SAMPLE_CLAUSES = [
    {
        "id": 1,
        "text": "This Agreement shall be governed by and construed in accordance with the laws of the State of Delaware.",
        "section": "General",
    },
    {
        "id": 2,
        "text": "During the term of this Agreement and for a period of two years thereafter, the Employee shall not directly or indirectly engage in any business that competes with the Company.",
        "section": "Restrictive Covenants",
    },
    {
        "id": 3,
        "text": "Either party may terminate this Agreement at any time upon thirty (30) days prior written notice to the other party.",
        "section": "Termination",
    },
    {
        "id": 4,
        "text": "The Company shall indemnify and hold harmless the Contractor from any claims, damages, or expenses arising from the Company's breach of this Agreement.",
        "section": "Liability",
    },
    {
        "id": 5,
        "text": "The Licensor grants to the Licensee a non-exclusive, non-transferable license to use the Software solely for internal business purposes.",
        "section": "License",
    },
]


def test_single_clause():
    """Test classifying a single clause."""
    result = classify_clause(SAMPLE_CLAUSES[0]["text"])
    print(f"Clause: {SAMPLE_CLAUSES[0]['text'][:80]}...")
    print(f"Type:   {result['clause_type']}")
    print(f"Conf:   {result['confidence']}")
    print(f"Why:    {result['reasoning']}")
    print()


def test_classification_node():
    """Test the full LangGraph node with multiple clauses."""
    state = {"clauses": SAMPLE_CLAUSES}
    result = classification_node(state)

    print(f"{'ID':<4} {'Classified As':<35} {'Confidence':<12} {'Section'}")
    print("-" * 80)
    for clause in result["classified_clauses"]:
        print(f"{clause['id']:<4} {clause['clause_type']:<35} {clause['confidence']:<12.2f} {clause['section']}")

    # Save results to JSON
    output_path = os.path.join(OUTPUT_DIR, "classification_results.json")
    with open(output_path, "w") as f:
        json.dump(result["classified_clauses"], f, indent=2)
    print(f"\nResults saved to {output_path}")


if __name__ == "__main__":
    print("Single Clause Test \n")
    test_single_clause()

    print("Full Node Test \n")
    test_classification_node()