timmy-tower/artifacts/mobile/context/TimmyContext.tsx

import React, {
  createContext,
  useCallback,
  useContext,
  useEffect,
  useMemo,
  useRef,
  useState,
} from "react";
import { AppState, Platform } from "react-native";
import AsyncStorage from "@react-native-async-storage/async-storage";

import { SERVER_URL_KEY } from "@/constants/storage-keys";

export type TimmyMood = "idle" | "thinking" | "working" | "speaking";

export type WsEvent = {
  id: string;
  type: string;
  timestamp: number;
  agentId?: string;
  jobId?: string;
  text?: string;
  state?: string;
  count?: number;
};

export type ConnectionStatus =
  | "connecting"
  | "connected"
  | "disconnected"
  | "reconnecting"
  | "error";

type TimmyContextValue = {
  timmyMood: TimmyMood;
  connectionStatus: ConnectionStatus;
  /** True when the WebSocket is fully open */
  isConnected: boolean;
  recentEvents: WsEvent[];
  send: (msg: object) => void;
  sendVisitorMessage: (text: string) => void;
  visitorId: string;
  /** Current API / WebSocket base domain */
  apiBaseUrl: string;
  /** Persist a new base URL and reconnect the WebSocket */
  setApiBaseUrl: (url: string) => void;
};

const TimmyContext = createContext<TimmyContextValue | null>(null);

const MAX_EVENTS = 100;
const ENV_DOMAIN = process.env["EXPO_PUBLIC_DOMAIN"] ?? "";
const VISITOR_ID =
  Date.now().toString() + Math.random().toString(36).substr(2, 9);

