frontend/src/hooks/useGeneration.js

import { useState, useCallback, useRef, useEffect } from 'react'

const DEFAULT_PARAMS = {
  musicgen: {
    prompt: '',
    duration: 15,
    seed: -1,
  },
  'ace-step': {
    task_type: 'repaint',
    caption: '',
    lyrics: '[Instrumental]',
    bpm: null,
    duration: 30,
    audio_cover_strength: 0.8,
    model_variant: 'xl-turbo',
    inference_steps: 15,
    seed: -1,
    scheduler: 'ode',
    cfg_scale: null,
    timestep_shift: null,
  },
}

export function useGeneration({
  sessionId,
  generate,
  extractSeed,
  generateAceStep,
  checkModels,
}) {
  const [selectedModel, setSelectedModel] = useState('musicgen')
  const [modelParams, setModelParams] = useState(() => JSON.parse(JSON.stringify(DEFAULT_PARAMS)))
  const [isGenerating, setIsGenerating] = useState(false)
  const [generationResult, setGenerationResult] = useState(null)
  const [generationProgress, setGenerationProgress] = useState(null)
  const [modelAvailability, setModelAvailability] = useState(null)
  const [seedData, setSeedData] = useState(null)
  const abortRef = useRef(null)

  // Check model availability and clear stale results when session changes
  useEffect(() => {
    // Clear previous generation when session changes
    setGenerationResult(null)
    setSeedData(null)
    setGenerationProgress(null)

    if (checkModels) {
      checkModels().then(data => {
        if (data?.models) {
          const avail = {}
          data.models.forEach(m => { avail[m.name] = m })
          setModelAvailability(avail)
        }
      })
    }
  }, [sessionId, checkModels])

  const setModelParam = useCallback((model, key, value) => {
    setModelParams(prev => ({
      ...prev,
      [model]: { ...prev[model], [key]: value },
    }))
  }, [])

  const resetParams = useCallback((model) => {
    setModelParams(prev => ({
      ...prev,
      [model]: JSON.parse(JSON.stringify(DEFAULT_PARAMS[model])),
    }))
  }, [])

  // Extract seed from region
  const doExtractSeed = useCallback(async (regionStart, regionEnd) => {
    if (!extractSeed) return null
    const result = await extractSeed(regionStart, regionEnd)
    if (result) {
      setSeedData(result)
      // Pre-fill BPM from analysis
      if (result.analysis?.bpm) {
        setModelParams(prev => ({
          ...prev,
          'ace-step': { ...prev['ace-step'], bpm: result.analysis.bpm },
        }))
      }
    }
    return result
  }, [extractSeed])

  // Start generation
  const startGeneration = useCallback(async (regionStart, regionEnd) => {
    setIsGenerating(true)
    setGenerationResult(null)
    setGenerationProgress({ status: 'starting' })

    try {
      let result = null

      if (selectedModel === 'musicgen') {
        const params = modelParams.musicgen
        result = await generate(regionStart, regionEnd, params.duration, params.prompt || null)
      } else if (selectedModel === 'ace-step') {
        if (!seedData?.seed_id) {
          // Auto-extract seed if not done
          const extracted = await doExtractSeed(regionStart, regionEnd)
          if (!extracted) throw new Error('Failed to extract seed')
        }
        const params = { ...modelParams['ace-step'] }
        params.seed_id = seedData.seed_id
        result = await generateAceStep(params)
      }

      if (result) {
        setGenerationResult({ ...result, model: selectedModel })
        setGenerationProgress({ status: 'done' })
      } else {
        setGenerationProgress({ status: 'error' })
      }

      return result
    } catch (err) {
      setGenerationProgress({ status: 'error', message: err.message })
      return null
    } finally {
      setIsGenerating(false)
    }
  }, [selectedModel, modelParams, seedData, generate, generateAceStep, doExtractSeed])

  const cancelGeneration = useCallback(() => {
    if (abortRef.current) {
      abortRef.current.abort()
      abortRef.current = null
    }
    setIsGenerating(false)
    setGenerationProgress(null)
  }, [])

  return {
    selectedModel,
    setSelectedModel,
    modelParams,
    setModelParam,
    resetParams,
    isGenerating,
    generationResult,
    setGenerationResult,
    generationProgress,
    modelAvailability,
    seedData,
    setSeedData,
    doExtractSeed,
    startGeneration,
    cancelGeneration,
  }
}