import { useState, useRef, useEffect, useCallback } from 'react';
import AppContainer from '../../components/AppContainer';
import { createArenaViewer, getArenaScene } from '../../lib/neatArena/arenaScene';
import { createSimulation, stepSimulation } from '../../lib/neatArena/simulation';
import { spinnerBotAction } from '../../lib/neatArena/baselineBots';
import { createPopulation, getPopulationStats, DEFAULT_EVOLUTION_CONFIG, type Population } from '../../lib/neatArena/evolution';
import { createNetwork } from '../../lib/neatArena/network';
import { generateObservation, observationToInputs } from '../../lib/neatArena/sensors';
import { exportGenome, downloadGenomeAsFile, uploadGenomeFromFile } from '../../lib/neatArena/exportImport';
import type { SimulationState, AgentAction, Genome } from '../../lib/neatArena/types';
import type { TrainingWorkerMessage, TrainingWorkerResponse } from '../../lib/neatArena/training.worker';
import FitnessGraph from './FitnessGraph';
import './NeatArena.css';

/**
 * NEAT Arena Miniapp
 * 
 * Trains AI agents using NEAT (NeuroEvolution of Augmenting Topologies) to play
 * a 2D top-down shooter arena via self-play.
 */
export default function NeatArena() {
    // Training state
    const [population, setPopulation] = useState<Population>(() => createPopulation(DEFAULT_EVOLUTION_CONFIG));
    const [isTraining, setIsTraining] = useState(false);
    const [showRays, setShowRays] = useState(false);
    const [mapSeed] = useState(12345);
    const [importedGenome, setImportedGenome] = useState<Genome | null>(null);
    const [fitnessHistory, setFitnessHistory] = useState<{ generation: number; best: number; avg: number }[]>([]);

    // Stats
    const stats = getPopulationStats(population);

    // Phaser game instance
    const phaserGameRef = useRef<Phaser.Game | null>(null);
    const phaserContainerRef = useRef<HTMLDivElement>(null);

    // Exhibition match state (visualizing champion)
    const simulationRef = useRef<SimulationState | null>(null);

    // Web Worker
    const workerRef = useRef<Worker | null>(null);

    // Initialize Web Worker
    useEffect(() => {
        const worker = new Worker(new URL('../../lib/neatArena/training.worker.ts', import.meta.url), {
            type: 'module'
        });

        worker.onmessage = (e: MessageEvent<TrainingWorkerResponse>) => {
            const response = e.data;

            switch (response.type) {
                case 'update':
                    if (response.population) {
                        setPopulation(response.population);
                        console.log('[UI] Stats?', response.stats ? 'YES' : 'NO', response.stats);

                        // Track fitness history for graph
                        if (response.stats) {
                            setFitnessHistory(prev => [...prev, {
                                generation: response.stats!.generation,
                                best: response.stats!.maxFitness,
                                avg: response.stats!.avgFitness,
                            }]);
                        }
                    }
                    break;

                case 'error':
                    console.error('Worker error:', response.error);
                    setIsTraining(false);
                    alert('Training error: ' + response.error);
                    break;

                case 'ready':
                    console.log('Worker ready');
                    break;
            }
        };

        // Initialize worker with config
        worker.postMessage({
            type: 'init',
            config: DEFAULT_EVOLUTION_CONFIG,
        } as TrainingWorkerMessage);

        workerRef.current = worker;

        return () => {
            worker.terminate();
            workerRef.current = null;
        };
    }, []);

    // Control worker based on training state
    useEffect(() => {
        if (!workerRef.current) return;

        if (isTraining) {
            workerRef.current.postMessage({
                type: 'start',
            } as TrainingWorkerMessage);
        } else {
            workerRef.current.postMessage({
                type: 'pause',
            } as TrainingWorkerMessage);
        }
    }, [isTraining]);

    // Initialize Phaser
    useEffect(() => {
        if (!phaserContainerRef.current) return;

        phaserContainerRef.current.innerHTML = '';
        const game = createArenaViewer(phaserContainerRef.current);
        phaserGameRef.current = game;

        simulationRef.current = createSimulation(mapSeed, 0);

        return () => {
            game.destroy(true);
            phaserGameRef.current = null;
        };
    }, [mapSeed]);

    // Exhibition match loop (visualizing best vs second-best AI)
    useEffect(() => {
        if (!phaserGameRef.current) return;

        const interval = setInterval(() => {
            if (!simulationRef.current) return;

            const sim = simulationRef.current;

            if (sim.isOver) {
                simulationRef.current = createSimulation(mapSeed, 0);
                return;
            }

            // Get best and second-best genomes
            const sortedGenomes = [...population.genomes].sort((a, b) => b.fitness - a.fitness);
            const genome0 = importedGenome || sortedGenomes[0] || null;
            const genome1 = sortedGenomes.length > 1 ? sortedGenomes[1] : null;

            // Agent 0: Best AI
            let action0: AgentAction;
            if (genome0) {
                const network = createNetwork(genome0);
                const obs = generateObservation(0, sim);
                const inputs = observationToInputs(obs);
                const outputs = network.activate(inputs);

                action0 = {
                    moveX: outputs[0],
                    moveY: outputs[1],
                    turn: outputs[2],
                    shoot: outputs[3],
                };
            } else {
                action0 = spinnerBotAction();
            }

            // Agent 1: Second-best AI (or spinner if not enough genomes)
            let action1: AgentAction;
            if (genome1) {
                const network = createNetwork(genome1);
                const obs = generateObservation(1, sim);
                const inputs = observationToInputs(obs);
                const outputs = network.activate(inputs);

                action1 = {
                    moveX: outputs[0],
                    moveY: outputs[1],
                    turn: outputs[2],
                    shoot: outputs[3],
                };
            } else {
                action1 = spinnerBotAction();
            }

            simulationRef.current = stepSimulation(sim, [action0, action1]);

            if (phaserGameRef.current) {
                const scene = getArenaScene(phaserGameRef.current);
                scene.updateSimulation(simulationRef.current);
                scene.setShowRays(showRays);
            }

        }, 1000 / 30);

        return () => clearInterval(interval);
    }, [showRays, mapSeed, population.genomes, importedGenome]);

    const handleReset = useCallback(() => {
        setIsTraining(false);
        setImportedGenome(null);
        setFitnessHistory([]);

        if (workerRef.current) {
            workerRef.current.postMessage({
                type: 'reset',
            } as TrainingWorkerMessage);
        }

        simulationRef.current = createSimulation(mapSeed, 0);

        if (phaserGameRef.current) {
            const scene = getArenaScene(phaserGameRef.current);
            scene.updateSimulation(simulationRef.current);
        }
    }, [mapSeed]);

    const handleStepGeneration = useCallback(() => {
        if (workerRef.current) {
            workerRef.current.postMessage({
                type: 'step',
            } as TrainingWorkerMessage);
        }
    }, []);

    const handleExport = useCallback(() => {
        if (!population.bestGenomeEver) {
            alert('No champion to export yet!');
            return;
        }

        const exported = exportGenome(
            population.bestGenomeEver,
            DEFAULT_EVOLUTION_CONFIG,
            {
                generation: stats.generation,
                fitness: stats.bestFitnessEver,
                speciesCount: stats.speciesCount,
            }
        );

        downloadGenomeAsFile(exported, `neat-champion-gen${stats.generation}.json`);
    }, [population.bestGenomeEver, stats]);

    const handleImport = useCallback(async () => {
        try {
            const exported = await uploadGenomeFromFile();
            setImportedGenome(exported.genome);
            alert(`Imported champion from generation ${exported.metadata?.generation || '?'} with fitness ${exported.metadata?.fitness?.toFixed(1) || '?'}`);
        } catch (err) {
            alert('Failed to import genome: ' + (err as Error).message);
        }
    }, []);

    return (
        <AppContainer title="NEAT Arena">
            <div className="neat-arena-layout">
                {/* Left Panel: Controls */}
                <div className="controls-panel">
                    <section className="control-section">
                        <h3>Training Controls</h3>
                        <div className="control-group">
                            <button
                                className={`btn-primary ${isTraining ? 'btn-stop' : 'btn-start'}`}
                                onClick={() => setIsTraining(!isTraining)}
                            >
                                {isTraining ? '⏸ Pause Training' : '▶ Start Training'}
                            </button>
                            <button
                                className="btn-secondary"
                                onClick={handleStepGeneration}
                                disabled={isTraining}
                            >
                                ⏭ Step Generation
                            </button>
                            <button
                                className="btn-secondary"
                                onClick={handleReset}
                                disabled={isTraining}
                            >
                                🔄 Reset
                            </button>
                        </div>
                        <p className="info-text">
                            {isTraining
                                ? '🟢 Training in background worker...'
                                : importedGenome
                                    ? '🎮 Watching imported champion vs Gen best'
                                    : population.genomes.length > 1
                                        ? `🎮 Watching Gen ${stats.generation}: Best vs 2nd-Best AI`
                                        : '⚪ Need at least 2 genomes for exhibition'}
                        </p>
                    </section>

                    <section className="control-section">
                        <h3>Evolution Stats</h3>
                        <div className="stats-grid">
                            <div className="stat-item">
                                <span className="stat-label">Generation</span>
                                <span className="stat-value">{stats.generation}</span>
                            </div>
                            <div className="stat-item">
                                <span className="stat-label">Species</span>
                                <span className="stat-value">{stats.speciesCount}</span>
                            </div>
                            <div className="stat-item">
                                <span className="stat-label">Best Fitness</span>
                                <span className="stat-value">{stats.maxFitness.toFixed(1)}</span>
                            </div>
                            <div className="stat-item">
                                <span className="stat-label">Avg Fitness</span>
                                <span className="stat-value">{stats.avgFitness.toFixed(1)}</span>
                            </div>
                            <div className="stat-item">
                                <span className="stat-label">Champion</span>
                                <span className="stat-value">{stats.bestFitnessEver.toFixed(1)}</span>
                            </div>
                            <div className="stat-item">
                                <span className="stat-label">Innovations</span>
                                <span className="stat-value">{stats.totalInnovations}</span>
                            </div>
                        </div>
                    </section>

                    {fitnessHistory.length > 0 && (
                        <section className="control-section">
                            <h3>Fitness Progress</h3>
                            <FitnessGraph history={fitnessHistory} />
                        </section>
                    )}

                    <section className="control-section">
                        <h3>Debug Options</h3>
                        <label className="checkbox-label">
                            <input
                                type="checkbox"
                                checked={showRays}
                                onChange={(e) => setShowRays(e.target.checked)}
                            />
                            <span>Show Ray Sensors</span>
                        </label>
                    </section>

                    <section className="control-section">
                        <h3>Export / Import</h3>
                        <div className="control-group">
                            <button
                                className="btn-secondary"
                                onClick={handleExport}
                                disabled={!population.bestGenomeEver}
                            >
                                💾 Export Champion
                            </button>
                            <button
                                className="btn-secondary"
                                onClick={handleImport}
                            >
                                📂 Import Genome
                            </button>
                        </div>
                        {importedGenome && (
                            <p className="info-text">
                                ✅ Imported genome loaded
                            </p>
                        )}
                    </section>

                    <section className="info-section">
                        <h4>NEAT Arena Status</h4>
                        <ul>
                            <li>✅ Deterministic 30Hz simulation</li>
                            <li>✅ Symmetric procedural maps</li>
                            <li>✅ Agent physics & bullets</li>
                            <li>✅ 360° ray sensors (53 inputs)</li>
                            <li>✅ NEAT evolution with speciation</li>
                            <li>✅ Self-play training (K=4 matches)</li>
                            <li>✅ Export/import genomes</li>
                            <li>✅ Web worker (no UI lag!)</li>
                        </ul>
                    </section>
                </div>

                {/* Right Panel: Phaser Viewer */}
                <div className="viewer-panel">
                    <div
                        ref={phaserContainerRef}
                        className="phaser-container"
                    />
                </div>
            </div>
        </AppContainer>
    );
}
