import type { Population } from './evolution';
import type { EvolutionConfig } from './evolution';
import { evaluatePopulation, DEFAULT_MATCH_CONFIG } from './selfPlay';
import { evolveGeneration, createPopulation, getPopulationStats } from './evolution';

/**
 * NEAT Training Worker
 * 
 * Runs training in a background thread to prevent UI blocking.
 * The main thread only handles visualization and UI updates.
 */

export interface TrainingWorkerMessage {
    type: 'start' | 'pause' | 'step' | 'reset' | 'init';
    config?: EvolutionConfig;
}

export interface TrainingWorkerResponse {
    type: 'update' | 'error' | 'ready';
    population?: Population;
    stats?: ReturnType<typeof getPopulationStats>;
    error?: string;
}

let population: Population | null = null;
let isRunning = false;
let config: EvolutionConfig | null = null;

/**
 * Handle messages from main thread
 */
self.onmessage = async (e: MessageEvent<TrainingWorkerMessage>) => {
    const message = e.data;
    
    try {
        switch (message.type) {
            case 'init':
                if (message.config) {
                    config = message.config;
                    population = createPopulation(config);
                    sendUpdate();
                    self.postMessage({ type: 'ready' } as TrainingWorkerResponse);
                }
                break;
                
            case 'start':
                isRunning = true;
                runTrainingLoop();
                break;
                
            case 'pause':
                isRunning = false;
                break;
                
            case 'step':
                if (population && config) {
                    const stats = await runSingleGeneration();
                    sendUpdate(stats);
                }
                break;
                
            case 'reset':
                if (config) {
                    population = createPopulation(config);
                    isRunning = false;
                    sendUpdate();
                }
                break;
        }
    } catch (error) {
        self.postMessage({
            type: 'error',
            error: error instanceof Error ? error.message : 'Unknown error',
        } as TrainingWorkerResponse);
    }
};

/**
 * Run continuous training loop
 */
async function runTrainingLoop() {
    while (isRunning && population && config) {
        const stats = await runSingleGeneration();
        sendUpdate(stats);
        
        // Yield to allow pause/stop messages to be processed
        await new Promise(resolve => setTimeout(resolve, 0));
    }
}

/**
 * Run a single generation
 */
async function runSingleGeneration(): Promise<ReturnType<typeof getPopulationStats> | null> {
    if (!population || !config) return null;
    
    console.log('[Worker] Starting generation', population.generation);
    
    // Evaluate population
    const evaluatedPop = evaluatePopulation(population, DEFAULT_MATCH_CONFIG);
    
    // Check fitness after evaluation
    const fitnesses = evaluatedPop.genomes.map(g => g.fitness);
    const avgFit = fitnesses.reduce((a, b) => a + b, 0) / fitnesses.length;
    const maxFit = Math.max(...fitnesses);
    console.log('[Worker] After evaluation - Avg fitness:', avgFit.toFixed(2), 'Max:', maxFit.toFixed(2));
    
    // Evolve to next generation
    population = evolveGeneration(evaluatedPop, config);
    
    console.log('[Worker] Generation', population.generation, 'complete');
    
    // IMPORTANT: Send stats from the EVALUATED population, not the evolved one
    // (evolved population has fitness reset to 0)
    return getPopulationStats(evaluatedPop);
}

/**
 * Send population update to main thread
 */
function sendUpdate(stats?: ReturnType<typeof getPopulationStats> | null) {
    if (!population) return;
    
    self.postMessage({
        type: 'update',
        population,
        stats: stats || undefined,
    } as TrainingWorkerResponse);
}