import pytest
from app.core.document_handlers.ner_processor import NerProcessor

def test_generate_masked_mapping():
    processor = NerProcessor()
    unique_entities = [
        {'text': '李强', 'type': '人名'},
        {'text': '李强', 'type': '人名'},  # Duplicate to test numbering
        {'text': '王小明', 'type': '人名'},
        {'text': 'Acme Manufacturing Inc.', 'type': '英文公司名', 'industry': 'manufacturing'},
        {'text': 'Google LLC', 'type': '英文公司名'},
        {'text': 'A公司', 'type': '公司名称'},
        {'text': 'B公司', 'type': '公司名称'},
        {'text': 'John Smith', 'type': '英文人名'},
        {'text': 'Elizabeth Windsor', 'type': '英文人名'},
        {'text': '华梦龙光伏项目', 'type': '项目名'},
        {'text': '案号12345', 'type': '案号'},
        {'text': '310101198802080000', 'type': '身份证号'},
        {'text': '9133021276453538XT', 'type': '社会信用代码'},
    ]
    linkage = {
        'entity_groups': [
            {
                'group_id': 'g1',
                'group_type': '公司名称',
                'entities': [
                    {'text': 'A公司', 'type': '公司名称', 'is_primary': True},
                    {'text': 'B公司', 'type': '公司名称', 'is_primary': False},
                ]
            },
            {
                'group_id': 'g2',
                'group_type': '人名',
                'entities': [
                    {'text': '李强', 'type': '人名', 'is_primary': True},
                    {'text': '李强', 'type': '人名', 'is_primary': False},
                ]
            }
        ]
    }
    mapping = processor._generate_masked_mapping(unique_entities, linkage)
    # 人名 - Updated for new Chinese name masking rules
    assert mapping['李强'] == '李Q'
    assert mapping['王小明'] == '王XM'
    # 英文公司名
    assert mapping['Acme Manufacturing Inc.'] == 'MANUFACTURING'
    assert mapping['Google LLC'] == 'COMPANY'
    # 公司名同组
    assert mapping['A公司'] == mapping['B公司']
    assert mapping['A公司'].endswith('公司')
    # 英文人名
    assert mapping['John Smith'] == 'J*** S***'
    assert mapping['Elizabeth Windsor'] == 'E*** W***'
    # 项目名
    assert mapping['华梦龙光伏项目'].endswith('项目')
    # 案号
    assert mapping['案号12345'] == '***'
    # 身份证号
    assert mapping['310101198802080000'] == 'XXXXXX'
    # 社会信用代码
    assert mapping['9133021276453538XT'] == 'XXXXXXXX'


def test_chinese_name_pinyin_masking():
    """Test Chinese name masking with pinyin functionality"""
    processor = NerProcessor()
    
    # Test basic Chinese name masking
    test_cases = [
        ("李强", "李Q"),
        ("张韶涵", "张SH"),
        ("张若宇", "张RY"),
        ("白锦程", "白JC"),
        ("王小明", "王XM"),
        ("陈志强", "陈ZQ"),
    ]
    
    surname_counter = {}
    
    for original_name, expected_masked in test_cases:
        masked = processor._mask_chinese_name(original_name, surname_counter)
        assert masked == expected_masked, f"Expected {expected_masked}, got {masked} for {original_name}"
    
    # Test duplicate handling
    duplicate_test_cases = [
        ("李强", "李Q"),
        ("李强", "李Q2"),  # Should be numbered
        ("李倩", "李Q3"),  # Should be numbered
        ("张韶涵", "张SH"),
        ("张韶涵", "张SH2"),  # Should be numbered
        ("张若宇", "张RY"),  # Different initials, should not be numbered
    ]
    
    surname_counter = {}  # Reset counter
    
    for original_name, expected_masked in duplicate_test_cases:
        masked = processor._mask_chinese_name(original_name, surname_counter)
        assert masked == expected_masked, f"Expected {expected_masked}, got {masked} for {original_name}"
    
    # Test edge cases
    edge_cases = [
        ("", ""),  # Empty string
        ("李", "李"),  # Single character
        ("李强强", "李QQ"),  # Multiple characters with same pinyin
    ]
    
    surname_counter = {}  # Reset counter
    
    for original_name, expected_masked in edge_cases:
        masked = processor._mask_chinese_name(original_name, surname_counter)
        assert masked == expected_masked, f"Expected {expected_masked}, got {masked} for {original_name}"


def test_chinese_name_integration():
    """Test Chinese name masking integrated with the full mapping process"""
    processor = NerProcessor()
    
    # Test Chinese names in the full mapping context
    unique_entities = [
        {'text': '李强', 'type': '人名'},
        {'text': '张韶涵', 'type': '人名'},
        {'text': '张若宇', 'type': '人名'},
        {'text': '白锦程', 'type': '人名'},
        {'text': '李强', 'type': '人名'},  # Duplicate
        {'text': '张韶涵', 'type': '人名'},  # Duplicate
    ]
    
    linkage = {
        'entity_groups': [
            {
                'group_id': 'g1',
                'group_type': '人名',
                'entities': [
                    {'text': '李强', 'type': '人名', 'is_primary': True},
                    {'text': '张韶涵', 'type': '人名', 'is_primary': True},
                    {'text': '张若宇', 'type': '人名', 'is_primary': True},
                    {'text': '白锦程', 'type': '人名', 'is_primary': True},
                ]
            }
        ]
    }
    
    mapping = processor._generate_masked_mapping(unique_entities, linkage)
    
    # Verify the mapping results
    assert mapping['李强'] == '李Q'
    assert mapping['张韶涵'] == '张SH'
    assert mapping['张若宇'] == '张RY'
    assert mapping['白锦程'] == '白JC'
    
    # Check that duplicates are handled correctly
    # The second occurrence should be numbered
    assert '李Q2' in mapping.values() or '张SH2' in mapping.values()


def test_lawyer_and_judge_names():
    """Test that lawyer and judge names follow the same Chinese name rules"""
    processor = NerProcessor()
    
    # Test lawyer and judge names
    test_entities = [
        {'text': '王律师', 'type': '律师姓名'},
        {'text': '李法官', 'type': '审判人员姓名'},
        {'text': '张检察官', 'type': '检察官姓名'},
    ]
    
    linkage = {
        'entity_groups': [
            {
                'group_id': 'g1',
                'group_type': '律师姓名',
                'entities': [{'text': '王律师', 'type': '律师姓名', 'is_primary': True}]
            },
            {
                'group_id': 'g2',
                'group_type': '审判人员姓名',
                'entities': [{'text': '李法官', 'type': '审判人员姓名', 'is_primary': True}]
            },
            {
                'group_id': 'g3',
                'group_type': '检察官姓名',
                'entities': [{'text': '张检察官', 'type': '检察官姓名', 'is_primary': True}]
            }
        ]
    }
    
    mapping = processor._generate_masked_mapping(test_entities, linkage)
    
    # These should follow the same Chinese name masking rules
    assert mapping['王律师'] == '王L'
    assert mapping['李法官'] == '李F'
    assert mapping['张检察官'] == '张JC'