frontend/src/hooks/useObservability.ts

"use client";

import { useState, useEffect, useMemo, useRef } from "react";
import { Message } from "./useConversation";

// ---- Types ----

export interface ToolCallEntry {
  id: string;
  name: string;
  startedAt: number;       // Date.now()
  duration: number | null;  // ms, null if still running
  status: "started" | "completed" | "error";
}

export interface ServerMetrics {
  connected_clients: number;
  session_id: string | null;
  server_start_time: string;
  uptime_seconds: number;
  total_events_broadcast: number;
}

export interface CostMetrics {
  input_text: number;
  input_audio: number;
  output_text: number;
  output_audio: number;
  cached_text: number;
  cached_audio: number;
  total_cost: number;
  response_count: number;
  cache_hit_rate: number;
}

export interface ObservabilityData {
  // Message counts
  userMessageCount: number;
  assistantMessageCount: number;
  toolMessageCount: number;

  // Tool timeline
  toolCalls: ToolCallEntry[];

  // Events per second (rolling 10s window)
  eventsPerSecond: number;

  // GenUI activity
  genuiComponentCount: number;
  genuiComponentNames: string[];

  // Session uptime (seconds, null if no session)
  sessionUptimeSeconds: number | null;

  // Server metrics (polled)
  serverMetrics: ServerMetrics | null;

  // Cost metrics (polled from /api/stream/cost-summary)
  costMetrics: CostMetrics | null;

  // Total WebSocket events received this session
  totalWsEvents: number;
}

// ---- Hook ----

interface UseObservabilityOptions {
  messages: Message[];
  isSessionActive: boolean;
  apiBaseUrl?: string;
}

