import { useCallback, useMemo, useRef } from 'react';
import { trpc } from '@/lib/trpc';
import type { PendingQuestions, ChangeSet } from '@codewalk-district/shared';

export type RefineAgentState = 'none' | 'running' | 'waiting' | 'completed' | 'crashed';

type RefineAgent = NonNullable<ReturnType<typeof trpc.getActiveRefineAgent.useQuery>['data']>;

export interface SpawnRefineAgentOptions {
  initiativeId: string;
  instruction?: string;
}

export interface UseRefineAgentResult {
  /** Current refine agent for the initiative */
  agent: RefineAgent | null;
  /** Current state of the refine agent */
  state: RefineAgentState;
  /** Questions from the agent (when state is 'waiting') */
  questions: PendingQuestions | null;
  /** Latest applied change set (when state is 'completed') */
  changeSet: ChangeSet | null;
  /** Raw result message (when state is 'completed') */
  result: string | null;
  /** Mutation for spawning a new refine agent */
  spawn: {
    mutate: (options: SpawnRefineAgentOptions) => void;
    isPending: boolean;
    error: Error | null;
  };
  /** Mutation for resuming agent with answers */
  resume: {
    mutate: (answers: Record<string, string>) => void;
    isPending: boolean;
    error: Error | null;
  };
  /** Stop the current agent (kills process, clears questions) */
  stop: {
    mutate: () => void;
    isPending: boolean;
  };
  /** Dismiss the current agent (sets userDismissedAt so it disappears) */
  dismiss: () => void;
  /** Whether any queries are loading */
  isLoading: boolean;
  /** Function to refresh agent data */
  refresh: () => void;
}

/**
 * Hook for managing refine agents for a specific initiative.
 *
 * Encapsulates the logic for finding, spawning, and interacting with refine agents
 * that analyze and suggest improvements to initiative content.
 */
