import React, {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useRef,
  useState,
} from "react";
import { AppState, AppStateStatus } from "react-native";
import { Message, WsIncoming, WsSession, PaiProject } from "../types";
import { useConnection } from "./ConnectionContext";
import { playAudio, encodeAudioToBase64, saveBase64Audio, canAutoplay } from "../services/audio";
import { requestNotificationPermissions, notifyIncomingMessage } from "../services/notifications";

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

// --- Message persistence ---
// Lazily import expo-file-system/legacy so a missing native module doesn't crash the app.

let _fsReady: Promise<typeof import("expo-file-system/legacy")> | null = null;
function getFs() {
  if (!_fsReady) _fsReady = import("expo-file-system/legacy");
  return _fsReady;
}

const MESSAGES_DIR = "pailot-messages";

/** Strip heavy fields (base64 images, audio URIs) before persisting.
 *  Voice messages keep their content (transcript) but lose audioUri
 *  since cache files won't survive app restarts. */
function lightMessage(m: Message): Message {
  const light = { ...m };
  if (light.imageBase64) light.imageBase64 = undefined;
  if (light.audioUri) light.audioUri = undefined;
  return light;
}

async function persistMessages(map: Record<string, Message[]>): Promise<void> {
  try {
    const fs = await getFs();
    const dir = `${fs.documentDirectory}${MESSAGES_DIR}/`;
    const dirInfo = await fs.getInfoAsync(dir);
    if (!dirInfo.exists) await fs.makeDirectoryAsync(dir, { intermediates: true });
    // Save each session's messages
    for (const [sessionId, msgs] of Object.entries(map)) {
      if (msgs.length === 0) continue;
      const light = msgs.map(lightMessage);
      await fs.writeAsStringAsync(`${dir}${sessionId}.json`, JSON.stringify(light));
    }
  } catch {
    // Persistence is best-effort
  }
}

async function loadMessages(): Promise<Record<string, Message[]>> {
  try {
    const fs = await getFs();
    const dir = `${fs.documentDirectory}${MESSAGES_DIR}/`;
    const dirInfo = await fs.getInfoAsync(dir);
    if (!dirInfo.exists) return {};
    const files = await fs.readDirectoryAsync(dir);
    const result: Record<string, Message[]> = {};
    for (const file of files) {
      if (!file.endsWith(".json")) continue;
      const sessionId = file.replace(".json", "");
      const content = await fs.readAsStringAsync(`${dir}${file}`);
      result[sessionId] = (JSON.parse(content) as Message[])
        // Drop voice messages with no audio and no content (empty chunks)
        .filter((m) => !(m.type === "voice" && !m.audioUri && !m.content))
        .map((m) => {
          // Voice messages without audio but with transcript → show as text
          if (m.type === "voice" && !m.audioUri && m.content) {
            return { ...m, type: "text" };
          }
          return m;
        });
    }
    return result;
  } catch {
    return {};
  }
}

async function deletePersistedSession(sessionId: string): Promise<void> {
  try {
    const fs = await getFs();
    const path = `${fs.documentDirectory}${MESSAGES_DIR}/${sessionId}.json`;
    const info = await fs.getInfoAsync(path);
    if (info.exists) await fs.deleteAsync(path);
  } catch {
    // Best-effort
  }
}

async function clearPersistedMessages(sessionId: string): Promise<void> {
  try {
    const fs = await getFs();
    await fs.writeAsStringAsync(
      `${fs.documentDirectory}${MESSAGES_DIR}/${sessionId}.json`,
      "[]"
    );
  } catch {
    // Best-effort
  }
}

// --- Debounced save ---
let saveTimer: ReturnType<typeof setTimeout> | null = null;
function debouncedSave(map: Record<string, Message[]>): void {
  if (saveTimer) clearTimeout(saveTimer);
  saveTimer = setTimeout(() => persistMessages(map), 1000);
}

const PAGE_SIZE = 50;

// --- Context ---

interface IncomingToast {
  sessionId: string;
  sessionName: string;
  preview: string;
}

