import React, {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useRef,
  useState,
} from "react";
import { Message, WsIncoming, WsSession } from "../types";
import { useConnection } from "./ConnectionContext";
import { playAudio, encodeAudioToBase64 } from "../services/audio";

function generateId(): string {
  return Date.now().toString(36) + Math.random().toString(36).slice(2);
}

interface ChatContextValue {
  messages: Message[];
  sendTextMessage: (text: string) => void;
  sendVoiceMessage: (audioUri: string, durationMs?: number) => void;
  clearMessages: () => void;
  // Session management
  sessions: WsSession[];
  requestSessions: () => void;
  switchSession: (sessionId: string) => void;
  renameSession: (sessionId: string, name: string) => void;
  // Screenshot / navigation
  latestScreenshot: string | null;
  requestScreenshot: () => void;
  sendNavKey: (key: string) => void;
}

const ChatContext = createContext<ChatContextValue | null>(null);

export function ChatProvider({ children }: { children: React.ReactNode }) {
  const [messages, setMessages] = useState<Message[]>([]);
  const [sessions, setSessions] = useState<WsSession[]>([]);
  const [latestScreenshot, setLatestScreenshot] = useState<string | null>(null);
  const {
    sendTextMessage: wsSend,
    sendVoiceMessage: wsVoice,
    sendCommand,
    onMessageReceived,
  } = useConnection();

  const addMessage = useCallback((msg: Message) => {
    setMessages((prev) => [...prev, msg]);
  }, []);

  const updateMessageStatus = useCallback(
    (id: string, status: Message["status"]) => {
      setMessages((prev) =>
        prev.map((m) => (m.id === id ? { ...m, status } : m))
      );
    },
    []
  );

  // Handle incoming WebSocket messages
  useEffect(() => {
    onMessageReceived.current = (data: WsIncoming) => {
      switch (data.type) {
        case "text": {
          const msg: Message = {
            id: generateId(),
            role: "assistant",
            type: "text",
            content: data.content,
            timestamp: Date.now(),
            status: "sent",
          };
          setMessages((prev) => [...prev, msg]);
          break;
        }
        case "voice": {
          const msg: Message = {
            id: generateId(),
            role: "assistant",
            type: "voice",
            content: data.content ?? "",
            audioUri: data.audioBase64
              ? `data:audio/mp4;base64,${data.audioBase64}`
              : undefined,
            timestamp: Date.now(),
            status: "sent",
          };
          setMessages((prev) => [...prev, msg]);
          if (msg.audioUri) {
            playAudio(msg.audioUri).catch(() => {});
          }
          break;
        }
        case "image": {
          // Store as latest screenshot for navigation mode
          setLatestScreenshot(data.imageBase64);
          // Also add to chat as an image message
          const msg: Message = {
            id: generateId(),
            role: "assistant",
            type: "image",
            content: data.caption ?? "Screenshot",
            imageBase64: data.imageBase64,
            timestamp: Date.now(),
            status: "sent",
          };
          setMessages((prev) => [...prev, msg]);
          break;
        }
        case "sessions": {
          setSessions(data.sessions);
          break;
        }
        case "session_switched": {
          const msg: Message = {
            id: generateId(),
            role: "system",
            type: "text",
            content: `Switched to ${data.name}`,
            timestamp: Date.now(),
          };
          setMessages((prev) => [...prev, msg]);
          break;
        }
        case "session_renamed": {
          const msg: Message = {
            id: generateId(),
            role: "system",
            type: "text",
            content: `Renamed to ${data.name}`,
            timestamp: Date.now(),
          };
          setMessages((prev) => [...prev, msg]);
          // Refresh sessions to show updated name
          sendCommand("sessions");
          break;
        }
        case "error": {
          const msg: Message = {
            id: generateId(),
            role: "system",
            type: "text",
            content: data.message,
            timestamp: Date.now(),
          };
          setMessages((prev) => [...prev, msg]);
          break;
        }
      }
    };

    return () => {
      onMessageReceived.current = null;
    };
  }, [onMessageReceived, sendCommand]);

  const sendTextMessage = useCallback(
    (text: string) => {
      const id = generateId();
      const msg: Message = {
        id,
        role: "user",
        type: "text",
        content: text,
        timestamp: Date.now(),
        status: "sending",
      };
      addMessage(msg);
      const sent = wsSend(text);
      updateMessageStatus(id, sent ? "sent" : "error");
    },
    [wsSend, addMessage, updateMessageStatus]
  );

  const sendVoiceMessage = useCallback(
    async (audioUri: string, durationMs?: number) => {
      const id = generateId();
      const msg: Message = {
        id,
        role: "user",
        type: "voice",
        content: "",
        audioUri,
        timestamp: Date.now(),
        status: "sending",
        duration: durationMs,
      };
      addMessage(msg);
      try {
        const base64 = await encodeAudioToBase64(audioUri);
        const sent = wsVoice(base64);
        updateMessageStatus(id, sent ? "sent" : "error");
      } catch (err) {
        console.error("Failed to encode audio:", err);
        updateMessageStatus(id, "error");
      }
    },
    [wsVoice, addMessage, updateMessageStatus]
  );

  const clearMessages = useCallback(() => {
    setMessages([]);
  }, []);

  // --- Session management ---
  const requestSessions = useCallback(() => {
    sendCommand("sessions");
  }, [sendCommand]);

  const switchSession = useCallback(
    (sessionId: string) => {
      sendCommand("switch", { sessionId });
    },
    [sendCommand]
  );

  const renameSession = useCallback(
    (sessionId: string, name: string) => {
      sendCommand("rename", { sessionId, name });
    },
    [sendCommand]
  );

  // --- Screenshot / navigation ---
  const requestScreenshot = useCallback(() => {
    sendCommand("screenshot");
  }, [sendCommand]);

  const sendNavKey = useCallback(
    (key: string) => {
      sendCommand("nav", { key });
    },
    [sendCommand]
  );

  return (
    <ChatContext.Provider
      value={{
        messages,
        sendTextMessage,
        sendVoiceMessage,
        clearMessages,
        sessions,
        requestSessions,
        switchSession,
        renameSession,
        latestScreenshot,
        requestScreenshot,
        sendNavKey,
      }}
    >
      {children}
    </ChatContext.Provider>
  );
}

export function useChat() {
  const ctx = useContext(ChatContext);
  if (!ctx) throw new Error("useChat must be used within ChatProvider");
  return ctx;
}