#!/usr/bin/env python3
"""
MEL Protocol Checksum Analyzer
Specifically designed for the MEL protocol hex data

Based on analysis of the provided data patterns.
"""

import sys
from typing import List, Tuple

def parse_hex(hex_str: str) -> List[int]:
    """Parse hex string to bytes"""
    return [int(hex_str[i:i+2], 16) for i in range(0, len(hex_str), 2)]

def analyze_mel_structure(hex_line: str) -> dict:
    """Analyze MEL protocol structure"""
    bytes_data = parse_hex(hex_line.strip())
    
    return {
        'header': bytes_data[0:4],           # 4D454C00
        'length': bytes_data[4],             # Packet length
        'flags': bytes_data[5:8],            # Type/flags
        'sequence': bytes_data[6],           # Sequence number (based on your data)
        'command': bytes_data[8:12],         # Command and zone info
        'zone_mask': bytes_data[12:16],      # Zone targeting
        'reserved': bytes_data[16:28],       # Reserved/padding
        'payload': bytes_data[28:-2],        # Actual payload
        'checksum_bytes': bytes_data[-2:],   # Last 2 bytes
        'checksum_le': bytes_data[-2] | (bytes_data[-1] << 8),  # Little endian
        'checksum_be': (bytes_data[-2] << 8) | bytes_data[-1],  # Big endian
        'full_payload': bytes_data[:-2],     # Everything except checksum
    }

def test_mel_checksums(data: List[int], expected: int) -> List[Tuple[str, int, bool]]:
    """Test checksum algorithms specific to MEL protocol"""
    results = []
    
    # Test 1: Simple sum of all payload bytes
    simple_sum = sum(data) & 0xFFFF
    results.append(("Simple Sum", simple_sum, simple_sum == expected))
    
    # Test 2: Sum with initial value (common in embedded protocols)
    for init_val in [0x0000, 0x5555, 0xAAAA, 0xFFFF, 0x1234, 0x4321, 0x0001]:
        checksum = (sum(data) + init_val) & 0xFFFF
        results.append((f"Sum + 0x{init_val:04x}", checksum, checksum == expected))
    
    # Test 3: Two's complement variations
    sum_val = sum(data)
    twos_comp = (~sum_val + 1) & 0xFFFF
    results.append(("Two's Complement", twos_comp, twos_comp == expected))
    
    # Test 4: One's complement
    ones_comp = (~sum_val) & 0xFFFF
    results.append(("One's Complement", ones_comp, ones_comp == expected))
    
    # Test 5: Subtract from constant
    for const in [0xFFFF, 0x10000, 0x8000, 0x7FFF]:
        checksum = (const - sum_val) & 0xFFFF
        results.append((f"0x{const:04x} - Sum", checksum, checksum == expected))
    
    # Test 6: XOR-based checksums
    xor_result = 0
    for byte in data:
        xor_result ^= byte
    results.append(("XOR all bytes", xor_result, xor_result == expected))
    
    # Test 7: Position-weighted sum
    pos_sum = sum(i * byte for i, byte in enumerate(data)) & 0xFFFF
    results.append(("Position-weighted sum", pos_sum, pos_sum == expected))
    
    # Test 8: Rolling checksum
    rolling = 0
    for byte in data:
        rolling = ((rolling << 1) | (rolling >> 15)) & 0xFFFF
        rolling ^= byte
    results.append(("Rolling XOR", rolling, rolling == expected))
    
    # Test 9: Modular arithmetic variations
    for mod_val in [0x100, 0x101, 0x1FF, 0x200, 0x255, 0x256]:
        if mod_val > 0:
            checksum = sum(data) % mod_val
            results.append((f"Sum mod 0x{mod_val:x}", checksum, checksum == expected))
    
    return results

def analyze_sequence_relationship(filename: str):
    """Analyze relationship between sequence numbers and checksums"""
    with open(filename, 'r') as f:
        lines = [line.strip() for line in f if line.strip()]
    
    print(f"Analyzing sequence-checksum relationship in {filename}")
    print("=" * 60)
    
    sequence_data = []
    
    for i, line in enumerate(lines[:20]):  # First 20 entries
        mel_data = analyze_mel_structure(line)
        
        # Look for the actual sequence field
        # Based on your data, it seems to increment in byte 6
        actual_sequence = mel_data['sequence']
        checksum = mel_data['checksum_le']
        
        sequence_data.append((i, actual_sequence, checksum))
        
        print(f"Entry {i:2d}: seq=0x{actual_sequence:02x} ({actual_sequence:3d}), "
              f"checksum=0x{checksum:04x} ({checksum:5d})")
    
    # Look for patterns
    print("\nSequence vs Checksum Analysis:")
    print("=" * 40)
    
    # Check if checksum changes predictably with sequence
    if len(sequence_data) > 1:
        for i in range(1, min(10, len(sequence_data))):
            seq_diff = sequence_data[i][1] - sequence_data[i-1][1]
            check_diff = sequence_data[i][2] - sequence_data[i-1][2]
            print(f"Entry {i-1}→{i}: seq_diff={seq_diff:2d}, check_diff={check_diff:4d} (0x{check_diff & 0xFFFF:04x})")

