import { defineStore } from 'pinia'
import { ref, computed } from 'vue'
import {
  runRagExperiment,
  createSSEConnection,
  type AIResponse,
  type RetrievalResult,
  type RagExperimentRequest,
  type RagExperimentResult
} from '@/api/rag'

export const useRagStore = defineStore('rag', () => {
  const retrievalResults = ref<RetrievalResult[]>([])
  const finalPrompt = ref('')
  const aiResponse = ref<AIResponse | null>(null)
  const totalLatencyMs = ref<number>(0)
  
  const loading = ref(false)
  const streaming = ref(false)
  const streamContent = ref('')
  const streamError = ref<string | null>(null)
  
  const hasResults = computed(() => retrievalResults.value.length > 0 || aiResponse.value !== null)
  
  const abortStream = ref<(() => void) | null>(null)
  
  const runExperiment = async (params: RagExperimentRequest) => {
    loading.value = true
    streamError.value = null
    
    try {
      const result: RagExperimentResult = await runRagExperiment(params)
      
      retrievalResults.value = result.retrieval_results || []
      finalPrompt.value = result.final_prompt || ''
      aiResponse.value = result.ai_response || null
      totalLatencyMs.value = result.total_latency_ms || 0
      
      return result
    } catch (error: any) {
      streamError.value = error?.message || '实验运行失败'
      throw error
    } finally {
      loading.value = false
    }
  }
  
  const startStream = (params: RagExperimentRequest) => {
    streaming.value = true
    streamContent.value = ''
    streamError.value = null
    aiResponse.value = null
    
    abortStream.value = createSSEConnection(
      '/admin/rag/experiments/stream',
      params,
      (data: string) => {
        try {
          const parsed = JSON.parse(data)
          
          if (parsed.type === 'content') {
            streamContent.value += parsed.content || ''
          } else if (parsed.type === 'retrieval') {
            retrievalResults.value = parsed.results || []
          } else if (parsed.type === 'prompt') {
            finalPrompt.value = parsed.prompt || ''
          } else if (parsed.type === 'complete') {
            aiResponse.value = {
              content: streamContent.value,
              prompt_tokens: parsed.prompt_tokens,
              completion_tokens: parsed.completion_tokens,
              total_tokens: parsed.total_tokens,
              latency_ms: parsed.latency_ms,
              model: parsed.model
            }
            totalLatencyMs.value = parsed.total_latency_ms || 0
          } else if (parsed.type === 'error') {
            streamError.value = parsed.message || '流式输出错误'
          }
        } catch {
          streamContent.value += data
        }
      },
      (error: Error) => {
        streaming.value = false
        streamError.value = error.message
      },
      () => {
        streaming.value = false
      }
    )
  }
  
  const stopStream = () => {
    if (abortStream.value) {
      abortStream.value()
      abortStream.value = null
    }
    streaming.value = false
  }
  
  const clearResults = () => {
    retrievalResults.value = []
    finalPrompt.value = ''
    aiResponse.value = null
    totalLatencyMs.value = 0
    streamContent.value = ''
    streamError.value = null
  }
  
  return {
    retrievalResults,
    finalPrompt,
    aiResponse,
    totalLatencyMs,
    loading,
    streaming,
    streamContent,
    streamError,
    hasResults,
    runExperiment,
    startStream,
    stopStream,
    clearResults
  }
})