interface ChatContextValue {
  messages: Message[];
  sendTextMessage: (text: string) => void;
  sendVoiceMessage: (audioUri: string, durationMs?: number) => void;
  sendImageMessage: (imageBase64: string, caption: string, mimeType: string) => void;
  deleteMessage: (id: string) => void;
  clearMessages: () => void;
  isTyping: boolean;
  sessions: WsSession[];
  activeSessionId: string | null;
  requestSessions: () => void;
  switchSession: (sessionId: string) => void;
  renameSession: (sessionId: string, name: string) => void;
  removeSession: (sessionId: string) => void;
  createSession: (opts?: { project?: string; path?: string }) => void;
  fetchProjects: () => void;
  projects: PaiProject[];
  loadMoreMessages: () => void;
  hasMoreMessages: boolean;
  unreadCounts: Record<string, number>;
  unreadSessions: Set<string>;
  incomingToast: IncomingToast | null;
  dismissToast: () => void;
  latestScreenshot: string | null;
  requestScreenshot: () => void;
  sendNavKey: (key: string) => void;
}

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

export function ChatProvider({ children }: { children: React.ReactNode }) {
  const [sessions, setSessions] = useState<WsSession[]>([]);
  const [activeSessionId, setActiveSessionId] = useState<string | null>(null);
  const activeSessionIdRef = useRef<string | null>(null);
  const [latestScreenshot, setLatestScreenshot] = useState<string | null>(null);
  const needsSync = useRef(true);

  // Sequence tracking for catch_up protocol
  const lastSeqRef = useRef(0);
  const seenSeqsRef = useRef(new Set<number>());

  // Per-session message storage
  const messagesMapRef = useRef<Record<string, Message[]>>({});
  // Messages for the active session (drives re-renders)
  const [messages, setMessages] = useState<Message[]>([]);
  // Unread counts for non-active sessions
  const [unreadCounts, setUnreadCounts] = useState<Record<string, number>>({});
  // Server-pushed unread indicators (sessions with new activity since last viewed)
  const [unreadSessions, setUnreadSessions] = useState<Set<string>>(new Set());
  // Per-session typing indicator (sessionId → boolean)
  const typingMapRef = useRef<Record<string, boolean>>({});
  const [isTyping, setIsTyping] = useState(false);
  // Toast queue for other-session incoming messages (show one at a time)
  const toastQueueRef = useRef<{ sessionId: string; sessionName: string; preview: string }[]>([]);
  const [incomingToast, setIncomingToast] = useState<{ sessionId: string; sessionName: string; preview: string } | null>(null);
  // PAI projects list
  const [projects, setProjects] = useState<PaiProject[]>([]);
  // Pagination: does the active session have more messages in storage?
  const [hasMoreMessages, setHasMoreMessages] = useState(false);

  const {
    status,
    sendTextMessage: wsSend,
    sendVoiceMessage: wsVoice,
    sendImageMessage: wsImageSend,
    sendCommand,
    onMessageReceived,
  } = useConnection();

  // Restore persisted messages on mount + request notification permissions
  useEffect(() => {
    loadMessages().then((loaded) => {
      if (Object.keys(loaded).length > 0) {
        messagesMapRef.current = loaded;
      }
    });
    requestNotificationPermissions();
  }, []);

  // Derive active session ID from sessions list when it arrives
  const syncActiveFromSessions = useCallback((incoming: WsSession[]) => {
    const active = incoming.find((s) => s.isActive);
    if (active) {
      setActiveSessionId((prev) => {
        if (prev !== active.id) {
          // No need to save prev — messagesMapRef is kept in sync by all mutators
          const all = messagesMapRef.current[active.id] ?? [];
          const page = all.length > PAGE_SIZE ? all.slice(-PAGE_SIZE) : all;
          setMessages(page);
          setHasMoreMessages(all.length > PAGE_SIZE);
          setUnreadCounts((u) => {
            if (!u[active.id]) return u;
            const next = { ...u };
            delete next[active.id];
            return next;
          });
          setUnreadSessions((prev) => {
            if (!prev.has(active.id)) return prev;
            const next = new Set(prev);
            next.delete(active.id);
            return next;
          });
          // Sync typing indicator for the new active session
          const activeTyping = typingMapRef.current[active.id] ?? false;
          setIsTyping(activeTyping);
        }
        activeSessionIdRef.current = active.id;
        return active.id;
      });
    }
  }, []);

  // On connect: ask gateway to sync sessions, then request catch_up for missed messages.
  useEffect(() => {
    if (status === "connected") {
      needsSync.current = true;
      const id = activeSessionIdRef.current;
      sendCommand("sync", id ? { activeSessionId: id } : undefined);
      // Request any messages we missed while disconnected/backgrounded
      sendCommand("catch_up", { lastSeq: lastSeqRef.current });
    } else if (status === "disconnected") {
      setIsTyping(false);
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps — only fire on status change
  }, [status, sendCommand]);

  // On foreground resume: request catch_up for any messages missed while backgrounded.
  // iOS keeps the WebSocket "open" at TCP level but suspends the app — messages sent
  // during that time are lost. catch_up replays them from the server's message log.
  useEffect(() => {
    let lastState: AppStateStatus = AppState.currentState;
    const sub = AppState.addEventListener("change", (nextState) => {
      if (lastState.match(/inactive|background/) && nextState === "active") {
        if (status === "connected") {
          sendCommand("catch_up", { lastSeq: lastSeqRef.current });
        }
      }
      lastState = nextState;
    });
    return () => sub.remove();
  }, [status, sendCommand]);

  // Helper: add a message to the active session
  const addMessageToActive = useCallback((msg: Message) => {
    setMessages((prev) => {
      const next = [...prev, msg];
      const id = activeSessionIdRef.current;
      if (id) {
        messagesMapRef.current[id] = next;
        debouncedSave(messagesMapRef.current);
      }
      return next;
    });
  }, []);

  // Helper: add a message to a specific session (may not be active)
  const addMessageToSession = useCallback((sessionId: string, msg: Message) => {
    const currentActive = activeSessionIdRef.current;
    if (sessionId === currentActive) {
      setMessages((prev) => {
        const next = [...prev, msg];
        messagesMapRef.current[sessionId] = next;
        debouncedSave(messagesMapRef.current);
        return next;
      });
    } else {
      const existing = messagesMapRef.current[sessionId] ?? [];
      messagesMapRef.current[sessionId] = [...existing, msg];
      debouncedSave(messagesMapRef.current);
      setUnreadCounts((u) => ({
        ...u,
        [sessionId]: (u[sessionId] ?? 0) + 1,
      }));
      // Queue toast for other-session messages (assistant only, skip system noise)
      if (msg.role === "assistant") {
        setSessions((prev) => {
          const session = prev.find((s) => s.id === sessionId);
          const name = session?.name ?? sessionId.slice(0, 8);
          const preview = msg.type === "voice" ? "🎤 Voice note" : msg.type === "image" ? "📷 Image" : (msg.content ?? "").slice(0, 60);
          const toast = { sessionId, sessionName: name, preview };
          // If no toast is showing, show immediately; otherwise queue
          setIncomingToast((current) => {
            if (current === null) return toast;
            toastQueueRef.current.push(toast);
            return current;
          });
          return prev;
        });
      }
    }
  }, []);

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

  // Update a message's content (e.g., voice transcript reflection)
  const updateMessageContent = useCallback((id: string, content: string) => {
    setMessages((prev) => {
      const next = prev.map((m) =>
        m.id === id ? { ...m, content } : m
      );
      const sessId = activeSessionIdRef.current;
      if (sessId) {
        messagesMapRef.current[sessId] = next;
        debouncedSave(messagesMapRef.current);
      }
      return next;
    });
  }, []);

  // Process a single incoming message (used by both live delivery and catch_up replay)
  const processIncoming = useCallback(async (data: WsIncoming, isCatchUp = false) => {
    // Dedup by seq: if we've seen this seq before, skip it
    const seq = (data as any).seq as number | undefined;
    if (seq) {
      if (seenSeqsRef.current.has(seq)) return;
      seenSeqsRef.current.add(seq);
      lastSeqRef.current = Math.max(lastSeqRef.current, seq);
      // Keep seen set bounded (last 500 seqs)
      if (seenSeqsRef.current.size > 500) {
        const arr = Array.from(seenSeqsRef.current).sort((a, b) => a - b);
        seenSeqsRef.current = new Set(arr.slice(-300));
      }
    }

    switch (data.type) {
      case "text": {
        if (!isCatchUp) setIsTyping(false);
        const msg: Message = {
          id: generateId(),
          role: "assistant",
          type: "text",
          content: data.content,
          timestamp: Date.now(),
          status: "sent",
        };
        if (data.sessionId) {
          addMessageToSession(data.sessionId, msg);
        } else {
          addMessageToActive(msg);
        }
        if (!isCatchUp) notifyIncomingMessage("PAILot", data.content ?? "New message");
        break;
      }
      case "voice": {
        if (!isCatchUp) setIsTyping(false);
        let audioUri: string | undefined;
        if (data.audioBase64) {
          try {
            audioUri = await saveBase64Audio(data.audioBase64);
          } catch {
            // fallback: no playable audio
          }
        }
        const msg: Message = {
          id: generateId(),
          role: "assistant",
          type: "voice",
          content: data.content ?? "",
          audioUri,
          timestamp: Date.now(),
          status: "sent",
        };
        const isForActive = !data.sessionId || data.sessionId === activeSessionIdRef.current;
        if (data.sessionId) {
          addMessageToSession(data.sessionId, msg);
        } else {
          addMessageToActive(msg);
        }
        if (!isCatchUp) notifyIncomingMessage("PAILot", data.content ?? "Voice message");
        // Only autoplay if live (not catch_up) and for the currently viewed session
        if (!isCatchUp && msg.audioUri && canAutoplay() && isForActive) {
          playAudio(msg.audioUri).catch(() => {});
        }
        break;
      }
      case "image": {
        setLatestScreenshot(data.imageBase64);
        const msg: Message = {
          id: generateId(),
          role: "assistant",
          type: "image",
          content: data.caption ?? "Screenshot",
          imageBase64: data.imageBase64,
          timestamp: Date.now(),
          status: "sent",
        };
        if (data.sessionId) {
          addMessageToSession(data.sessionId, msg);
        } else {
          addMessageToActive(msg);
        }
        if (!isCatchUp) notifyIncomingMessage("PAILot", data.caption ?? "New image");
        break;
      }
      case "sessions": {
        const incoming = data.sessions as WsSession[];
        setSessions(incoming);
        syncActiveFromSessions(incoming);
        needsSync.current = false;
        break;
      }
      case "session_switched": {
        sendCommand("sessions");
        break;
      }
      case "session_renamed": {
        sendCommand("sessions");
        break;
      }
      case "transcript": {
        updateMessageContent(data.messageId, data.content);
        break;
      }
      case "typing": {
        const typingSession = (data.sessionId as string) || activeSessionIdRef.current || "_global";
        typingMapRef.current[typingSession] = !!data.typing;
        const activeTyping = typingMapRef.current[activeSessionIdRef.current ?? ""] ?? false;
        setIsTyping(activeTyping);
        break;
      }
      case "status": {
        break;
      }
      case "projects": {
        setProjects(data.projects ?? []);
        break;
      }
      case "unread": {
        const targetId = data.sessionId as string;
        if (targetId && targetId !== activeSessionIdRef.current) {
          setUnreadSessions((prev) => {
            if (prev.has(targetId)) return prev;
            const next = new Set(prev);
            next.add(targetId);
            return next;
          });
        }
        break;
      }
      case "error": {
        const errMsg: Message = {
          id: generateId(),
          role: "system",
          type: "text",
          content: data.message,
          timestamp: Date.now(),
        };
        addMessageToActive(errMsg);
        break;
      }
    }
  }, [addMessageToActive, addMessageToSession, sendCommand, syncActiveFromSessions, updateMessageContent]);

  // Handle incoming WebSocket messages
  useEffect(() => {
    onMessageReceived.current = async (data: WsIncoming) => {
      // Handle catch_up response: replay all missed messages
      if (data.type === "catch_up") {
        const messages = (data as any).messages as WsIncoming[];
        const serverSeq = (data as any).serverSeq as number | undefined;
        if (serverSeq) lastSeqRef.current = Math.max(lastSeqRef.current, serverSeq);
        if (messages && messages.length > 0) {
          for (const msg of messages) {
            await processIncoming(msg, true);
          }
        }
        return;
      }
      // Live message — process normally
      await processIncoming(data);
    };

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

  const sendTextMessage = useCallback(
    (text: string) => {
      const id = generateId();
      const msg: Message = {
        id,
        role: "user",
        type: "text",
        content: text,
        timestamp: Date.now(),
        status: "sending",
      };
      addMessageToActive(msg);
      const sent = wsSend(text, activeSessionIdRef.current ?? undefined);
      updateMessageStatus(id, sent ? "sent" : "error");
    },
    [wsSend, addMessageToActive, 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,
      };
      addMessageToActive(msg);
      try {
        const base64 = await encodeAudioToBase64(audioUri);
        const sent = wsVoice(base64, "", id, activeSessionIdRef.current ?? undefined);
        updateMessageStatus(id, sent ? "sent" : "error");
      } catch (err) {
        console.error("Failed to encode audio:", err);
        updateMessageStatus(id, "error");
      }
    },
    [wsVoice, addMessageToActive, updateMessageStatus]
  );

  const sendImageMessage = useCallback(
    (imageBase64: string, caption: string, mimeType: string) => {
      const id = generateId();
      const msg: Message = {
        id,
        role: "user",
        type: "image",
        content: caption || "Photo",
        imageBase64,
        timestamp: Date.now(),
        status: "sending",
      };
      addMessageToActive(msg);
      const sent = wsImageSend(imageBase64, caption, mimeType, activeSessionIdRef.current ?? undefined);
      updateMessageStatus(id, sent ? "sent" : "error");
    },
    [wsImageSend, addMessageToActive, updateMessageStatus]
  );

  const deleteMessage = useCallback((id: string) => {
    setMessages((prev) => {
      const next = prev.filter((m) => m.id !== id);
      const sessId = activeSessionIdRef.current;
      if (sessId) {
        messagesMapRef.current[sessId] = next;
        debouncedSave(messagesMapRef.current);
      }
      return next;
    });
  }, []);

  const clearMessages = useCallback(() => {
    setMessages([]);
    const id = activeSessionIdRef.current;
    if (id) {
      messagesMapRef.current[id] = [];
      clearPersistedMessages(id);
    }
  }, []);

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

  const switchSession = useCallback(
    (sessionId: string) => {
      // messagesMapRef is already kept in sync by all mutators — no need to save here
      sendCommand("switch", { sessionId });
      // Clear the server-pushed unread indicator immediately on user intent
      setUnreadSessions((prev) => {
        if (!prev.has(sessionId)) return prev;
        const next = new Set(prev);
        next.delete(sessionId);
        return next;
      });
    },
    [sendCommand]
  );

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

  const removeSession = useCallback(
    (sessionId: string) => {
      sendCommand("remove", { sessionId });
      delete messagesMapRef.current[sessionId];
      deletePersistedSession(sessionId);
      setUnreadCounts((u) => {
        if (!u[sessionId]) return u;
        const next = { ...u };
        delete next[sessionId];
        return next;
      });
      setUnreadSessions((prev) => {
        if (!prev.has(sessionId)) return prev;
        const next = new Set(prev);
        next.delete(sessionId);
        return next;
      });
    },
    [sendCommand]
  );

  const createSession = useCallback((opts?: { project?: string; path?: string }) => {
    sendCommand("create", opts ?? {});
  }, [sendCommand]);

  const fetchProjects = useCallback(() => {
    sendCommand("projects");
  }, [sendCommand]);

  const dismissToast = useCallback(() => {
    // Show next queued toast, or clear
    const next = toastQueueRef.current.shift();
    setIncomingToast(next ?? null);
  }, []);

  const loadMoreMessages = useCallback(() => {
    const sessId = activeSessionIdRef.current;
    if (!sessId) return;
    const all = messagesMapRef.current[sessId] ?? [];
    setMessages((current) => {
      if (current.length >= all.length) {
        setHasMoreMessages(false);
        return current;
      }
      const nextSize = Math.min(current.length + PAGE_SIZE, all.length);
      const page = all.slice(-nextSize);
      setHasMoreMessages(nextSize < all.length);
      return page;
    });
  }, []);

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

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

  return (
    <ChatContext.Provider
      value={{
        messages,
        sendTextMessage,
        sendVoiceMessage,
        sendImageMessage,
        deleteMessage,
        clearMessages,
        isTyping,
        sessions,
        activeSessionId,
        requestSessions,
        switchSession,
        renameSession,
        removeSession,
        createSession,
        fetchProjects,
        projects,
        loadMoreMessages,
        hasMoreMessages,
        unreadCounts,
        unreadSessions,
        incomingToast,
        dismissToast,
        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;
}