#!/usr/bin/env python3

import logging
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
import config
from utils import process_house, call_model, load_input, llm_answer_to_img, image_answer_to_img

# Setup minimal logging
logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')


def main():
    dataset_dir = Path("./dataset")
    houses = [d.name for d in dataset_dir.iterdir() if d.is_dir()]

    print(f"Processing {len(houses)} houses with {len(config.MODELS)} models for {config.EPOCHS} epochs")

    tasks = []
    for epoch in range(1, config.EPOCHS + 1):
        for house in houses:
            for model in config.MODELS:
                tasks.append((house, model, epoch))

    # Run in parallel with statistics tracking
    successful = 0
    failed = 0
    skipped = 0
    start_time = datetime.now()

    with ThreadPoolExecutor(max_workers=config.NUM_WORKERS) as executor:
        futures = {executor.submit(process_house, house, model, epoch): (house, model, epoch)
                   for house, model, epoch in tasks}

        for future in as_completed(futures):
            house, model, epoch = futures[future]
            try:
                result = future.result()
                if result is None:  # Indicates success or skip
                    successful += 1
                else:
                    failed += 1
            except Exception as e:
                logging.error(f"Task {house}+{model}+{epoch} failed: {e}")
                failed += 1

    # Print summary
    end_time = datetime.now()
    duration = end_time - start_time
    print(f"\n{'='*60}")
    print("EXECUTION SUMMARY")
    print(f"{'='*60}")
    print(f"Total tasks: {len(tasks)}")
    print(f"Successful: {successful}")
    print(f"Failed: {failed}")
    print(f"Duration: {duration}")
    print(f"{'='*60}")

if __name__ == "__main__":
    main()
