import { useQuery, useQueryClient } from '@tanstack/react-query';
import { agentApi } from '@/lib/api';
import type { UseQueryResult } from '@tanstack/react-query';
import { useState } from 'react';

export interface HeartbeatStatus {
  enabled: boolean;
  until: string | null;
  active: boolean;
  duration_minutes: number;
  source?: string | null;
}

type UseHeartbeatResult = UseQueryResult<HeartbeatStatus, Error> & {
  recentlyTriggered: boolean;
  setRecentlyTriggered: (value: boolean) => void;
};

export const useHeartbeatStatus = (agentId: string, enabled: boolean = true): UseHeartbeatResult => {
  const [recentlyTriggered, setRecentlyTriggered] = useState(false);

  const query = useQuery({
    queryKey: ['heartbeat', agentId],
    queryFn: () => agentApi.getHeartbeatStatus(agentId),
    enabled: enabled && !!agentId,
    refetchInterval: (query) => {
      // Fast polling after button click (wait for agent to report)
      if (recentlyTriggered) {
        return 5000; // 5 seconds
      }

      // Fast polling during active operations
      const heartbeatData = query.state.data as HeartbeatStatus | undefined;
      if (heartbeatData?.enabled && heartbeatData?.active) {
        return 10000; // 10 seconds
      }

      // Slow polling when idle
      return 120000; // 2 minutes
    },
    refetchOnWindowFocus: true,
  });

  // Clear flag when agent reports active
  if (recentlyTriggered && query.data?.active) {
    setRecentlyTriggered(false);
  }

  return {
    ...query,
    recentlyTriggered,
    setRecentlyTriggered,
  };
};

// Hook to manually invalidate heartbeat cache (used after commands)
export const useInvalidateHeartbeat = () => {
  const queryClient = useQueryClient();

  return (agentId: string) => {
    // Invalidate heartbeat cache
    queryClient.invalidateQueries({ queryKey: ['heartbeat', agentId] });

    // Also invalidate agent cache to synchronize data
    queryClient.invalidateQueries({ queryKey: ['agent', agentId] });
    queryClient.invalidateQueries({ queryKey: ['agents'] });
  };
};

// Hook to synchronize agent data when heartbeat status changes
export const useHeartbeatAgentSync = (agentId: string, heartbeatStatus?: HeartbeatStatus) => {
  const queryClient = useQueryClient();

  // Sync agent data when heartbeat status changes
  return () => {
    if (agentId && heartbeatStatus) {
      // Invalidate agent cache to get updated last_seen and status
      queryClient.invalidateQueries({ queryKey: ['agent', agentId] });
      queryClient.invalidateQueries({ queryKey: ['agents'] });
    }
  };
};