import json
import sys
import concurrent.futures
from pathlib import Path
from copienator.commands import correction

def get_missing_tasks():
    """
    Identifies tasks (groups) where NONE of the student IDs in that group
    appear in the existing results for that label in correction.json.
    """
    missing = []

    # correction.results is already loaded from correction.json during 'from copienator.commands import correction'
    # correction.tasks is populated with (filepath, label) during 'from copienator.commands import correction'

    for task in correction.tasks:
        file_path, label = task
        # Find the group metadata file (Group_X.json) to know which IDs are inside
        meta_path = Path(file_path).with_suffix('.json')

        if not meta_path.exists():
            print("Missing meta_path :", meta_path)
            continue

        with open(meta_path, 'r', encoding="utf-8") as f:
            # group_data entries: [pid, ymin, ymax, width_ratio]
            group_data = json.load(f)

        pids_in_group = [str(item[0]) for item in group_data]

        # Check correction.json results for this specific label
        label_results = correction.results.get(label, [])

        # Collect all student IDs that have already been processed for this label
        covered_ids = set()
        for result_list in label_results:
            for entry in result_list:
                covered_ids.add(str(entry.get('id')))

        # Logic: Only process if EVERY ID in this group is missing from correction.json
        if all(pid not in covered_ids for pid in pids_in_group):
            missing.append(task)

    return missing

def main(argv=None):
    if argv:
        print("missing-correction does not accept command-line arguments.", file=sys.stderr)
        return 2
    missing_tasks = get_missing_tasks()
    print("\n Total nb of tasks : ", len(correction.tasks))

    if not missing_tasks:
        print("All groups are already present in correction.json. Nothing to do.")
        return

    print("\nThe following groups are missing from correction.json:")
    for path, label in missing_tasks:
        print(f" - [{label}] {path}")

    confirm = input(f"\nFound {len(missing_tasks)} missing groups. Start processing? (y/N): ")
    if confirm.lower() != 'y':
        print("Aborted.")
        return

    print(f"Processing {len(missing_tasks)} tasks with {correction.NB_THREADS} threads...")

    with concurrent.futures.ThreadPoolExecutor(max_workers=correction.NB_THREADS) as executor:
        # Map tasks to the processing function defined in correction.py
        futures = {executor.submit(correction.process_single_task, t): t for t in missing_tasks}

        for future in concurrent.futures.as_completed(futures):
            try:
                # Handle potential sub-tasks (like label errors) generated during processing
                new_generated_tasks = future.result()
                if new_generated_tasks:
                    for nt in new_generated_tasks:
                        executor.submit(correction.process_single_task, nt)
            except Exception as e:
                t = futures[future]
                print(f"Error processing {t[0]}: {e}")

    print("\nProcessing complete.")

if __name__ == "__main__":
    raise SystemExit(main())