#!/usr/bin/env python3
"""
ARC-AGI NeuroGolf Championship - Main Entry Point

Usage:
  python -m neurogolf_solver.main --data_dir ARC-AGI/data/training/ --output_dir submission
  python -m neurogolf_solver.main --kaggle --output_dir /kaggle/working/submission
  python -m neurogolf_solver.main --data_dir ARC-AGI/data/training/ --arcgen_dir ARC-GEN-100K/ --use_wandb
"""

import argparse
import os
import sys
import time
import onnxruntime as ort
from .config import get_providers
from .data_loader import load_tasks_dir, load_tasks_kaggle
from .submission import run_tasks, generate_submission, print_summary
from .profiler import score_network
from .constants import EXCLUDED_TASKS, MAX_ONNX_FILESIZE

try:
    import wandb
except ImportError:
    wandb = None


def check_all_models(output_dir, strict_size, strict_score):
    """Check all .onnx files for size limit and scoreability."""
    size_problems = []
    score_problems = []
    for f in sorted(os.listdir(output_dir)):
        if not f.endswith('.onnx'):
            continue
        fpath = os.path.join(output_dir, f)
        fsize = os.path.getsize(fpath)
        if fsize > MAX_ONNX_FILESIZE:
            size_problems.append((f, fsize))
        macs, memory, params = score_network(fpath)
        if macs is None or memory is None or params is None:
            score_problems.append(f)
    if size_problems:
        print(f"\n{'!'*70}")
        print(f"FATAL: {len(size_problems)} .onnx files exceed 1.44MB limit:")
        for f, sz in size_problems:
            print(f"  {f}: {sz:,} bytes ({sz/1024:.1f} KB)")
        print(f"{'!'*70}")
        if strict_size:
            sys.exit(1)
    if score_problems:
        print(f"\nWARNING: {len(score_problems)} .onnx files unscorable by onnx_tool:")
        for f in score_problems:
            print(f"  {f}")
        if strict_score:
            print("Stopping (--strict_score is on).")
            sys.exit(1)
    if not size_problems and not score_problems:
        print(f"\nAll .onnx files pass size and score checks.")


def main():
    parser = argparse.ArgumentParser(description='NeuroGolf Solver v5')
    parser.add_argument('--data_dir', default='ARC-AGI/data/training/')
    parser.add_argument('--arcgen_dir', default='', help='Path to ARC-GEN-100K/ directory')
    parser.add_argument('--output_dir', default='/kaggle/working/submission')
    parser.add_argument('--kaggle', action='store_true', help='Use Kaggle task format')
    parser.add_argument('--conv_budget', type=float, default=30.0, help='Seconds per conv solver per task')
    parser.add_argument('--tasks', type=str, default='', help='Comma-separated task numbers')
    parser.add_argument('--device', type=str, default='auto', choices=['auto', 'cpu', 'cuda'])
    parser.add_argument('--use_wandb', action='store_true', help='Enable W&B logging')
    parser.add_argument('--strict_size', type=bool, default=True, help='Halt if any .onnx > 1.44MB (default: True)')
    parser.add_argument('--strict_score', type=bool, default=False, help='Halt if any model unscorable (default: False)')
    args = parser.parse_args()
    providers = get_providers(args.device)
    config = {"device": args.device, "conv_budget": args.conv_budget, "data_dir": args.data_dir, "arcgen_dir": args.arcgen_dir, "tasks": args.tasks}
    ort.set_default_logger_severity(3)
    print(f"Using providers: {providers}")
    print(f"Strict size: {args.strict_size} | Strict score: {args.strict_score}")
    print(f"Max .onnx file size: {MAX_ONNX_FILESIZE:,} bytes")
    if args.kaggle:
        tasks = load_tasks_kaggle(args.data_dir)
    else:
        arcgen = args.arcgen_dir if args.arcgen_dir else None
        tasks = load_tasks_dir(args.data_dir, arcgen_dir=arcgen)
    total_arcgen = sum(len(t['data'].get('arc-gen', [])) for t in tasks.values())
    print(f"Loaded {len(tasks)} tasks ({total_arcgen} ARC-GEN examples)")
    task_nums = [int(t) for t in args.tasks.split(',')] if args.tasks else sorted(tasks.keys())
    print(f"Solving {len(task_nums)} tasks")
    print(f"Conv budget: {args.conv_budget}s per task")
    print("=" * 70)
    t0 = time.time()
    if args.use_wandb and wandb is not None:
        with wandb.init(project="neurogolf", name="solver_run", config=config):
            results, costs_dict, total_score = run_tasks(task_nums, tasks, args.output_dir, providers, args.conv_budget, EXCLUDED_TASKS, use_wandb=True)
    else:
        results, costs_dict, total_score = run_tasks(task_nums, tasks, args.output_dir, providers, args.conv_budget, EXCLUDED_TASKS, use_wandb=False)
    elapsed = time.time() - t0
    check_all_models(args.output_dir, args.strict_size, args.strict_score)
    submission_info = generate_submission(args.output_dir, results, costs_dict, task_nums)
    print_summary(results, submission_info, elapsed)

if __name__ == '__main__':
    main()