export function useRefineAgent(initiativeId: string): UseRefineAgentResult {
  const utils = trpc.useUtils();

  // Query only the active refine agent for this initiative (server-side filtered)
  const agentQuery = trpc.getActiveRefineAgent.useQuery({ initiativeId });
  const agent = agentQuery.data ?? null;

  const state: RefineAgentState = useMemo(() => {
    if (!agent) return 'none';
    switch (agent.status) {
      case 'running':
        return 'running';
      case 'waiting_for_input':
        return 'waiting';
      case 'idle':
        return 'completed';
      case 'crashed':
        return 'crashed';
      default:
        return 'none';
    }
  }, [agent]);

  // Fetch questions when waiting for input
  const questionsQuery = trpc.getAgentQuestions.useQuery(
    { id: agent?.id ?? '' },
    { enabled: state === 'waiting' && !!agent },
  );

  // Fetch change sets from DB when completed
  const changeSetsQuery = trpc.listChangeSets.useQuery(
    { agentId: agent?.id ?? '' },
    { enabled: state === 'completed' && !!agent },
  );

  // Fetch result when completed
  const resultQuery = trpc.getAgentResult.useQuery(
    { id: agent?.id ?? '' },
    { enabled: state === 'completed' && !!agent },
  );

  // Get latest applied change set
  const changeSet = useMemo(() => {
    if (!changeSetsQuery.data || changeSetsQuery.data.length === 0) return null;
    return changeSetsQuery.data.find((cs) => cs.status === 'applied') ?? null;
  }, [changeSetsQuery.data]);

  const result = useMemo(() => {
    if (!resultQuery.data?.success || !resultQuery.data.message) return null;
    return resultQuery.data.message;
  }, [resultQuery.data]);

  // Spawn mutation
  const spawnMutation = trpc.spawnArchitectRefine.useMutation({
    onMutate: async ({ initiativeId, instruction }) => {
      // Cancel outgoing refetches
      await utils.listAgents.cancel();
      await utils.getActiveRefineAgent.cancel({ initiativeId });

      // Snapshot previous values
      const previousAgents = utils.listAgents.getData();
      const previousRefineAgent = utils.getActiveRefineAgent.getData({ initiativeId });

      // Optimistically add a temporary agent
      const tempAgent = {
        id: `temp-${Date.now()}`,
        name: 'refine',
        mode: 'refine' as const,
        status: 'running' as const,
        initiativeId,
        taskId: null,
        phaseId: null,
        provider: 'claude',
        accountId: null,
        instruction: instruction || null,
        createdAt: new Date().toISOString(),
        updatedAt: new Date().toISOString(),
        userDismissedAt: null,
        completedAt: null,
      };

      utils.listAgents.setData(undefined, (old = []) => [tempAgent, ...old]);
      utils.getActiveRefineAgent.setData({ initiativeId }, tempAgent as any);

      return { previousAgents, previousRefineAgent };
    },
    onSuccess: () => {
      // Agent will appear in the list after invalidation
    },
    onError: (err, variables, context) => {
      // Revert optimistic updates
      if (context?.previousAgents) {
        utils.listAgents.setData(undefined, context.previousAgents);
      }
      if (context?.previousRefineAgent !== undefined) {
        utils.getActiveRefineAgent.setData({ initiativeId }, context.previousRefineAgent);
      }
    },
    onSettled: () => {
      void utils.listAgents.invalidate();
      void utils.getActiveRefineAgent.invalidate({ initiativeId });
    },
  });

  // Resume mutation
  const resumeMutation = trpc.resumeAgent.useMutation({
    onSuccess: () => {
      void utils.listAgents.invalidate();
    },
  });

  // Stop mutation — kills process and clears pending questions
  const stopMutation = trpc.stopAgent.useMutation({
    onSuccess: () => {
      void utils.listAgents.invalidate();
      void utils.listWaitingAgents.invalidate();
    },
  });

  // Dismiss mutation — sets userDismissedAt so agent disappears from the list
  const dismissMutation = trpc.dismissAgent.useMutation({
    onMutate: async ({ id }) => {
      // Cancel outgoing refetches
      await utils.listAgents.cancel();
      await utils.getActiveRefineAgent.cancel({ initiativeId });

      // Snapshot previous values
      const previousAgents = utils.listAgents.getData();
      const previousRefineAgent = utils.getActiveRefineAgent.getData({ initiativeId });

      // Optimistically remove the agent from the list
      utils.listAgents.setData(undefined, (old = []) =>
        old.filter(a => a.id !== id)
      );
      // Optimistically clear the active refine agent so the banner disappears immediately
      utils.getActiveRefineAgent.setData({ initiativeId }, null);

      return { previousAgents, previousRefineAgent };
    },
    onSuccess: () => {},
    onError: (err, variables, context) => {
      // Revert optimistic updates
      if (context?.previousAgents) {
        utils.listAgents.setData(undefined, context.previousAgents);
      }
      if (context?.previousRefineAgent !== undefined) {
        utils.getActiveRefineAgent.setData({ initiativeId }, context.previousRefineAgent);
      }
    },
    onSettled: () => {
      void utils.listAgents.invalidate();
      void utils.getActiveRefineAgent.invalidate({ initiativeId });
    },
  });

  // Keep mutation functions in refs so the returned spawn/resume objects are
  // stable across renders.
  const spawnMutateRef = useRef(spawnMutation.mutate);
  spawnMutateRef.current = spawnMutation.mutate;
  const agentRef = useRef(agent);
  agentRef.current = agent;
  const resumeMutateRef = useRef(resumeMutation.mutate);
  resumeMutateRef.current = resumeMutation.mutate;
  const stopMutateRef = useRef(stopMutation.mutate);
  stopMutateRef.current = stopMutation.mutate;
  const dismissMutateRef = useRef(dismissMutation.mutate);
  dismissMutateRef.current = dismissMutation.mutate;

  const spawnFn = useCallback(({ initiativeId, instruction }: SpawnRefineAgentOptions) => {
    spawnMutateRef.current({
      initiativeId,
      instruction: instruction?.trim() || undefined,
    });
  }, []);

  const spawn = useMemo(() => ({
    mutate: spawnFn,
    isPending: spawnMutation.isPending,
    error: spawnMutation.error,
  }), [spawnFn, spawnMutation.isPending, spawnMutation.error]);

  const resumeFn = useCallback((answers: Record<string, string>) => {
    const a = agentRef.current;
    if (a) {
      resumeMutateRef.current({ id: a.id, answers });
    }
  }, []);

  const resume = useMemo(() => ({
    mutate: resumeFn,
    isPending: resumeMutation.isPending,
    error: resumeMutation.error,
  }), [resumeFn, resumeMutation.isPending, resumeMutation.error]);

  const stopFn = useCallback(() => {
    const a = agentRef.current;
    if (a) {
      stopMutateRef.current({ id: a.id });
    }
  }, []);

  const stop = useMemo(() => ({
    mutate: stopFn,
    isPending: stopMutation.isPending,
  }), [stopFn, stopMutation.isPending]);

  const dismiss = useCallback(() => {
    const a = agentRef.current;
    if (a) {
      dismissMutateRef.current({ id: a.id });
    }
  }, []);

  const refresh = useCallback(() => {
    void utils.getActiveRefineAgent.invalidate({ initiativeId });
    void utils.listChangeSets.invalidate();
  }, [utils, initiativeId]);

  const isLoading = agentQuery.isLoading ||
    (state === 'waiting' && questionsQuery.isLoading) ||
    (state === 'completed' && (resultQuery.isLoading || changeSetsQuery.isLoading));

  return {
    agent,
    state,
    questions: questionsQuery.data ?? null,
    changeSet,
    result,
    spawn,
    resume,
    stop,
    dismiss,
    isLoading,
    refresh,
  };
}
