import pytest
from cleverswarm_python_client.libs.evaluation.detailed_eval import (
    calculate_metrics, normalize_string, normalize_triple,
    evaluate_entity_presence, evaluate_entity_linking,
    evaluate_relation_linking, calculate_precision_recall_f1,
    evaluate_triples, read_jsonl, write_jsonl, load_ontology,
    calculate_average_metrics, evaluate_ontology, CustomEncoder,
    EvaluateResults
)
import json
import tempfile
import os

@pytest.fixture
def sample_ontology():
    return {
        "Relations": {
            "works_at": [["Person", "Organization"]],
            "part_of": [["Organization", "Organization"]],
            "located_in": [["Organization", "Location"]]
        }
    }

@pytest.fixture
def sample_ground_truth():
    return {
        "id": "test1",
        "sent": "John works at Apple in California",
        "triples": [
            {"sub": "John", "rel": "works_at", "obj": "Apple"},
            {"sub": "Apple", "rel": "located_in", "obj": "California"}
        ]
    }

@pytest.fixture
def sample_response():
    return {
        "id": "test1",
        "ents": [
            {"text": "John", "class": "Person"},
            {"text": "Apple", "class": "Organization"},
            {"text": "California", "class": "Location"}
        ],
        "triples": [
            ["John", "works_at", "Apple"],
            ["Apple", "located_in", "California"]
        ]
    }

def test_calculate_metrics():
    # Test normal case
    precision, recall, f1 = calculate_metrics(tp=5, fp=2, fn=1)
    assert precision == round(5/7, 2)
    assert recall == round(5/6, 2)
    assert f1 == round(2 * (5/7 * 5/6) / (5/7 + 5/6), 2)

    # Test edge case with zero values
    precision, recall, f1 = calculate_metrics(tp=0, fp=0, fn=0)
    assert precision == 0
    assert recall == 0
    assert f1 == 0

def test_normalize_string():
    test_cases = [
        ("The Test String", "teststring"),
        ("Test (with parentheses)", "test"),
        ("Test123abc", "test123"),
        ("Company, Inc.", "company"),
        ("Multiple    Spaces", "multiplespace"),
        ("Numbers123", "numbers123"),
        ("The Company", "company"),
        ("Products & Services", "productsservice")
    ]
    
    for input_str, expected in test_cases:
        result = normalize_string(input_str)
        assert result == expected, f"Failed for input '{input_str}': expected '{expected}' but got '{result}'"

def test_normalize_triple():
    test_cases = [
        ("Subject", "Relation", "Object", "subjectrelationobject"),
        ("The Company", "has_product", "Item", "companyhasproductitem"),
        ("Person", "works_at", None, "personworksatnone"),
        ("Test", "test", "", "testtest")
    ]
    
    for sub, rel, obj, expected in test_cases:
        result = normalize_triple(sub, rel, obj)
        assert result == expected, f"Failed for input ({sub}, {rel}, {obj}): expected '{expected}' but got '{result}'"

def test_normalize_string_edge_cases():
    test_cases = [
        ("", ""),  # Empty string
        ("   ", ""),  # Only spaces
        ("The", "the"),  # Single word with 'The'
        ("A B C", "abc"),  # Multiple single letters
        ("Test123Test", "test123"),  # Mixed alphanumeric
        ("Test-123_Test", "test123"),  # With special characters
    ]
    
    for input_str, expected in test_cases:
        result = normalize_string(input_str)
        assert result == expected, f"Failed for input '{input_str}': expected '{expected}' but got '{result}'"

def test_normalize_triple_edge_cases():
    test_cases = [
        ("", "", "", ""),  # All empty
        ("A", "B", "C", "abc"),  # Single letters
        ("The A", "The B", "The C", "abc"),  # With articles
        ("Test(1)", "Test,2", "Test.3", "testtesttest3"),  # With special chars and numbers
    ]
    
    for sub, rel, obj, expected in test_cases:
        result = normalize_triple(sub, rel, obj)
        assert result == expected, f"Failed for input ({sub}, {rel}, {obj}): expected '{expected}' but got '{result}'"

def test_normalize_triple_none_handling():
    # Test with None as object - this works as the function handles it
    result = normalize_triple("Person", "works_at", None)
    assert result == "personworksatnone", f"Failed for None object: expected 'personworksatnone' but got '{result}'"
    
    # Test with empty string instead of None for subject - since function doesn't handle None for subject
    result = normalize_triple("", "works_at", "Company")
    assert result == "worksatcompany", f"Failed for empty subject: expected 'worksatcompany' but got '{result}'"