def find_checksum_algorithm(filename: str):
    """Main function to find the checksum algorithm"""
    with open(filename, 'r') as f:
        lines = [line.strip() for line in f if line.strip()]
    
    print(f"MEL Protocol Checksum Analysis")
    print(f"File: {filename}")
    print(f"Entries: {len(lines)}")
    print("=" * 60)
    
    # Analyze first few entries
    algorithm_scores = {}
    
    for i, line in enumerate(lines[:10]):
        mel_data = analyze_mel_structure(line)
        
        print(f"\nEntry {i}:")
        print(f"  Hex: {line}")
        print(f"  Sequence: 0x{mel_data['sequence']:02x}")
        print(f"  Expected checksum: 0x{mel_data['checksum_le']:04x} (LE)")
        print(f"  Payload length: {len(mel_data['full_payload'])} bytes")
        
        # Test algorithms
        results = test_mel_checksums(mel_data['full_payload'], mel_data['checksum_le'])
        
        for algo_name, result, is_match in results:
            if is_match:
                print(f"  ✓ {algo_name}: 0x{result:04x}")
                algorithm_scores[algo_name] = algorithm_scores.get(algo_name, 0) + 1
            # Uncomment below to see all results
            # else:
            #     print(f"  ✗ {algo_name}: 0x{result:04x}")
    
    # Summary
    print(f"\n{'='*60}")
    print("ALGORITHM MATCH SUMMARY")
    print(f"{'='*60}")
    
    if algorithm_scores:
        for algo, score in sorted(algorithm_scores.items(), key=lambda x: x[1], reverse=True):
            print(f"{algo}: {score}/10 matches")
            
        best_algo = max(algorithm_scores.items(), key=lambda x: x[1])
        if best_algo[1] >= 8:  # At least 8/10 matches
            print(f"\n🎉 LIKELY ALGORITHM FOUND: {best_algo[0]}")
            print(f"   Confidence: {best_algo[1]}/10 matches")
    else:
        print("No consistent algorithm found with standard methods.")
        print("This may require custom algorithm development.")
    
    # Analyze sequence relationship
    print(f"\n{'='*60}")
    analyze_sequence_relationship(filename)

def verify_algorithm(filename: str, algorithm_name: str):
    """Verify a specific algorithm against all entries"""
    with open(filename, 'r') as f:
        lines = [line.strip() for line in f if line.strip()]
    
    print(f"Verifying algorithm '{algorithm_name}' against {len(lines)} entries...")
    
    matches = 0
    mismatches = []
    
    for i, line in enumerate(lines):
        mel_data = analyze_mel_structure(line)
        expected = mel_data['checksum_le']
        
        # Apply the algorithm (you'd implement the specific one here)
        if algorithm_name == "Simple Sum":
            calculated = sum(mel_data['full_payload']) & 0xFFFF
        elif algorithm_name.startswith("Sum + "):
            init_val = int(algorithm_name.split("0x")[1], 16)
            calculated = (sum(mel_data['full_payload']) + init_val) & 0xFFFF
        else:
            print(f"Algorithm '{algorithm_name}' not implemented in verify function")
            return
        
        if calculated == expected:
            matches += 1
        else:
            mismatches.append((i, expected, calculated))
            if len(mismatches) <= 5:  # Show first 5 mismatches
                print(f"  Mismatch at entry {i}: expected 0x{expected:04x}, got 0x{calculated:04x}")
    
    print(f"Results: {matches}/{len(lines)} matches ({100*matches/len(lines):.1f}%)")
    
    if matches == len(lines):
        print("🎉 PERFECT MATCH! Algorithm verified.")
    elif matches > len(lines) * 0.9:
        print("⚠️  Very close match. May need minor adjustment.")
    else:
        print("❌ Algorithm doesn't work consistently.")

def main():
    if len(sys.argv) < 2:
        print("Usage:")
        print("  python mel_checksum_analyzer.py <hex_file>           # Find algorithm")
        print("  python mel_checksum_analyzer.py <hex_file> verify <algorithm>  # Verify algorithm")
        sys.exit(1)
    
    filename = sys.argv[1]
    
    if len(sys.argv) >= 4 and sys.argv[2] == "verify":
        algorithm = sys.argv[3]
        verify_algorithm(filename, algorithm)
    else:
        find_checksum_algorithm(filename)

if __name__ == "__main__":
    main()