import { useState, useEffect, useMemo, useCallback } from 'react'
import type { ACPClient } from '../acp/client'
import type { ModelInfo, SessionModelState } from '../acp/types'

export interface UseModelsResult {
  /** Whether model selection is supported by the current agent */
  supportsModelSelection: boolean
  /** List of available models */
  availableModels: ModelInfo[]
  /** The currently selected model ID */
  currentModelId: string | null
  /** The currently selected model info */
  currentModel: ModelInfo | null
  /** Set the model for the current session */
  setModel: (modelId: string) => Promise<void>
  /** Whether a model change is in progress */
  isLoading: boolean
}

/**
 * Hook to manage model selection state.
 * Reference: Zed's AcpModelSelector reads from state.available_models and state.current_model_id
 *
 * Uses event-driven updates instead of polling:
 * - setModelStateChangedHandler: called on session create/disconnect
 * - setModelChangedHandler: called when model selection changes
 */
export function useModels(client: ACPClient): UseModelsResult {
  const [modelState, setModelState] = useState<SessionModelState | null>(
    client.modelState,
  )
  const [isLoading, setIsLoading] = useState(false)

  // Subscribe to model state changes (session created/destroyed)
  // This replaces the previous 500ms polling approach
  useEffect(() => {
    // Handler for when model state changes (session created or disconnected)
    const handleModelStateChanged = (state: SessionModelState | null) => {
      setModelState(state)
      // Auto-restore previously selected model when a new session is created
      if (state && state.availableModels.length > 0) {
        const saved = localStorage.getItem('acp_model_id')
        if (
          saved &&
          saved !== state.currentModelId &&
          state.availableModels.some(m => m.modelId === saved)
        ) {
          client.setSessionModel(saved).catch(() => {})
        }
      }
    }

    // Handler for when current model changes within a session
    const handleModelChanged = (modelId: string) => {
      setModelState(prev => {
        if (!prev) return null
        return {
          ...prev,
          currentModelId: modelId,
        }
      })
      setIsLoading(false)
    }

    // Register handlers - setModelStateChangedHandler immediately calls with current state
    client.setModelStateChangedHandler(handleModelStateChanged)
    client.setModelChangedHandler(handleModelChanged)

    return () => {
      // Clear handlers on unmount
      client.setModelStateChangedHandler(() => {})
      client.setModelChangedHandler(() => {})
    }
  }, [client])

  const availableModels = useMemo(
    () => modelState?.availableModels ?? [],
    [modelState],
  )

  const currentModelId = modelState?.currentModelId ?? null

  const currentModel = useMemo(
    () => availableModels.find(m => m.modelId === currentModelId) ?? null,
    [availableModels, currentModelId],
  )

  const setModel = useCallback(
    async (modelId: string) => {
      if (!modelState) {
        throw new Error('Model selection not supported')
      }
      setIsLoading(true)
      try {
        await client.setSessionModel(modelId)
        localStorage.setItem('acp_model_id', modelId)
        // The model_changed event will update the state
      } catch (error) {
        setIsLoading(false)
        throw error
      }
    },
    [client, modelState],
  )

  return {
    supportsModelSelection: modelState !== null && availableModels.length > 0,
    availableModels,
    currentModelId,
    currentModel,
    setModel,
    isLoading,
  }
}