F_Pro/src/hooks/useWebRTC.ts

import { useEffect, useRef, useState, useCallback } from 'react';
import { Client } from '@stomp/stompjs';
import SockJS from 'sockjs-client';
import { SignalingMessage } from '../types';

import { getWsUrl } from '../lib/api/config';

const ICE_SERVERS: RTCConfiguration = {
  iceServers: [
    { urls: 'stun:stun.l.google.com:19302' },
    { urls: 'stun:stun1.l.google.com:19302' },
    { urls: 'stun:stun2.l.google.com:19302' },
    {
      urls: 'turn:openrelay.metered.ca:80',
      username: 'openrelayproject',
      credential: 'openrelayproject',
    },
    {
      urls: 'turn:openrelay.metered.ca:443',
      username: 'openrelayproject',
      credential: 'openrelayproject',
    },
    {
      urls: 'turn:openrelay.metered.ca:443?transport=tcp',
      username: 'openrelayproject',
      credential: 'openrelayproject',
    },
  ],
};

export function useWebRTC(sessionId: string, userId: string) {
  const [localStream, setLocalStream] = useState<MediaStream | null>(null);
  const [remoteStream, setRemoteStream] = useState<MediaStream | null>(null);
  const [isConnected, setIsConnected] = useState(false);
  const [isCameraAvailable, setIsCameraAvailable] = useState(true);
  const [messages, setMessages] = useState<{ sender: string; message: string; time: string }[]>([]);
  const [connectionQuality, setConnectionQuality] = useState<'good' | 'fair' | 'poor' | 'unknown'>('unknown');

  const peerConnection = useRef<RTCPeerConnection | null>(null);
  const stompClient = useRef<Client | null>(null);
  const localStreamRef = useRef<MediaStream | null>(null);
  // Tracks whether we're mid-negotiation to avoid duplicate offer/answer rounds
  const isNegotiatingRef = useRef(false);
  // Queued ICE candidates received before remote description was set
  const pendingCandidates = useRef<RTCIceCandidateInit[]>([]);

  // ─── Signaling ────────────────────────────────────────────────────────────────
  const sendSignal = useCallback((type: SignalingMessage['type'] | 'chat', data: any) => {
    if (stompClient.current && stompClient.current.connected) {
      const message: any = { type, sender: userId, sessionId, data };
      console.log(`Sending ${type} to signaling server...`);
      stompClient.current.publish({
        destination: '/app/signal',
        body: JSON.stringify(message),
      });
    } else {
      console.warn('Cannot send signal: Stomp client not connected');
    }
  }, [sessionId, userId]);

  const sendChatMessage = (text: string) => {
    sendSignal('chat', text);
    setMessages(prev => [...prev, {
      sender: 'You',
      message: text,
      time: new Date().toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' }),
    }]);
  };

  // ─── Peer Connection ──────────────────────────────────────────────────────────
  const createPeerConnection = useCallback(() => {
    // Close any existing connection before creating a new one
    if (peerConnection.current) {
      peerConnection.current.close();
      peerConnection.current = null;
    }
    isNegotiatingRef.current = false;
    pendingCandidates.current = [];

    console.log('Creating RTCPeerConnection...');
    const pc = new RTCPeerConnection(ICE_SERVERS);

    pc.onicecandidate = (event) => {
      if (event.candidate) {
        console.log('Generated ICE candidate, sending...');
        sendSignal('ice-candidate', event.candidate);
      }
    };

    pc.onconnectionstatechange = () => {
      console.log('WebRTC Connection State:', pc.connectionState);
      if (pc.connectionState === 'connected') {
        isNegotiatingRef.current = false;
      }
    };

    pc.onsignalingstatechange = () => {
      console.log('WebRTC Signaling State:', pc.signalingState);
    };

    pc.ontrack = (event) => {
      console.log('Received remote track!', event.streams[0]);
      setRemoteStream(event.streams[0]);
    };

    peerConnection.current = pc;
    return pc;
  }, [sendSignal]);

  const addLocalTracks = useCallback((pc: RTCPeerConnection) => {
    if (!localStreamRef.current) return;
    const existingSenders = pc.getSenders();
    localStreamRef.current.getTracks().forEach((track) => {
      const alreadyAdded = existingSenders.some(s => s.track === track);
      if (!alreadyAdded) {
        pc.addTrack(track, localStreamRef.current!);
      }
    });
  }, []);

  // Drain any ICE candidates that arrived before remote description was set
  const drainPendingCandidates = useCallback(async (pc: RTCPeerConnection) => {
    for (const candidate of pendingCandidates.current) {
      try {
        await pc.addIceCandidate(new RTCIceCandidate(candidate));
      } catch (e) {
        console.warn('Error adding queued ice candidate', e);
      }
    }
    pendingCandidates.current = [];
  }, []);

  // ─── Offer / Answer handlers ──────────────────────────────────────────────────
  const handleOffer = useCallback(async (offer: RTCSessionDescriptionInit) => {
    console.log('Received Offer, creating Answer...');
    let pc = peerConnection.current;

    // If we already have a stable connection, ignore stale re-offers
    if (pc && pc.signalingState === 'stable' && pc.connectionState === 'connected') {
      console.warn('Ignoring offer — peer connection already connected and stable.');
      return;
    }

    // If there's no peer connection yet, or it's in a bad state, create a fresh one
    if (!pc || pc.connectionState === 'failed' || pc.connectionState === 'closed') {
      pc = createPeerConnection();
    }

    addLocalTracks(pc);

    try {
      await pc.setRemoteDescription(new RTCSessionDescription(offer));
      await drainPendingCandidates(pc);
      const answer = await pc.createAnswer();
      await pc.setLocalDescription(answer);
      sendSignal('answer', answer);
      isNegotiatingRef.current = false;
    } catch (e) {
      console.error('Error handling offer:', e);
      isNegotiatingRef.current = false;
    }
  }, [createPeerConnection, addLocalTracks, drainPendingCandidates, sendSignal]);

  const handleAnswer = useCallback(async (answer: RTCSessionDescriptionInit) => {
    const pc = peerConnection.current;
    if (!pc) return;

    // Only accept an answer when we're actually waiting for one
    if (pc.signalingState !== 'have-local-offer') {
      console.warn(
        `Ignoring answer — unexpected signaling state: ${pc.signalingState}`
      );
      return;
    }

    console.log('Received Answer, setting remote description...');
    try {
      await pc.setRemoteDescription(new RTCSessionDescription(answer));
      await drainPendingCandidates(pc);
    } catch (e) {
      console.error('Error handling answer:', e);
    } finally {
      isNegotiatingRef.current = false;
    }
  }, [drainPendingCandidates]);

  const handleIceCandidate = useCallback(async (candidate: RTCIceCandidateInit) => {
    const pc = peerConnection.current;
    if (!pc) return;

    // If remote description isn't set yet, queue the candidate
    if (!pc.remoteDescription) {
      pendingCandidates.current.push(candidate);
      return;
    }

    try {
      await pc.addIceCandidate(new RTCIceCandidate(candidate));
    } catch (e) {
      console.error('Error adding received ice candidate', e);
    }
  }, []);

  const startCall = useCallback(async () => {
    // Prevent duplicate simultaneous offer creation
    if (isNegotiatingRef.current) {
      console.warn('Already negotiating — ignoring duplicate startCall.');
      return;
    }

    const pc = peerConnection.current;
    // Don't start a new call if we're already connected
    if (pc && pc.connectionState === 'connected') {
      console.warn('Already connected — ignoring startCall.');
      return;
    }

    isNegotiatingRef.current = true;
    console.log('Starting call (sending Offer)...');

    const freshPc = createPeerConnection();
    addLocalTracks(freshPc);

    try {
      const offer = await freshPc.createOffer();
      await freshPc.setLocalDescription(offer);
      sendSignal('offer', offer);
    } catch (e) {
      console.error('Error starting call:', e);
      isNegotiatingRef.current = false;
    }
  }, [createPeerConnection, addLocalTracks, sendSignal]);

  // ─── Connection Quality Monitor ───────────────────────────────────────────────
  useEffect(() => {
    const interval = setInterval(async () => {
      if (!isConnected) { setConnectionQuality('unknown'); return; }

      if (peerConnection.current && peerConnection.current.connectionState === 'connected') {
        try {
          const stats = await peerConnection.current.getStats();
          let currentRtt = 0;
          stats.forEach(report => {
            if (report.type === 'candidate-pair' && report.state === 'succeeded' && report.currentRoundTripTime) {
              currentRtt = report.currentRoundTripTime * 1000;
            }
          });
          if (currentRtt > 0) {
            if (currentRtt < 150) setConnectionQuality('good');
            else if (currentRtt < 400) setConnectionQuality('fair');
            else setConnectionQuality('poor');
          } else {
            setConnectionQuality('good');
          }
        } catch {
          setConnectionQuality('good');
        }
      } else {
        setConnectionQuality('unknown');
      }
    }, 3000);
    return () => clearInterval(interval);
  }, [isConnected]);

  // ─── Main Init Effect ─────────────────────────────────────────────────────────
  useEffect(() => {
    if (sessionId === 'unknown') return;

    const init = async () => {
      // 1. Acquire camera/mic
      try {
        console.log('Requesting camera/mic permissions...');
        const stream = await navigator.mediaDevices.getUserMedia({ video: true, audio: true });
        setLocalStream(stream);
        setIsCameraAvailable(true);
        localStreamRef.current = stream;
      } catch (err) {
        console.warn('Camera/Mic access denied or busy. Proceeding with signaling only.', err);
        setIsCameraAvailable(false);
      }

      // 2. Connect to signaling server
      try {
        const wsBase = getWsUrl();

        const client = new Client({
          webSocketFactory: () => {
            // On Lightning AI (and similar cloud proxies) raw WebSocket upgrades
            // are often blocked. Force SockJS to use xhr-streaming / xhr-polling
            // so it never attempts a raw wss:// connection.
            const isCloudProxy =
              typeof window !== 'undefined' &&
              (window.location.hostname.includes('litng.ai') ||
               window.location.hostname.includes('lightning.ai'));
            return isCloudProxy
              ? new SockJS(`${wsBase}/ws`, null, { transports: ['xhr-streaming', 'xhr-polling'] })
              : new SockJS(`${wsBase}/ws`);
          },
          reconnectDelay: 5000,

          onConnect: () => {
            console.log('Connected to signaling server for session:', sessionId);
            setIsConnected(true);

            client.subscribe(`/topic/session/${sessionId}`, (message) => {
              const signal: any = JSON.parse(message.body);
              if (signal.sender === userId) return; // ignore own signals

              console.log(`Received ${signal.type} from ${signal.sender}`);

              switch (signal.type) {
                case 'offer':         handleOffer(signal.data);         break;
                case 'answer':        handleAnswer(signal.data);        break;
                case 'ice-candidate': handleIceCandidate(signal.data);  break;
                case 'chat':
                  setMessages(prev => [...prev, {
                    sender: signal.sender,
                    message: signal.data,
                    time: new Date().toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' }),
                  }]);
                  break;
                case 'join':
                  // Only the SECOND joiner starts the call to avoid both peers
                  // creating offers simultaneously.
                  console.log('Peer joined:', signal.sender, '— starting offer as second joiner');
                  startCall();
                  break;
              }
            });

            // Announce presence — the already-waiting peer will trigger startCall
            client.publish({
              destination: '/app/signal',
              body: JSON.stringify({ type: 'join', sender: userId, sessionId }),
            });
          },

          onDisconnect: () => {
            console.warn('Disconnected from signaling server, retrying in 5s...');
            setIsConnected(false);
            // Reset negotiation state so reconnect can re-initiate cleanly
            isNegotiatingRef.current = false;
          },

          onStompError: (frame) => {
            console.error('STOMP protocol error:', frame.headers['message']);
          },
        });

        client.activate();
        stompClient.current = client;
      } catch (err) {
        console.error('Failed to connect to signaling server:', err);
      }
    };

    init();

    return () => {
      stompClient.current?.deactivate();
      peerConnection.current?.close();
      localStreamRef.current?.getTracks().forEach(track => track.stop());
    };
  }, [sessionId, userId, handleOffer, handleAnswer, handleIceCandidate, startCall]);

  // ─── Media Toggles ────────────────────────────────────────────────────────────
  const toggleVideo = () => {
    const track = localStreamRef.current?.getVideoTracks()[0];
    if (track) track.enabled = !track.enabled;
  };

  const toggleAudio = () => {
    const track = localStreamRef.current?.getAudioTracks()[0];
    if (track) track.enabled = !track.enabled;
  };

  return {
    localStream,
    remoteStream,
    isConnected,
    isCameraAvailable,
    messages,
    sendChatMessage,
    toggleVideo,
    toggleAudio,
    connectionQuality,
  };
}