import ast
import os

model_files = [
    "app/ml/models/firco_ensemble_model.py",
    "app/ml/models/firco_tfidf_model.py",
    "app/ml/models/non_firco_roberta_model.py",
    "app/ml/models/firco_roberta_model.py",
    "app/ml/models/firco_bert_model.py"
]

expected_kwargs = {
    'columns', 
    'df', 
    'columns_to_use_in_reasoning', 
    'ground_truth_column_for_reasoning', 
    'language', 
    'number_of_reasonings'
}

print("=" * 80)
print("AST-Based Predict Function Signature Checker")
print("Target: predict_task.py calling conventions")
print("=" * 80)

for file_path in model_files:
    full_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), file_path)
    print(f"\nChecking File: {file_path}")
    
    if not os.path.exists(full_path):
        print(f"  Status: ERROR - File not found! ({full_path})")
        continue
        
    with open(full_path, "r", encoding="utf-8") as f:
        tree = ast.parse(f.read(), filename=full_path)
        
    # Find classes and their predict methods
    predict_found = False
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            for item in node.body:
                if isinstance(item, ast.FunctionDef) and item.name == 'predict':
                    predict_found = True
                    className = node.name
                    print(f"  Class: {className}")
                    
                    # Extract arg names
                    args = [arg.arg for arg in item.args.args]
                    kwonlyargs = [arg.arg for arg in item.args.kwonlyargs]
                    all_args = args + kwonlyargs
                    
                    if 'self' in all_args:
                        all_args.remove('self')
                        
                    print(f"    Signature args: {all_args}")
                    
                    missing = expected_kwargs - set(all_args)
                    # For *args / **kwargs if any
                    has_vararg = item.args.vararg is not None
                    has_kwarg = item.args.kwarg is not None
                    
                    if has_kwarg:
                        missing = set() # **kwargs will absorb everything
                    
                    if not missing:
                        if expected_kwargs.issubset(set(all_args)) or has_kwarg:
                            print("    Status: PERFECT MATCH (can accept all expected kwargs)")
                        else:
                            print("    Status: PARTIAL MATCH (has **kwargs or extra args)")
                    else:
                        print("    Status: Signature MISMATCH with predict_task.py")
                        print(f"      Missing kwargs expected by predict_task: {missing}")
                        
    if not predict_found:
        print("  Status: ERROR - no predict() method found in any class in this file!")

print("\n" + "=" * 80)