function buildWsUrl(domain: string): string {
  let d = domain.trim();
  if (!d) d = "localhost:8080";
  d = d.replace(/^https?:\/\//, "");
  d = d.replace(/\/$/, "");
  const proto = d.startsWith("localhost") ? "ws" : "wss";
  return `${proto}://${d}/api/ws`;
}

function deriveMood(agentStates: Record<string, string>): TimmyMood {
  if (agentStates["gamma"] === "working") return "working";
  if (
    agentStates["beta"] === "thinking" ||
    agentStates["alpha"] === "thinking"
  )
    return "thinking";
  if (Object.values(agentStates).some((s) => s !== "idle")) return "working";
  return "idle";
}

export function TimmyProvider({ children }: { children: React.ReactNode }) {
  const [apiBaseUrl, setApiBaseUrlState] = useState(ENV_DOMAIN);
  const [timmyMood, setTimmyMood] = useState<TimmyMood>("idle");
  const [connectionStatus, setConnectionStatus] =
    useState<ConnectionStatus>("connecting");
  const [recentEvents, setRecentEvents] = useState<WsEvent[]>([]);

  const wsRef = useRef<WebSocket | null>(null);
  const retryTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const retryCountRef = useRef(0);
  const agentStatesRef = useRef<Record<string, string>>({
    alpha: "idle",
    beta: "idle",
    gamma: "idle",
    delta: "idle",
  });
  const speakingTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  // Stable ref so WebSocket callbacks always read the current URL
  const apiBaseUrlRef = useRef(apiBaseUrl);
  // Stable refs to break the connectWs ↔ scheduleRetry circular dependency
  const connectWsRef = useRef<() => void>(() => {});
  const scheduleRetryRef = useRef<() => void>(() => {});

  // ── Load persisted URL on mount ────────────────────────────────────────

  useEffect(() => {
    AsyncStorage.getItem(SERVER_URL_KEY)
      .then((stored) => {
        if (stored) {
          setApiBaseUrlState(stored);
          apiBaseUrlRef.current = stored;
        }
      })
      .catch(() => {});
  }, []);

  const setApiBaseUrl = useCallback((url: string) => {
    setApiBaseUrlState(url);
    apiBaseUrlRef.current = url;
    AsyncStorage.setItem(SERVER_URL_KEY, url).catch(() => {});
  }, []);

  // ── WebSocket helpers ──────────────────────────────────────────────────

  const addEvent = useCallback((evt: Omit<WsEvent, "id" | "timestamp">) => {
    const entry: WsEvent = {
      id: Date.now().toString() + Math.random().toString(36).substr(2, 6),
      timestamp: Date.now(),
      ...evt,
    };
    setRecentEvents((prev) => [entry, ...prev].slice(0, MAX_EVENTS));
  }, []);

  const connectWs = useCallback(() => {
    if (wsRef.current) {
      wsRef.current.onclose = null;
      wsRef.current.onerror = null;
      wsRef.current.close();
      wsRef.current = null;
    }
    const url = buildWsUrl(apiBaseUrlRef.current);
    setConnectionStatus("connecting");
    let ws: WebSocket;
    try {
      ws = new WebSocket(url);
    } catch {
      setConnectionStatus("error");
      scheduleRetryRef.current();
      return;
    }
    wsRef.current = ws;

    ws.onopen = () => {
      retryCountRef.current = 0;
      setConnectionStatus("connected");
      ws.send(
        JSON.stringify({
          type: "visitor_enter",
          visitorId: VISITOR_ID,
          visitorName: "Mobile Visitor",
        })
      );
    };

    ws.onmessage = (e) => {
      let msg: Record<string, unknown>;
      try {
        msg = JSON.parse(e.data);
      } catch {
        return;
      }
      const type = msg.type as string;

      if (type === "ping") {
        ws.send(JSON.stringify({ type: "pong" }));
        return;
      }

      if (type === "world_state") {
        const states = (msg.agentStates as Record<string, string>) ?? {};
        agentStatesRef.current = { ...agentStatesRef.current, ...states };
        setTimmyMood(deriveMood(agentStatesRef.current));
        return;
      }

      if (type === "agent_state") {
        const agentId = msg.agentId as string;
        const state = msg.state as string;
        agentStatesRef.current = {
          ...agentStatesRef.current,
          [agentId]: state,
        };
        setTimmyMood(deriveMood(agentStatesRef.current));
        addEvent({ type, agentId, state });
        return;
      }

      if (type === "chat") {
        const agentId = msg.agentId as string;
        const text = msg.text as string;
        addEvent({ type, agentId, text });

        if (agentId === "timmy" || !agentId) {
          if (speakingTimerRef.current) clearTimeout(speakingTimerRef.current);
          setTimmyMood("speaking");
          const duration = Math.max(2000, (text?.length ?? 50) * 50);
          speakingTimerRef.current = setTimeout(() => {
            setTimmyMood(deriveMood(agentStatesRef.current));
          }, duration);
        }
        return;
      }

      if (type === "job_started" || type === "job_completed") {
        addEvent({
          type,
          jobId: msg.jobId as string,
          agentId: msg.agentId as string,
        });
        return;
      }

      if (type === "visitor_count") {
        addEvent({ type, count: msg.count as number });
        return;
      }
    };

    ws.onclose = () => {
      setConnectionStatus("disconnected");
      scheduleRetryRef.current();
    };

    ws.onerror = () => {
      setConnectionStatus("error");
    };
  }, [addEvent]);

  const scheduleRetry = useCallback(() => {
    if (retryTimerRef.current) clearTimeout(retryTimerRef.current);
    const delay = Math.min(1000 * Math.pow(2, retryCountRef.current), 30000);
    retryCountRef.current += 1;
    retryTimerRef.current = setTimeout(() => {
      connectWsRef.current();
    }, delay);
  }, []);

  // Keep the stable refs current after every render
  connectWsRef.current = connectWs;
  scheduleRetryRef.current = scheduleRetry;

  // ── Initial connect ────────────────────────────────────────────────────

  useEffect(() => {
    connectWs();
    return () => {
      if (retryTimerRef.current) clearTimeout(retryTimerRef.current);
      if (speakingTimerRef.current) clearTimeout(speakingTimerRef.current);
      if (wsRef.current) {
        wsRef.current.onclose = null;
        wsRef.current.close();
      }
    };
  }, [connectWs]);

  // Reconnect when apiBaseUrl changes (skip the very first render)
  const isFirstRenderRef = useRef(true);
  useEffect(() => {
    if (isFirstRenderRef.current) {
      isFirstRenderRef.current = false;
      return;
    }
    retryCountRef.current = 0;
    connectWs();
  }, [apiBaseUrl, connectWs]);

  // ── AppState-aware reconnect on foreground ─────────────────────────────

  useEffect(() => {
    if (Platform.OS === "web") return;

    const appStateRef = { current: AppState.currentState };

    const subscription = AppState.addEventListener("change", (nextAppState) => {
      const wasBackground =
        appStateRef.current === "background" ||
        appStateRef.current === "inactive";
      const isNowActive = nextAppState === "active";

      if (wasBackground && isNowActive) {
        const ws = wsRef.current;
        if (!ws || ws.readyState !== WebSocket.OPEN) {
          if (retryTimerRef.current) {
            clearTimeout(retryTimerRef.current);
            retryTimerRef.current = null;
          }
          retryCountRef.current = 0;
          setConnectionStatus("reconnecting");
          connectWsRef.current();
        }
      } else if (nextAppState === "background") {
        if (retryTimerRef.current) {
          clearTimeout(retryTimerRef.current);
          retryTimerRef.current = null;
        }
        if (wsRef.current) {
          wsRef.current.onclose = null;
          wsRef.current.onerror = null;
          wsRef.current.close();
          wsRef.current = null;
        }
        setConnectionStatus("disconnected");
      }

      appStateRef.current = nextAppState;
    });

    return () => {
      subscription.remove();
    };
  }, []);

  // ── Outbound messages ──────────────────────────────────────────────────

  const send = useCallback((msg: object) => {
    if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(msg));
    }
  }, []);

  const sendVisitorMessage = useCallback((text: string) => {
    if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) {
      wsRef.current.send(
        JSON.stringify({
          type: "visitor_message",
          visitorId: VISITOR_ID,
          text,
        })
      );
      setTimmyMood("thinking");
    }
  }, []);

  // ── Context value ──────────────────────────────────────────────────────

  const value = useMemo<TimmyContextValue>(
    () => ({
      timmyMood,
      connectionStatus,
      isConnected: connectionStatus === "connected",
      recentEvents,
      send,
      sendVisitorMessage,
      visitorId: VISITOR_ID,
      apiBaseUrl,
      setApiBaseUrl,
    }),
    [
      timmyMood,
      connectionStatus,
      recentEvents,
      send,
      sendVisitorMessage,
      apiBaseUrl,
      setApiBaseUrl,
    ]
  );

  return (
    <TimmyContext.Provider value={value}>{children}</TimmyContext.Provider>
  );
}

export function useTimmy() {
  const ctx = useContext(TimmyContext);
  if (!ctx) throw new Error("useTimmy must be used within TimmyProvider");
  return ctx;
}