def test_evaluate_entity_presence(sample_ground_truth, sample_response):
    result = evaluate_entity_presence(sample_ground_truth, sample_response)
    assert isinstance(result, dict)
    assert all(key in result for key in ['precision', 'recall', 'f1', 'has_entities'])
    assert result['has_entities'] == True
    assert result['precision'] > 0
    assert result['recall'] > 0
    assert result['f1'] > 0

def test_evaluate_entity_linking(sample_ground_truth, sample_response, sample_ontology):
    result = evaluate_entity_linking(sample_ground_truth, sample_response, sample_ontology)
    assert isinstance(result, dict)
    assert all(key in result for key in ['accuracy', 'correct_links', 'total_links', 
                                       'gt_class_instances', 'correctly_linked_entities', 
                                       'has_links'])
    assert result['has_links'] == True
    assert result['accuracy'] > 0

def test_evaluate_relation_linking(sample_ground_truth, sample_response):
    correctly_linked_entities = ["john", "apple", "california"]
    result = evaluate_relation_linking(sample_ground_truth, sample_response, correctly_linked_entities)
    assert isinstance(result, dict)
    assert all(key in result for key in ['correct', 'missing', 'wrong', 'correct_triples',
                                       'missing_triples', 'wrong_triples', 'total_triples',
                                       'has_triples'])
    assert result['has_triples'] == True

def test_evaluate_triples(sample_ground_truth, sample_response):
    result = evaluate_triples(sample_ground_truth, sample_response)
    assert isinstance(result, dict)
    assert all(key in result for key in ['triples_precision', 'triples_recall', 'triples_f1'])
    assert result['triples_precision'] > 0
    assert result['triples_recall'] > 0
    assert result['triples_f1'] > 0

def test_calculate_precision_recall_f1():
    gold = {"triple1", "triple2", "triple3"}
    pred = {"triple1", "triple2", "triple4"}
    
    precision, recall, f1 = calculate_precision_recall_f1(gold, pred)
    assert precision == 2/3
    assert recall == 2/3
    assert f1 == 2/3

    # Test empty prediction
    precision, recall, f1 = calculate_precision_recall_f1(gold, set())
    assert precision == 0
    assert recall == 0
    assert f1 == 0

def test_jsonl_operations():
    test_data = [{"test": 1}, {"test": 2}]
    with tempfile.NamedTemporaryFile(mode='w+', delete=False) as f:
        write_jsonl(test_data, f.name)
        result = read_jsonl(f.name)
        assert result == test_data
    os.unlink(f.name)

def test_calculate_average_metrics():
    results = [
        {
            'triples_evaluation': {'triples_precision': 0.8, 'triples_recall': 0.7, 'triples_f1': 0.75},
            'entity_presence': {'precision': 0.9, 'recall': 0.85, 'f1': 0.87},
            'entity_linking': {'accuracy': 0.8, 'has_links': True},
            'relation_linking': {'correct': 0.7, 'missing': 0.2, 'wrong': 0.1, 'has_triples': True}
        }
    ]
    
    avg_metrics = calculate_average_metrics(results, True)
    assert isinstance(avg_metrics, dict)
    assert avg_metrics['triples_evaluation_precision'] == 0.8
    assert avg_metrics['entity_presence_f1'] == 0.87
    assert avg_metrics['entity_linking_accuracy'] == 0.8 

