#!/usr/bin/env python3
"""Submission file generation and task running with W&B logging."""

import os
import csv
import io
import math
import zipfile
from collections import Counter
from .profiler import score_network
from .constants import MAX_ONNX_FILESIZE, EXCLUDED_TASKS

try:
    import wandb
except ImportError:
    wandb = None


def run_tasks(task_nums, tasks, output_dir, providers, conv_budget, excluded_tasks, use_wandb):
    """Run all tasks and collect results.
    
    Returns: (results, costs_dict, total_score)
    """
    from .solvers.solver_registry import solve_task

    results = {}
    costs_dict = {}
    total_score = 0

    for tn in task_nums:
        if tn not in tasks:
            continue

        td = tasks[tn]['data']
        ok, sname, sz, t_task, model_path = solve_task(
            tn, td, output_dir, providers, conv_budget, excluded_tasks
        )

        if ok:
            profile_result = score_network(model_path)
            if profile_result is None or profile_result[0] is None:
                memory, params = 0, 0
            else:
                memory, params = profile_result[0], profile_result[1]
            
            cost = memory + params
            score = max(1.0, 25.0 - math.log(max(1, cost)))
            total_score += score

            # Check per-file size limit
            if sz and sz > MAX_ONNX_FILESIZE:
                print(f"Task {tn:3d}: {sname:25s} OVER SIZE LIMIT ({sz:,} > {MAX_ONNX_FILESIZE:,})")
                continue

            results[tn] = (sname, t_task, sz)
            costs_dict[tn] = cost
            print(f"Task {tn:3d}: {sname:25s} {score:7.3f} {cost:>12} {t_task:7.3f}s  ({sz:>8,} bytes)")
        else:
            score = 0
            cost = 0
            print(f"Task {tn:3d}: UNSOLVED  {t_task:7.3f}s")

        if use_wandb and wandb is not None:
            wandb.log({
                "task_id": tn,
                "solver": sname if ok else "unsolved",
                "onnx_bytes": sz if ok else 0,
                "task_time_sec": t_task,
                "cost": cost,
                "score": score,
            })

    return results, costs_dict, total_score


def generate_submission(output_dir, results, costs_dict, active_tasks):
    """Generate submission.zip and submission.csv.
    
    Returns dict with submission info.
    """
    n_files = len([f for f in os.listdir(output_dir) if f.endswith('.onnx')])
    total_size = sum(os.path.getsize(os.path.join(output_dir, f))
                     for f in os.listdir(output_dir) if f.endswith('.onnx'))

    # Check per-file size limits
    oversized = []
    for f in os.listdir(output_dir):
        if f.endswith('.onnx'):
            fsize = os.path.getsize(os.path.join(output_dir, f))
            if fsize > MAX_ONNX_FILESIZE:
                oversized.append((f, fsize))

    # Create submission.zip
    parent_dir = os.path.dirname(output_dir) or '/kaggle/working/'
    zip_path = os.path.join(parent_dir, 'submission.zip')
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
        for f in sorted(os.listdir(output_dir)):
            if f.endswith('.onnx'):
                zf.write(os.path.join(output_dir, f), f)
    zip_bytes = buf.getvalue()
    with open(zip_path, 'wb') as f:
        f.write(zip_bytes)
    zip_size = len(zip_bytes)

    # Create submission.csv
    csv_path = os.path.join(parent_dir, 'submission.csv')
    with open(csv_path, 'w', newline='') as f:
        w = csv.writer(f)
        w.writerow(['task_id', 'total_cost'])
        for tn in sorted(costs_dict.keys()):
            w.writerow([f'task{tn:03d}', costs_dict[tn]])

    unsolved_count = len(active_tasks) - len(results)
    total_score = sum(max(1.0, 25.0 - math.log(max(1, cost))) for cost in costs_dict.values())
    total_cost = sum(costs_dict.values())

    return {
        'n_files': n_files,
        'total_size': total_size,
        'zip_path': zip_path,
        'zip_size': zip_size,
        'csv_path': csv_path,
        'total_score': total_score,
        'total_cost': total_cost,
        'unsolved_count': unsolved_count,
        'oversized': oversized,
    }


def print_summary(results, submission_info, elapsed):
    """Print summary statistics."""
    active_count = submission_info['unsolved_count'] + len(results)

    print(f"\n{'=' * 70}")
    print(f"Solved: {len(results)}/{active_count} tasks in {elapsed:.0f}s")
    solver_names = [v[0] for v in results.values()]
    sc = Counter(solver_names)
    for s, c in sc.most_common():
        print(f"  {s}: {c}")

    print(f"\n{submission_info['n_files']} ONNX files, {submission_info['total_size'] / 1024:.1f} KB uncompressed")
    print(f"ZIP size: {submission_info['zip_size'] / 1024:.1f} KB")

    if submission_info['oversized']:
        print(f"WARNING: {len(submission_info['oversized'])} files exceed 1.44MB limit:")
        for f, sz in submission_info['oversized']:
            print(f"  {f}: {sz / 1024:.1f} KB")

    print(f"\nEstimated LB score: {submission_info['total_score']:.1f}")
    print(f"Total cost: {submission_info['total_cost']:,}")
    print(f"Solved: {len(results)} | Unsolved: {submission_info['unsolved_count']}")
    print(f"Written: {submission_info['zip_path']} | {submission_info['csv_path']}")