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

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

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;
  recentEvents: WsEvent[];
  send: (msg: object) => void;
  sendVisitorMessage: (text: string) => void;
  visitorId: string;
};

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

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

function getWsUrl(): string {
  let domain = BASE_URL;
  if (!domain) {
    domain = "localhost:8080";
  }
  domain = domain.replace(/^https?:\/\//, "");
  domain = domain.replace(/\/$/, "");
  const proto = domain.startsWith("localhost") ? "ws" : "wss";
  return `${proto}://${domain}/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 [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);

  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 = getWsUrl();
    setConnectionStatus("connecting");
    let ws: WebSocket;
    try {
      ws = new WebSocket(url);
    } catch {
      setConnectionStatus("error");
      scheduleRetry();
      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");
      scheduleRetry();
    };

    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(() => {
      connectWs();
    }, delay);
  }, [connectWs]);

  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]);

  // AppState-aware WebSocket 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) {
        // App returned to foreground — check if WS is still alive
        const ws = wsRef.current;
        if (!ws || ws.readyState !== WebSocket.OPEN) {
          // Cancel any pending retry so we don't create duplicates
          if (retryTimerRef.current) {
            clearTimeout(retryTimerRef.current);
            retryTimerRef.current = null;
          }
          retryCountRef.current = 0;
          setConnectionStatus("reconnecting");
          connectWs();
        }
      } else if (nextAppState === "background") {
        // Proactively close the WS to avoid OS killing it mid-frame
        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();
    };
  }, [connectWs]);

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

  const value = useMemo(
    () => ({
      timmyMood,
      connectionStatus,
      recentEvents,
      send,
      sendVisitorMessage,
      visitorId: VISITOR_ID,
    }),
    [timmyMood, connectionStatus, recentEvents, send, sendVisitorMessage]
  );

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