def test_evaluate_ontology(tmp_path):
    # Create temporary test files
    config = {
        'path_patterns': {
            'sys': str(tmp_path / 'outputs' / '$$onto$$_output.jsonl'),
            'gt': str(tmp_path / 'ground_truth' / '$$onto$$_gt.jsonl'),
            'onto': str(tmp_path / 'ontologies' / '$$onto$$.json'),
            'output': str(tmp_path / 'metrics' / 'ont_$$onto$$_eval_results.jsonl')
        }
    }
    
    # Create necessary directories
    os.makedirs(tmp_path / 'outputs', exist_ok=True)
    os.makedirs(tmp_path / 'ground_truth', exist_ok=True)
    os.makedirs(tmp_path / 'ontologies', exist_ok=True)
    
    # Create test files
    test_output = [
        {
            "id": "test1",
            "ents": [{"text": "John", "class": "Person"}],
            "triples": [["John", "works_at", "Company"]]
        }
    ]
    test_gt = [
        {
            "id": "test1",
            "sent": "John works at Company",
            "triples": [{"sub": "John", "rel": "works_at", "obj": "Company"}]
        }
    ]
    test_onto = {
        "Relations": {
            "works_at": [["Person", "Organization"]]
        }
    }
    
    # Write test files
    write_jsonl(test_output, str(tmp_path / 'outputs' / 'test_output.jsonl'))
    write_jsonl(test_gt, str(tmp_path / 'ground_truth' / 'test_gt.jsonl'))
    with open(str(tmp_path / 'ontologies' / 'test.json'), 'w') as f:
        json.dump(test_onto, f)
    
    results = evaluate_ontology('test', config)
    assert isinstance(results, list)
    assert len(results) > 0
    assert 'triples_evaluation' in results[0]
    assert 'entity_presence' in results[0]
    assert 'entity_linking' in results[0]

def test_evaluate_results_with_config(tmp_path):
    config = {
        'onto_list': ['test'],
        'path_patterns': {
            'sys': str(tmp_path / 'outputs' / '$$onto$$_output.jsonl'),
            'gt': str(tmp_path / 'ground_truth' / '$$onto$$_gt.jsonl'),
            'onto': str(tmp_path / 'ontologies' / '$$onto$$.json'),
            'output': str(tmp_path / 'metrics' / 'ont_$$onto$$_eval_results.jsonl')
        },
        'overall_output': str(tmp_path / 'overall_metrics.json')
    }
    
    # Create test files similar to test_evaluate_ontology
    os.makedirs(tmp_path / 'outputs', exist_ok=True)
    os.makedirs(tmp_path / 'ground_truth', exist_ok=True)
    os.makedirs(tmp_path / 'ontologies', exist_ok=True)
    
    test_output = [{"id": "test1", "ents": [], "triples": []}]
    test_gt = [{"id": "test1", "sent": "test", "triples": []}]
    test_onto = {"Relations": {}}
    
    write_jsonl(test_output, str(tmp_path / 'outputs' / 'test_output.jsonl'))
    write_jsonl(test_gt, str(tmp_path / 'ground_truth' / 'test_gt.jsonl'))
    with open(str(tmp_path / 'ontologies' / 'test.json'), 'w') as f:
        json.dump(test_onto, f)
    
    # Test both detailed and non-detailed evaluations
    from cleverswarm_python_client.libs.evaluation.detailed_eval import EvaluateResults
    EvaluateResults(config, is_detailed_metrics=True)
    assert os.path.exists(config['overall_output'])
    
    EvaluateResults(config, is_detailed_metrics=False)
    assert os.path.exists(config['overall_output'])

def test_custom_encoder():
    test_data = {"set_field": {1, 2, 3}}
    encoded = json.dumps(test_data, cls=CustomEncoder)
    decoded = json.loads(encoded)
    assert isinstance(decoded['set_field'], list)
    assert set(decoded['set_field']) == {1, 2, 3}

def test_calculate_average_metrics_edge_cases():
    # Test with empty results
    empty_results = []
    avg_metrics = calculate_average_metrics(empty_results, True)
    assert isinstance(avg_metrics, dict)
    
    # Test with all required fields with minimal data
    incomplete_results = [{
        'triples_evaluation': {
            'triples_precision': 0.8,
            'triples_recall': 0.7,
            'triples_f1': 0.75
        },
        'entity_presence': {
            'precision': 0.7,
            'recall': 0.8,
            'f1': 0.75,
            'has_entities': True
        },
        'entity_linking': {
            'accuracy': 0.6,
            'has_links': True
        },
        'relation_linking': {
            'correct': 0.5,
            'missing': 0.3,
            'wrong': 0.2,
            'has_triples': True
        }
    }]
    
    # Test with detailed metrics
    avg_metrics = calculate_average_metrics(incomplete_results, True)
    assert isinstance(avg_metrics, dict)
    assert 'triples_evaluation_precision' in avg_metrics
    assert 'triples_evaluation_recall' in avg_metrics
    assert 'entity_linking_accuracy' in avg_metrics
    
    # Test with non-detailed metrics
    avg_metrics = calculate_average_metrics(incomplete_results, False)
    assert isinstance(avg_metrics, dict)
    assert 'entity_presence_precision' in avg_metrics
    assert 'triples_evaluation_precision' in avg_metrics