import React, {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useRef,
  useState,
} from "react";
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 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>;
  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 [latestScreenshot, setLatestScreenshot] = useState<string | null>(null);
  const needsSync = useRef(true);

  // 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>>({});
  // Typing indicator from server
  const [isTyping, setIsTyping] = useState(false);
  // 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) {
          if (prev) {
            messagesMapRef.current[prev] = messages;
          }
          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;
          });
        }
        return active.id;
      });
    }
  }, [messages]);

  // On connect: ask gateway to sync sessions. If we already had a session
  // selected, tell the gateway so it preserves our selection instead of
  // jumping to whatever iTerm has focused on the Mac.
  useEffect(() => {
    if (status === "connected") {
      needsSync.current = true;
      sendCommand("sync", activeSessionId ? { activeSessionId } : undefined);
    } else if (status === "disconnected") {
      setIsTyping(false);
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps — only fire on status change
  }, [status, sendCommand]);

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

  // Helper: add a message to a specific session (may not be active)
  const addMessageToSession = useCallback((sessionId: string, msg: Message) => {
    setActiveSessionId((currentActive) => {
      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,
        }));
      }
      return currentActive;
    });
  }, []);

  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
      );
      setActiveSessionId((sessId) => {
        if (sessId) {
          messagesMapRef.current[sessId] = next;
          debouncedSave(messagesMapRef.current);
        }
        return sessId;
      });
      return next;
    });
  }, []);

  // Handle incoming WebSocket messages
  useEffect(() => {
    onMessageReceived.current = async (data: WsIncoming) => {
      switch (data.type) {
        case "text": {
          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);
          }
          notifyIncomingMessage("PAILot", data.content ?? "New message");
          break;
        }
        case "voice": {
          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",
          };
          if (data.sessionId) {
            addMessageToSession(data.sessionId, msg);
          } else {
            addMessageToActive(msg);
          }
          notifyIncomingMessage("PAILot", data.content ?? "Voice message");
          if (msg.audioUri && canAutoplay()) {
            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);
          }
          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": {
          const msg: Message = {
            id: generateId(),
            role: "system",
            type: "text",
            content: `Switched to ${data.name}`,
            timestamp: Date.now(),
          };
          addMessageToActive(msg);
          sendCommand("sessions");
          break;
        }
        case "session_renamed": {
          const msg: Message = {
            id: generateId(),
            role: "system",
            type: "text",
            content: `Renamed to ${data.name}`,
            timestamp: Date.now(),
          };
          addMessageToActive(msg);
          sendCommand("sessions");
          break;
        }
        case "transcript": {
          // Voice → text reflection: replace voice bubble with transcribed text
          updateMessageContent(data.messageId, data.content);
          break;
        }
        case "typing": {
          setIsTyping(data.typing);
          break;
        }
        case "status": {
          // Connection status update — ignore for now
          break;
        }
        case "projects": {
          setProjects(data.projects ?? []);
          break;
        }
        case "error": {
          const msg: Message = {
            id: generateId(),
            role: "system",
            type: "text",
            content: data.message,
            timestamp: Date.now(),
          };
          addMessageToActive(msg);
          break;
        }
      }
    };

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

  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);
      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);
        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);
      updateMessageStatus(id, sent ? "sent" : "error");
    },
    [wsImageSend, addMessageToActive, updateMessageStatus]
  );

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

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

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

  const switchSession = useCallback(
    (sessionId: string) => {
      setActiveSessionId((prev) => {
        if (prev) {
          messagesMapRef.current[prev] = messages;
          debouncedSave(messagesMapRef.current);
        }
        return prev;
      });
      sendCommand("switch", { sessionId });
    },
    [sendCommand, messages]
  );

  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;
      });
    },
    [sendCommand]
  );

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

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

  const loadMoreMessages = useCallback(() => {
    setActiveSessionId((sessId) => {
      if (!sessId) return sessId;
      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;
      });
      return sessId;
    });
  }, []);

  // --- 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,
        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;
}