#!/usr/bin/env python3
"""
TFT Configuration Validator

Validates that all TFTConfig declarations in test files have correct feature counts.
Usage: python3 scripts/validate_tft_configs.py
"""

import re
import glob
import sys
from pathlib import Path

def validate_tft_configs():
    """Validate all TFT configurations in test files."""

    # Find all TFT test files
    test_dir = Path("ml/tests")
    test_files = list(test_dir.glob("tft*.rs")) + list(test_dir.glob("test_tft*.rs"))
    test_files += list(test_dir.glob("gpu_4_model_stress_test.rs"))
    test_files += list(test_dir.glob("ensemble_tft*.rs"))

    total_configs = 0
    valid_configs = 0
    invalid_configs = []

    for filepath in test_files:
        with open(filepath, 'r') as f:
            content = f.read()

        # Find all TFTConfig blocks
        configs = re.finditer(r'TFTConfig\s*\{([^}]+)\}', content, re.DOTALL)

        for match in configs:
            config_text = match.group(1)

            # Extract values
            input_dim_match = re.search(r'input_dim:\s*(\d+)', config_text)
            static_match = re.search(r'num_static_features:\s*(\d+)', config_text)
            known_match = re.search(r'num_known_features:\s*(\d+)', config_text)
            unknown_match = re.search(r'num_unknown_features:\s*(\d+)', config_text)

            if all([input_dim_match, static_match, known_match, unknown_match]):
                total_configs += 1

                input_dim = int(input_dim_match.group(1))
                static = int(static_match.group(1))
                known = int(known_match.group(1))
                unknown = int(unknown_match.group(1))
                total = static + known + unknown

                if total == input_dim:
                    valid_configs += 1
                else:
                    line_num = content[:match.start()].count('\n') + 1
                    invalid_configs.append({
                        'file': str(filepath),
                        'line': line_num,
                        'input_dim': input_dim,
                        'static': static,
                        'known': known,
                        'unknown': unknown,
                        'total': total
                    })

    # Print results
    print("=" * 70)
    print("TFT Configuration Validation Report")
    print("=" * 70)
    print(f"\nTotal configurations found: {total_configs}")
    print(f"Valid configurations: {valid_configs}")
    print(f"Invalid configurations: {len(invalid_configs)}")
    print(f"\nSuccess rate: {valid_configs}/{total_configs} ({100*valid_configs//total_configs if total_configs > 0 else 0}%)")

    if invalid_configs:
        print("\n" + "=" * 70)
        print("INVALID CONFIGURATIONS FOUND:")
        print("=" * 70)

        for config in invalid_configs:
            print(f"\n{config['file']}:{config['line']}")
            print(f"  input_dim={config['input_dim']}")
            print(f"  static({config['static']}) + known({config['known']}) + unknown({config['unknown']}) = {config['total']}")
            print(f"  ❌ Mismatch: {config['total']} != {config['input_dim']}")
            print(f"  Fix: Change num_unknown_features to {config['input_dim'] - config['static'] - config['known']}")

        print("\n" + "=" * 70)
        print("VALIDATION FAILED")
        print("=" * 70)
        return 1
    else:
        print("\n" + "=" * 70)
        print("✓ ALL CONFIGURATIONS VALID")
        print("=" * 70)
        return 0

if __name__ == "__main__":
    sys.exit(validate_tft_configs())