import { useState, useEffect } from 'react';
import { useAuth } from '../App';
import { Upload, Database, Save, Activity, CheckCircle, AlertTriangle, RefreshCw, Lightbulb } from 'lucide-react';

const API_URL = import.meta.env.VITE_API_URL || 'http://localhost:3001';

interface TrainingData {
    id: string;
    audioUrl: string;
    transcription: string; // generated by Whisper
    manualCorrection?: string;
    rawWER?: number;
    normalizedWER?: number;
    status: string;
    createdAt: string;
}

export default function TrainingLab() {
    const { apiKey, logout } = useAuth();
    const [mode, setMode] = useState<'db' | 'upload' | 'suggestions'>('db');
    const [audios, setAudios] = useState<TrainingData[]>([]);
    const [selectedAudio, setSelectedAudio] = useState<TrainingData | null>(null);
    const [manualCorrection, setManualCorrection] = useState('');
    const [loading, setLoading] = useState(false);
    const [submitting, setSubmitting] = useState(false);
    const [result, setResult] = useState<{ rawWER: number, normalizedWER: number, missingWords: string[] } | null>(null);

    const [suggestions, setSuggestions] = useState<{ original: string, replacement: string, count: number }[]>([]);
    const [selectedSuggestions, setSelectedSuggestions] = useState<Set<string>>(new Set());
    const [recalculating, setRecalculating] = useState(false);
    const [recalcResult, setRecalcResult] = useState<{ processed: number, avgRawWER: number, avgNormalizedWER: number, improvementPercent: number } | null>(null);

    const fetchAudios = async () => {
        setLoading(true);
        try {
            const res = await fetch(`${API_URL}/v1/admin/training/audios`, {
                headers: { 'Authorization': `Bearer ${apiKey}` }
            });
            if (res.status === 401) return logout();
            const data = await res.json();
            setAudios(data);
        } catch (err) {
            console.error(err);
        } finally {
            setLoading(false);
        }
    };

    useEffect(() => {
        if (mode === 'db') {
            fetchAudios();
        }
    }, [mode, apiKey, logout]);

    const handleSubmit = async () => {
        if (!selectedAudio || !manualCorrection.trim()) return;
        setSubmitting(true);
        setResult(null);
        try {
            const res = await fetch(`${API_URL}/v1/admin/training/submit`, {
                method: 'POST',
                headers: {
                    'Authorization': `Bearer ${apiKey}`,
                    'Content-Type': 'application/json'
                },
                body: JSON.stringify({
                    id: selectedAudio.id,
                    audioUrl: selectedAudio.audioUrl,
                    transcription: selectedAudio.transcription,
                    manualCorrection
                })
            });
            if (res.status === 401) return logout();
            const json = await res.json();
            if (json.error) {
                alert('Erreur: ' + JSON.stringify(json.error));
            } else {
                setResult({
                    rawWER: json.rawWER,
                    normalizedWER: json.normalizedWER,
                    missingWords: json.missingWords
                });
                // Remove from list
                setAudios(prev => prev.filter(a => a.id !== selectedAudio.id));
            }
        } catch (err) {
            console.error(err);
            alert('Erreur serveur.');
        } finally {
            setSubmitting(false);
        }
    };

    const fetchSuggestions = async () => {
        setLoading(true);
        try {
            const res = await fetch(`${API_URL}/v1/admin/training/suggestions`, {
                headers: { 'Authorization': `Bearer ${apiKey}` }
            });
            if (res.status === 401) return logout();
            const data = await res.json();
            setSuggestions(data);
            setSelectedSuggestions(new Set(data.map((d: any) => `${d.original}->${d.replacement}`)));
        } catch (err) {
            console.error(err);
        } finally {
            setLoading(false);
        }
    };

    const applySuggestions = async () => {
        const payload = suggestions.filter(s => selectedSuggestions.has(`${s.original}->${s.replacement}`));
        if (payload.length === 0) return;
        setSubmitting(true);
        try {
            const res = await fetch(`${API_URL}/v1/admin/training/apply-suggestions`, {
                method: 'POST',
                headers: { 'Authorization': `Bearer ${apiKey}`, 'Content-Type': 'application/json' },
                body: JSON.stringify({ suggestions: payload })
            });
            if (res.status === 401) return logout();
            const json = await res.json();
            alert(`Succès! ${json.injectedCount} règles ont été injectées dans le dictionnaire.`);
            fetchSuggestions();
        } catch (err) {
            console.error(err);
        } finally {
            setSubmitting(false);
        }
    };

    const recalculateWER = async () => {
        setRecalculating(true);
        try {
            const res = await fetch(`${API_URL}/v1/admin/training/recalculate-wer`, {
                method: 'POST',
                headers: { 'Authorization': `Bearer ${apiKey}` }
            });
            if (res.status === 401) return logout();
            const json = await res.json();
            setRecalcResult(json);
        } catch (err) {
            console.error(err);
        } finally {
            setRecalculating(false);
        }
    };

    return (
        <div className="p-8 max-w-5xl mx-auto">
            <div className="flex items-center gap-3 mb-8">
                <Activity className="w-8 h-8 text-purple-600" />
                <h1 className="text-3xl font-bold text-slate-800">Training Lab (WER)</h1>
            </div>

            <div className="bg-white p-2 rounded-xl border border-slate-200 inline-flex mb-8 shadow-sm">
                <button
                    onClick={() => setMode('db')}
                    className={`flex items-center gap-2 px-6 py-2.5 rounded-lg text-sm font-medium transition ${mode === 'db' ? 'bg-slate-900 text-white shadow' : 'text-slate-600 hover:bg-slate-50'}`}
                >
                    <Database className="w-4 h-4" /> Audios de la BDD
                </button>
                <button
                    onClick={() => setMode('upload')}
                    className={`flex items-center gap-2 px-6 py-2.5 rounded-lg text-sm font-medium transition ${mode === 'upload' ? 'bg-slate-900 text-white shadow' : 'text-slate-600 hover:bg-slate-50'}`}
                >
                    <Upload className="w-4 h-4" /> Upload Manuel
                </button>
                <button
                    onClick={() => { setMode('suggestions'); fetchSuggestions(); }}
                    className={`flex items-center gap-2 px-6 py-2.5 rounded-lg text-sm font-medium transition ${mode === 'suggestions' ? 'bg-slate-900 text-white shadow' : 'text-slate-600 hover:bg-slate-50'}`}
                >
                    <Lightbulb className="w-4 h-4" /> Suggestions Auto-Normalisation
                </button>
            </div>

            <div className="grid grid-cols-1 lg:grid-cols-3 gap-8">
                {mode === 'suggestions' && (
                    <div className="lg:col-span-3 space-y-6">
                        <div className="bg-white rounded-2xl border border-slate-200 shadow-sm p-8">
                            <div className="flex items-center justify-between mb-6">
                                <div>
                                    <h2 className="text-xl font-bold text-slate-800 flex items-center gap-2">
                                        <Lightbulb className="w-5 h-5 text-amber-500" />
                                        Auto-Normalisation (Top 20)
                                    </h2>
                                    <p className="text-sm text-slate-500 mt-1">Ces mots ont été fréquemment corrigés manuellement. Validez-les pour les injecter dans le dictionnaire Wolof.</p>
                                </div>
                                <div className="flex gap-3">
                                    <button
                                        onClick={recalculateWER}
                                        disabled={recalculating}
                                        className="flex items-center gap-2 px-4 py-2 bg-slate-100 hover:bg-slate-200 text-slate-700 font-medium rounded-lg transition disabled:opacity-50"
                                    >
                                        <RefreshCw className={`w-4 h-4 ${recalculating ? 'animate-spin' : ''}`} />
                                        Recalculer WER Global
                                    </button>
                                    <button
                                        onClick={applySuggestions}
                                        disabled={submitting || selectedSuggestions.size === 0}
                                        className="flex items-center gap-2 px-4 py-2 bg-indigo-600 hover:bg-indigo-700 text-white font-medium rounded-lg transition disabled:opacity-50"
                                    >
                                        <Save className="w-4 h-4" />
                                        Injecter ({selectedSuggestions.size}) Règles
                                    </button>
                                </div>
                            </div>

                            {recalcResult && (
                                <div className="mb-8 p-4 bg-emerald-50 border border-emerald-100 rounded-xl flex items-start gap-4">
                                    <CheckCircle className="w-6 h-6 text-emerald-500 shrink-0" />
                                    <div>
                                        <h3 className="font-bold text-emerald-800 mb-1">Benchmark Terminé ({recalcResult.processed} audios)</h3>
                                        <div className="flex gap-6 mt-2">
                                            <div>
                                                <p className="text-xs text-emerald-600 uppercase">WER Brut Moy</p>
                                                <p className="text-xl font-bold text-emerald-900">{Math.round(recalcResult.avgRawWER * 100)}%</p>
                                            </div>
                                            <div>
                                                <p className="text-xs text-emerald-600 uppercase">WER Normalisé Moy</p>
                                                <p className="text-xl font-bold text-emerald-900">{Math.round(recalcResult.avgNormalizedWER * 100)}%</p>
                                            </div>
                                            <div>
                                                <p className="text-xs text-emerald-600 uppercase">Gain de Précision</p>
                                                <p className="text-xl font-bold text-emerald-900">+{recalcResult.improvementPercent.toFixed(2)}%</p>
                                            </div>
                                        </div>
                                    </div>
                                </div>
                            )}

                            {loading ? <p className="text-slate-500 py-10 text-center flex justify-center"><RefreshCw className="animate-spin text-indigo-500" /></p> : suggestions.length === 0 ? (
                                <div className="text-center py-12 px-4 bg-slate-50 rounded-xl border border-dashed border-slate-200">
                                    <CheckCircle className="w-12 h-12 text-slate-300 mx-auto mb-3" />
                                    <p className="text-slate-500">Aucune nouvelle suggestion détectée.</p>
                                </div>
                            ) : (
                                <div className="overflow-x-auto">
                                    <table className="w-full text-left text-sm text-slate-600">
                                        <thead className="text-xs text-slate-500 uppercase bg-slate-50">
                                            <tr>
                                                <th className="px-6 py-3 rounded-tl-xl"><input type="checkbox" checked={selectedSuggestions.size === suggestions.length} onChange={(e) => setSelectedSuggestions(e.target.checked ? new Set(suggestions.map(s => `${s.original}->${s.replacement}`)) : new Set())} /></th>
                                                <th className="px-6 py-3">Erreur (Whisper)</th>
                                                <th className="px-6 py-3">Correction (Humain)</th>
                                                <th className="px-6 py-3 rounded-tr-xl">Fréquence</th>
                                            </tr>
                                        </thead>
                                        <tbody>
                                            {suggestions.map(s => {
                                                const key = `${s.original}->${s.replacement}`;
                                                return (
                                                    <tr key={key} className="border-b border-slate-100 last:border-0 hover:bg-slate-50">
                                                        <td className="px-6 py-4">
                                                            <input type="checkbox" checked={selectedSuggestions.has(key)} onChange={(e) => {
                                                                const newSet = new Set(selectedSuggestions);
                                                                if (e.target.checked) newSet.add(key); else newSet.delete(key);
                                                                setSelectedSuggestions(newSet);
                                                            }} />
                                                        </td>
                                                        <td className="px-6 py-4 font-mono text-orange-600">{s.original}</td>
                                                        <td className="px-6 py-4 font-mono text-emerald-600 font-bold">{s.replacement}</td>
                                                        <td className="px-6 py-4">
                                                            <span className="bg-indigo-50 text-indigo-700 px-2.5 py-1 rounded-full text-xs font-semibold">{s.count} fois</span>
                                                        </td>
                                                    </tr>
                                                );
                                            })}
                                        </tbody>
                                    </table>
                                </div>
                            )}
                        </div>
                    </div>
                )}

                {/* Left Sidebar: List */}
                <div className={`bg-white rounded-2xl border border-slate-200 overflow-hidden shadow-sm h-[600px] flex flex-col ${mode === 'suggestions' ? 'hidden' : ''}`}>
                    <div className="p-4 border-b border-slate-100 bg-slate-50 flex items-center justify-between">
                        <h2 className="font-semibold text-slate-800 flex items-center gap-2">
                            <Database className="w-4 h-4 text-slate-400" />
                            File d'attente ({audios.length})
                        </h2>
                        <button onClick={fetchAudios} disabled={loading} className="text-slate-400 hover:text-slate-600">
                            <RefreshCw className={`w-4 h-4 ${loading ? 'animate-spin' : ''}`} />
                        </button>
                    </div>

                    <div className="flex-1 overflow-auto p-2 space-y-2">
                        {loading && <p className="text-center text-slate-400 py-6 text-sm">Chargement...</p>}
                        {!loading && audios.length === 0 && (
                            <div className="text-center py-10 px-4">
                                <CheckCircle className="w-12 h-12 text-emerald-400 mx-auto mb-3 opacity-50" />
                                <p className="text-sm text-slate-500">Aucun audio en attente de révision.</p>
                            </div>
                        )}
                        {audios.map(audio => (
                            <button
                                key={audio.id}
                                onClick={() => { setSelectedAudio(audio); setManualCorrection(''); setResult(null); }}
                                className={`w-full text-left p-4 rounded-xl transition border text-sm ${selectedAudio?.id === audio.id ? 'bg-indigo-50 border-indigo-200 ring-1 ring-indigo-200' : 'bg-white border-transparent hover:border-slate-200 hover:bg-slate-50'}`}
                            >
                                <p className="font-medium text-slate-800 truncate mb-1">{audio.transcription || 'Sans transcription'}</p>
                                <p className="text-xs text-slate-400 truncate">{new Date(audio.createdAt).toLocaleString()}</p>
                            </button>
                        ))}
                    </div>
                </div>

                {/* Right Area: Editor */}
                <div className={`lg:col-span-2 space-y-6 ${mode === 'suggestions' ? 'hidden' : ''}`}>
                    {mode === 'upload' && !selectedAudio && (
                        <div className="bg-white rounded-2xl border border-slate-200 shadow-sm p-10 text-center">
                            <Upload className="w-12 h-12 text-slate-300 mx-auto mb-4" />
                            <h3 className="text-lg font-medium text-slate-800 mb-2">Upload Manuel (Bientôt disponible)</h3>
                            <p className="text-sm text-slate-500 mb-6">Uploadez un fichier .wav/.mp3 pour le transcrire et l'ajouter au dataset d'entraînement local.</p>
                            <input type="file" className="block w-full max-w-sm mx-auto text-sm text-slate-500 file:mr-4 file:py-2 file:px-4 file:rounded-full file:border-0 file:text-sm file:font-semibold file:bg-indigo-50 file:text-indigo-700 hover:file:bg-indigo-100" />
                        </div>
                    )}

                    {selectedAudio && (
                        <div className="bg-white rounded-2xl border border-slate-200 shadow-sm p-6">
                            <div className="mb-6">
                                <h3 className="text-sm font-semibold text-slate-500 uppercase tracking-wider mb-3">Audio Source</h3>
                                <audio src={selectedAudio.audioUrl} controls className="w-full h-12 outline-none" />
                            </div>

                            <div className="mb-6">
                                <h3 className="text-sm font-semibold text-slate-500 uppercase tracking-wider mb-3">Whisper Genération v1</h3>
                                <div className="p-4 bg-slate-50 rounded-xl border border-slate-100 text-slate-800 text-sm leading-relaxed">
                                    {selectedAudio.transcription}
                                </div>
                            </div>

                            <div className="mb-6">
                                <h3 className="text-sm font-semibold text-slate-500 uppercase tracking-wider mb-3 flex items-center justify-between">
                                    Vérité Terrain (Ground Truth)
                                    <span className="text-xs font-normal text-indigo-500 bg-indigo-50 px-2 py-0.5 rounded-full">Wolof Standardisé</span>
                                </h3>
                                <textarea
                                    value={manualCorrection}
                                    onChange={e => setManualCorrection(e.target.value)}
                                    placeholder="Écrivez la transcription manuelle parfaite ici..."
                                    className="w-full min-h-[120px] p-4 bg-white border border-slate-200 rounded-xl text-sm leading-relaxed outline-none focus:ring-2 focus:ring-indigo-500 focus:border-indigo-500 resize-y"
                                />
                            </div>

                            <div className="flex items-center justify-end">
                                <button
                                    onClick={handleSubmit}
                                    disabled={submitting || !manualCorrection.trim()}
                                    className="flex items-center gap-2 px-6 py-3 bg-indigo-600 hover:bg-indigo-700 text-white font-medium rounded-xl transition disabled:opacity-50 disabled:cursor-not-allowed shadow-sm"
                                >
                                    {submitting ? <RefreshCw className="w-4 h-4 animate-spin" /> : <Save className="w-4 h-4" />}
                                    Enregistrer & Recalculer WER
                                </button>
                            </div>
                        </div>
                    )}

                    {result && (
                        <div className="bg-gradient-to-br from-indigo-900 to-slate-900 rounded-2xl p-6 text-white shadow-lg overflow-hidden relative">
                            <div className="absolute top-0 right-0 p-8 opacity-5">
                                <Activity className="w-48 h-48" />
                            </div>

                            <h3 className="text-lg font-bold mb-6 flex items-center gap-2">
                                <CheckCircle className="w-5 h-5 text-emerald-400" />
                                Entraînement enregistré !
                            </h3>

                            <div className="grid grid-cols-2 gap-4 mb-6 relative z-10">
                                <div className="bg-white/10 rounded-xl p-4 border border-white/5">
                                    <p className="text-xs font-medium text-slate-400 uppercase tracking-wider mb-1">WER Brut (Whisper)</p>
                                    <p className="text-3xl font-bold">{Math.round(result.rawWER * 100)}%</p>
                                </div>
                                <div className="bg-white/10 rounded-xl p-4 border border-white/5">
                                    <p className="text-xs font-medium text-slate-400 uppercase tracking-wider mb-1">WER Normalisé (Dico)</p>
                                    <p className="text-3xl font-bold text-emerald-400">{Math.round(result.normalizedWER * 100)}%</p>
                                </div>
                            </div>

                            {result.missingWords && result.missingWords.length > 0 && (
                                <div className="bg-orange-500/10 border border-orange-500/20 rounded-xl p-4 relative z-10">
                                    <div className="flex items-start gap-3">
                                        <AlertTriangle className="w-5 h-5 text-orange-400 shrink-0 mt-0.5" />
                                        <div>
                                            <p className="text-sm font-medium text-orange-300 mb-2">Mots absents du dictionnaire Wolof :</p>
                                            <div className="flex flex-wrap gap-2">
                                                {result.missingWords.map(w => (
                                                    <span key={w} className="px-2 py-1 bg-orange-500/20 text-orange-200 text-xs rounded-md font-mono">{w}</span>
                                                ))}
                                            </div>
                                            <p className="text-xs text-orange-400/70 mt-3">Suggérez d'ajouter ces mots dans `normalizeWolof.ts` pour améliorer le taux de reconnaissance global de la plateforme.</p>
                                        </div>
                                    </div>
                                </div>
                            )}
                        </div>
                    )}
                </div>
            </div>
        </div>
    );
}