kaboot/hooks/useGame.ts

import { useState, useEffect, useRef, useCallback } from 'react';
import { Quiz, Player, GameState, GameRole, NetworkMessage, AnswerOption, Question, GenerateQuizOptions, ProcessedDocument, GameConfig, DEFAULT_GAME_CONFIG } from '../types';
import { generateQuiz } from '../services/geminiService';
import { QUESTION_TIME, QUESTION_TIME_MS, PLAYER_COLORS, calculatePoints, getPlayerRank } from '../constants';
import { Peer, DataConnection } from 'peerjs';

const BACKEND_URL = import.meta.env.VITE_BACKEND_URL || 'http://localhost:3001';

export const useGame = () => {
  const [role, setRole] = useState<GameRole>('HOST');
  const [gameState, setGameState] = useState<GameState>('LANDING');
  const [quiz, setQuiz] = useState<Quiz | null>(null);
  const [currentQuestionIndex, setCurrentQuestionIndex] = useState(0);
  const [players, setPlayers] = useState<Player[]>([]);
  const [timeLeft, setTimeLeft] = useState(0);
  const [error, setError] = useState<string | null>(null);
  const [hasAnswered, setHasAnswered] = useState(false);
  const [gamePin, setGamePin] = useState<string | null>(null);
  const [currentCorrectShape, setCurrentCorrectShape] = useState<string | null>(null);
  const [lastPointsEarned, setLastPointsEarned] = useState<number | null>(null);
  const [selectedOption, setSelectedOption] = useState<AnswerOption | null>(null);
  const [currentPlayerScore, setCurrentPlayerScore] = useState(0);
  const [currentStreak, setCurrentStreak] = useState(0);
  const [currentPlayerId, setCurrentPlayerId] = useState<string | null>(null);
  const [currentPlayerName, setCurrentPlayerName] = useState<string | null>(null);
  const [pendingQuizToSave, setPendingQuizToSave] = useState<{ quiz: Quiz; topic: string } | null>(null);
  const [sourceQuizId, setSourceQuizId] = useState<string | null>(null);
  const [gameConfig, setGameConfig] = useState<GameConfig>(DEFAULT_GAME_CONFIG);
  const [firstCorrectPlayerId, setFirstCorrectPlayerId] = useState<string | null>(null);
  
  const timerRef = useRef<ReturnType<typeof setInterval> | null>(null);
  const peerRef = useRef<Peer | null>(null);
  const connectionsRef = useRef<Map<string, DataConnection>>(new Map());
  const hostConnectionRef = useRef<DataConnection | null>(null);
  
  // Refs for callbacks/async functions to access latest state
  const timeLeftRef = useRef(0);
  const playersRef = useRef<Player[]>([]);
  const currentQuestionIndexRef = useRef(0);
  const quizRef = useRef<Quiz | null>(null);
  const gameConfigRef = useRef<GameConfig>(DEFAULT_GAME_CONFIG);

  useEffect(() => { timeLeftRef.current = timeLeft; }, [timeLeft]);
  useEffect(() => { playersRef.current = players; }, [players]);
  useEffect(() => { currentQuestionIndexRef.current = currentQuestionIndex; }, [currentQuestionIndex]);
  useEffect(() => { quizRef.current = quiz; }, [quiz]);
  useEffect(() => { gameConfigRef.current = gameConfig; }, [gameConfig]);

  const generateGamePin = () => Math.floor(Math.random() * 900000) + 100000 + "";

  // -- HOST LOGIC --

  const uploadDocument = async (file: File, useOcr: boolean = false): Promise<ProcessedDocument> => {
    const formData = new FormData();
    formData.append('document', file);
    formData.append('useOcr', String(useOcr));
    
    const response = await fetch(`${BACKEND_URL}/api/upload`, {
      method: 'POST',
      body: formData
    });
    
    if (!response.ok) {
      const error = await response.json();
      throw new Error(error.error || 'Failed to upload document');
    }
    
    return response.json();
  };

  const startQuizGen = async (options: { topic?: string; questionCount?: number; files?: File[]; useOcr?: boolean }) => {
    try {
      setGameState('GENERATING');
      setError(null);
      setRole('HOST');
      
      let documents: ProcessedDocument[] | undefined;
      if (options.files && options.files.length > 0) {
        documents = await Promise.all(
          options.files.map(file => uploadDocument(file, options.useOcr))
        );
      }
      
      const generateOptions: GenerateQuizOptions = {
        topic: options.topic,
        questionCount: options.questionCount,
        documents
      };
      
      const generatedQuiz = await generateQuiz(generateOptions);
      const saveLabel = options.topic || options.files?.map(f => f.name).join(', ') || '';
      setPendingQuizToSave({ quiz: generatedQuiz, topic: saveLabel });
      setQuiz(generatedQuiz);
      setGameState('EDITING');
    } catch (e) {
      const message = e instanceof Error ? e.message : "Failed to generate quiz.";
      setError(message);
      setGameState('LANDING');
    }
  };

  const dismissSavePrompt = () => {
    setPendingQuizToSave(null);
  };

  const startManualCreation = () => {
    setRole('HOST');
    setGameState('CREATING');
  };

  const finalizeManualQuiz = (manualQuiz: Quiz, saveToLibrary: boolean = false) => {
    if (saveToLibrary) {
      setPendingQuizToSave({ quiz: manualQuiz, topic: '' });
    }
    initializeHostGame(manualQuiz);
  };

  const loadSavedQuiz = (savedQuiz: Quiz, quizId?: string) => {
    setRole('HOST');
    setQuiz(savedQuiz);
    setSourceQuizId(quizId || null);
    setGameState('EDITING');
  };
  
  const updateQuizFromEditor = (updatedQuiz: Quiz) => {
    setQuiz(updatedQuiz);
    setPendingQuizToSave(prev => prev ? { ...prev, quiz: updatedQuiz } : { quiz: updatedQuiz, topic: '' });
  };
  
  const startGameFromEditor = (finalQuiz: Quiz, config: GameConfig) => {
    setQuiz(finalQuiz);
    setGameConfig(config);
    initializeHostGame(finalQuiz, config.hostParticipates);
  };
  
  const backFromEditor = () => {
    setQuiz(null);
    setPendingQuizToSave(null);
    setSourceQuizId(null);
    setGameState('LANDING');
  };

  // We use a ref to hold the current handleHostData function
  // This prevents stale closures in the PeerJS event listeners
  const handleHostDataRef = useRef<(conn: DataConnection, data: NetworkMessage) => void>(() => {});

  const initializeHostGame = (newQuiz: Quiz, hostParticipates: boolean = true) => {
    setQuiz(newQuiz);
    const pin = generateGamePin();
    setGamePin(pin);
    
    const peer = new Peer(`kaboot-${pin}`);
    peerRef.current = peer;

    peer.on('open', (id) => {
      if (hostParticipates) {
        const hostPlayer: Player = {
          id: 'host',
          name: 'Host',
          score: 0,
          streak: 0,
          lastAnswerCorrect: null,
          isBot: false,
          avatarSeed: Math.random(),
          color: PLAYER_COLORS[0]
        };
        setPlayers([hostPlayer]);
        setCurrentPlayerId('host');
        setCurrentPlayerName('Host');
      } else {
        setPlayers([]);
        setCurrentPlayerId(null);
        setCurrentPlayerName(null);
      }
      setGameState('LOBBY');
    });

    peer.on('connection', (conn) => {
      conn.on('data', (data: any) => {
          // Delegate to the ref to ensure we always use the latest function closure
          handleHostDataRef.current(conn, data);
      });
    });
    
    peer.on('error', () => {
      setError("Network error. Try a different topic or reload.");
      setGameState('LANDING');
    });
  };

  const handleHostData = (conn: DataConnection, data: NetworkMessage) => {
    if (data.type === 'JOIN') {
      setPlayers(prev => {
        if (prev.find(p => p.id === conn.peer)) return prev;
        const colorIndex = prev.length % PLAYER_COLORS.length;
        const newPlayer: Player = {
          id: conn.peer,
          name: data.payload.name,
          score: 0,
          streak: 0,
          lastAnswerCorrect: null,
          isBot: false,
          avatarSeed: Math.random(),
          color: PLAYER_COLORS[colorIndex]
        };
        return [...prev, newPlayer];
      });
      connectionsRef.current.set(conn.peer, conn);
      conn.send({ type: 'WELCOME', payload: { playerId: conn.peer, quizTitle: 'Kaboot', players: [] } });
    }
    
    if (data.type === 'ANSWER') {
      const { playerId, isCorrect } = data.payload;
      
      const currentPlayer = playersRef.current.find(p => p.id === playerId);
      
      if (!currentPlayer || currentPlayer.lastAnswerCorrect !== null) return;

      const isFirstCorrect = isCorrect && firstCorrectPlayerId === null;
      if (isFirstCorrect) {
        setFirstCorrectPlayerId(playerId);
      }

      const newStreak = isCorrect ? currentPlayer.streak + 1 : 0;
      const playerRank = getPlayerRank(playerId, playersRef.current);

      const points = calculatePoints({
        isCorrect,
        timeLeftMs: timeLeftRef.current,
        questionTimeMs: QUESTION_TIME_MS,
        streak: newStreak,
        playerRank,
        isFirstCorrect,
        config: gameConfigRef.current,
      });
      const newScore = Math.max(0, currentPlayer.score + points);

      setPlayers(prev => prev.map(p => {
        if (p.id !== playerId) return p;
        return { ...p, score: newScore, streak: newStreak, lastAnswerCorrect: isCorrect };
      }));

      conn.send({ type: 'RESULT', payload: { isCorrect, scoreAdded: points, newScore } });
    }
  };

  // Update the ref whenever handleHostData changes (which happens on render)
  useEffect(() => {
      handleHostDataRef.current = handleHostData;
  });

  const broadcast = (msg: NetworkMessage) => {
    connectionsRef.current.forEach(conn => { if (conn.open) conn.send(msg); });
  };

  const startHostGame = () => {
    setCurrentQuestionIndex(0);
    broadcast({ type: 'GAME_START', payload: {} });
    startCountdown();
  };

  const startCountdown = () => {
    setGameState('COUNTDOWN');
    broadcast({ type: 'START_COUNTDOWN', payload: { duration: 3 } });
    
    let count = 3;
    setTimeLeft(count);
    
    if (timerRef.current) clearInterval(timerRef.current);
    timerRef.current = setInterval(() => {
      count--;
      setTimeLeft(count);
      if (count <= 0) { 
          if (timerRef.current) clearInterval(timerRef.current); 
          startQuestion(); 
      }
    }, 1000);
  };

  const startQuestion = () => {
    setGameState('QUESTION');
    setHasAnswered(false);
    setLastPointsEarned(null);
    setSelectedOption(null);
    setTimeLeft(QUESTION_TIME_MS);
    setFirstCorrectPlayerId(null);
    setPlayers(prev => prev.map(p => ({ ...p, lastAnswerCorrect: null })));

    const currentQuiz = quizRef.current;
    const currentIndex = currentQuestionIndexRef.current;

    if (currentQuiz) {
        const currentQ = currentQuiz.questions[currentIndex];
        const options = currentQ.options || [];
        const correctOpt = options.find(o => o.isCorrect);
        const correctShape = correctOpt?.shape || 'triangle';
        setCurrentCorrectShape(correctShape);

        const optionsForClient = options.map(o => ({
            ...o,
            isCorrect: false
        }));

        broadcast({ 
            type: 'QUESTION_START', 
            payload: { 
                totalQuestions: currentQuiz.questions.length, 
                currentQuestionIndex: currentIndex,
                timeLimit: QUESTION_TIME,
                correctShape,
                questionText: currentQ.text,
                options: optionsForClient
            } 
        });
    }

    if (timerRef.current) clearInterval(timerRef.current);
    let tickCount = 0;
    timerRef.current = setInterval(() => {
      setTimeLeft(prev => {
        if (prev <= 100) { endQuestion(); return 0; }
        const newTime = prev - 100;
        tickCount++;
        if (tickCount % 10 === 0) {
          broadcast({ type: 'TIME_SYNC', payload: { timeLeft: newTime } });
        }
        return newTime;
      });
    }, 100);
  };

  const endQuestion = () => {
    if (timerRef.current) clearInterval(timerRef.current);
    setGameState('REVEAL');
    broadcast({ type: 'TIME_UP', payload: {} });
  };
  
  const showScoreboard = () => {
    setGameState('SCOREBOARD');
    broadcast({ type: 'SHOW_SCOREBOARD', payload: { players: playersRef.current } });
  };

  const nextQuestion = () => {
    const currentQuiz = quizRef.current;
    const currentIndex = currentQuestionIndexRef.current;
    
    if (!currentQuiz) return;
    
    if (currentIndex < currentQuiz.questions.length - 1) {
      setCurrentQuestionIndex(prev => prev + 1);
      setTimeout(() => startCountdown(), 0); 
    } else {
      setGameState('PODIUM');
      broadcast({ type: 'GAME_OVER', payload: { players } });
    }
  };

  // -- CLIENT LOGIC --

  const joinGame = (pin: string, name: string) => {
    setRole('CLIENT');
    setError(null);
    setGamePin(pin);
    setCurrentPlayerName(name);
    const peer = new Peer();
    peerRef.current = peer;

    peer.on('open', (id) => {
        setCurrentPlayerId(id);
        const conn = peer.connect(`kaboot-${pin}`);
        hostConnectionRef.current = conn;
        conn.on('open', () => {
            conn.send({ type: 'JOIN', payload: { name } });
            setGameState('LOBBY');
        });
        conn.on('data', (data: any) => handleClientData(data));
        conn.on('close', () => { setError("Disconnected"); setGameState('LANDING'); });
        setTimeout(() => { if (!conn.open) setError("Check PIN"); }, 5000);
    });
  };

  const handleClientData = (data: NetworkMessage) => {
    if (data.type === 'WELCOME') setQuiz({ title: data.payload.quizTitle, questions: [] });
    
    if (data.type === 'START_COUNTDOWN') {
        setGameState('COUNTDOWN');
        setTimeLeft(data.payload.duration);
        if (timerRef.current) clearInterval(timerRef.current);
        timerRef.current = setInterval(() => {
             setTimeLeft(prev => Math.max(0, prev - 1));
        }, 1000);
    }

    if (data.type === 'QUESTION_START') {
        setGameState('QUESTION');
        setHasAnswered(false);
        setLastPointsEarned(null);
        setSelectedOption(null);
        setCurrentQuestionIndex(data.payload.currentQuestionIndex);
        setTimeLeft(data.payload.timeLimit * 1000);
        setCurrentCorrectShape(data.payload.correctShape);
        
        setQuiz(prev => {
            const tempQuestions = prev ? [...prev.questions] : [];
            while (tempQuestions.length < data.payload.totalQuestions) {
                 tempQuestions.push({ id: `loading-${tempQuestions.length}`, text: '', options: [], timeLimit: 0 });
            }
            tempQuestions[data.payload.currentQuestionIndex] = {
                id: `q-${data.payload.currentQuestionIndex}`,
                text: data.payload.questionText,
                options: data.payload.options,
                timeLimit: data.payload.timeLimit
            };
            return { title: prev?.title || 'Quiz', questions: tempQuestions };
        });

        if (timerRef.current) clearInterval(timerRef.current);
        timerRef.current = setInterval(() => setTimeLeft(prev => Math.max(0, prev - 100)), 100);
    }

    if (data.type === 'RESULT') {
        setLastPointsEarned(data.payload.scoreAdded);
        setCurrentPlayerScore(data.payload.newScore);
        if (data.payload.isCorrect) {
            setCurrentStreak(prev => prev + 1);
        } else {
            setCurrentStreak(0);
        }
    }
    
    if (data.type === 'TIME_SYNC') {
        setTimeLeft(data.payload.timeLeft);
    }
    
    if (data.type === 'TIME_UP') {
        if (timerRef.current) clearInterval(timerRef.current);
        setGameState('REVEAL');
    }

    if (data.type === 'SHOW_SCOREBOARD') {
        setGameState('SCOREBOARD');
        setPlayers(data.payload.players);
    }
    if (data.type === 'GAME_OVER') {
        setGameState('PODIUM');
        setPlayers(data.payload.players);
    }
  };

  const handleAnswer = (arg: boolean | AnswerOption) => {
    if (hasAnswered || gameState !== 'QUESTION') return;
    if (role === 'HOST' && !gameConfigRef.current.hostParticipates) return;
    
    setHasAnswered(true);

    if (role === 'HOST') {
        const option = arg as AnswerOption;
        const isCorrect = option.isCorrect;
        setSelectedOption(option);
        
        const hostPlayer = playersRef.current.find(p => p.id === 'host');
        const currentStrk = hostPlayer?.streak || 0;
        const newStreak = isCorrect ? currentStrk + 1 : 0;
        
        const isFirstCorrect = isCorrect && firstCorrectPlayerId === null;
        if (isFirstCorrect) {
          setFirstCorrectPlayerId('host');
        }
        
        const playerRank = getPlayerRank('host', playersRef.current);
        
        const points = calculatePoints({
          isCorrect,
          timeLeftMs: timeLeftRef.current,
          questionTimeMs: QUESTION_TIME_MS,
          streak: newStreak,
          playerRank,
          isFirstCorrect,
          config: gameConfigRef.current,
        });
        
        setLastPointsEarned(points);
        const newScore = Math.max(0, (hostPlayer?.score || 0) + points);
        setCurrentPlayerScore(newScore);
        setCurrentStreak(newStreak);
        
        setPlayers(prev => prev.map(p => {
          if (p.id !== 'host') return p;
          return { ...p, score: newScore, streak: newStreak, lastAnswerCorrect: isCorrect };
        }));
    } else {
        const option = arg as AnswerOption;
        setSelectedOption(option);
        const isCorrect = option.shape === currentCorrectShape;
        hostConnectionRef.current?.send({ type: 'ANSWER', payload: { playerId: peerRef.current?.id, isCorrect } });
    }
  };

  useEffect(() => {
    if (role === 'HOST' && (gameState === 'SCOREBOARD' || gameState === 'PODIUM')) {
        broadcast({ type: gameState === 'SCOREBOARD' ? 'SHOW_SCOREBOARD' : 'GAME_OVER', payload: { players } });
    }
  }, [gameState, players, role]);

  useEffect(() => {
    return () => {
      if (timerRef.current) clearInterval(timerRef.current);
      if (peerRef.current) peerRef.current.destroy();
    };
  }, []);

  return {
    role, gameState, quiz, players, currentQuestionIndex, timeLeft, error, gamePin, hasAnswered, lastPointsEarned, currentCorrectShape, selectedOption, currentPlayerScore, currentStreak, currentPlayerId, gameConfig,
    pendingQuizToSave, dismissSavePrompt, sourceQuizId,
    startQuizGen, startManualCreation, finalizeManualQuiz, loadSavedQuiz, joinGame, startGame: startHostGame, handleAnswer, nextQuestion, showScoreboard,
    updateQuizFromEditor, startGameFromEditor, backFromEditor
  };
};