export function useObservability({
  messages,
  isSessionActive,
  apiBaseUrl = "http://localhost:7860",
}: UseObservabilityOptions): ObservabilityData {
  // ---- Message counts (derived from messages) ----
  const userMessageCount = messages.filter((m) => m.role === "user" && !m.isPartial).length;
  const assistantMessageCount = messages.filter((m) => m.role === "assistant").length;
  const toolMessageCount = messages.filter((m) => m.role === "tool").length;

  // ---- GenUI activity ----
  const genuiMessages = messages.filter((m) => m.component);
  const genuiComponentCount = genuiMessages.length;
  const genuiComponentNames = [...new Set(genuiMessages.map((m) => m.component!.name))];

  // ---- Tool call timeline ----
  const toolCalls = useMemo(() => {
    const toolMessages = messages.filter((m) => m.role === "tool" && m.tool);
    return toolMessages.map((m) => ({
      id: m.id,
      name: m.tool!.name,
      startedAt: new Date(m.timestamp).getTime(),
      duration: null,  // Real duration not available from message data
      status: m.tool!.status,
    })).slice(-10) as ToolCallEntry[];
  }, [messages]);

  // ---- Events per second (rolling window) ----
  const eventTimestampsRef = useRef<number[]>([]);
  const [eventsPerSecond, setEventsPerSecond] = useState(0);
  const [totalWsEvents, setTotalWsEvents] = useState(0);

  // Count every message change as events
  const prevMessageCountRef = useRef(0);
  useEffect(() => {
    const currentCount = messages.length;
    if (currentCount > prevMessageCountRef.current) {
      const newEvents = currentCount - prevMessageCountRef.current;
      const now = Date.now();
      for (let i = 0; i < newEvents; i++) {
        eventTimestampsRef.current.push(now);
      }
      setTotalWsEvents((prev) => prev + newEvents);
    }
    prevMessageCountRef.current = currentCount;
  }, [messages]);

  // Update events/sec every second
  useEffect(() => {
    const interval = setInterval(() => {
      const now = Date.now();
      const windowMs = 10_000; // 10-second rolling window
      eventTimestampsRef.current = eventTimestampsRef.current.filter(
        (t) => now - t < windowMs
      );
      setEventsPerSecond(
        Math.round((eventTimestampsRef.current.length / (windowMs / 1000)) * 10) / 10
      );
    }, 1000);
    return () => clearInterval(interval);
  }, []);

  // ---- Session uptime ----
  const sessionStartRef = useRef<number | null>(null);
  const [sessionUptimeSeconds, setSessionUptimeSeconds] = useState<number | null>(null);

  useEffect(() => {
    if (isSessionActive) {
      if (!sessionStartRef.current) {
        sessionStartRef.current = Date.now();
      }
    } else {
      sessionStartRef.current = null;
      // Reset is handled via the tick interval below
    }
  }, [isSessionActive]);

  // Tick session uptime every second
  useEffect(() => {
    if (!isSessionActive) return;
    const interval = setInterval(() => {
      if (sessionStartRef.current) {
        setSessionUptimeSeconds(
          Math.floor((Date.now() - sessionStartRef.current) / 1000)
        );
      }
    }, 1000);
    return () => clearInterval(interval);
  }, [isSessionActive]);

  // ---- Server metrics (via WebSocket, fallback initial fetch) ----
  const [serverMetrics, setServerMetrics] = useState<ServerMetrics | null>(null);

  useEffect(() => {
    // Initial fetch for data before WebSocket connects
    const fetchMetrics = async () => {
      try {
        const res = await fetch(`${apiBaseUrl}/api/stream/observability`);
        if (res.ok) {
          const data = await res.json();
          setServerMetrics(data);
        }
      } catch {
        // Silently ignore — server may be down
      }
    };
    fetchMetrics();

    // Listen for WS-pushed observability snapshots
    const handleMetrics = (e: Event) => {
      const detail = (e as CustomEvent).detail;
      if (detail) {
        setServerMetrics({
          connected_clients: detail.connected_clients,
          session_id: detail.session_id,
          server_start_time: detail.server_start_time,
          uptime_seconds: detail.uptime_seconds,
          total_events_broadcast: detail.total_events_broadcast,
        });
      }
    };
    window.addEventListener("observability-metrics", handleMetrics);
    return () => window.removeEventListener("observability-metrics", handleMetrics);
  }, [apiBaseUrl]);

  // ---- Cost metrics (via WebSocket, fallback initial fetch) ----
  const [costMetrics, setCostMetrics] = useState<CostMetrics | null>(null);

  useEffect(() => {
    // Initial fetch for data before WebSocket connects
    const fetchCost = async () => {
      try {
        const res = await fetch(`${apiBaseUrl}/api/stream/cost-summary`);
        if (res.ok) {
          const data = await res.json();
          setCostMetrics(data);
        }
      } catch {
        // Silently ignore — server may be down
      }
    };
    fetchCost();

    // Listen for WS-pushed cost updates with aggregated summary
    const handleCost = (e: Event) => {
      const detail = (e as CustomEvent).detail;
      const summary = detail?.summary;
      if (summary) {
        setCostMetrics({
          input_text: summary.input_text ?? 0,
          input_audio: summary.input_audio ?? 0,
          output_text: summary.output_text ?? 0,
          output_audio: summary.output_audio ?? 0,
          cached_text: summary.cached_text ?? 0,
          cached_audio: summary.cached_audio ?? 0,
          total_cost: summary.total_cost ?? 0,
          response_count: summary.response_count ?? 0,
          cache_hit_rate: summary.cache_hit_rate ?? 0,
        });
      }
    };
    window.addEventListener("cost-updated", handleCost);
    return () => window.removeEventListener("cost-updated", handleCost);
  }, [apiBaseUrl]);

  return {
    userMessageCount,
    assistantMessageCount,
    toolMessageCount,
    toolCalls,
    eventsPerSecond,
    genuiComponentCount,
    genuiComponentNames,
    sessionUptimeSeconds,
    serverMetrics,
    costMetrics,
    totalWsEvents,
  };
}