import request from '../utils/request'

const RAG_TIMEOUT = 90000

const unwrap = (promise) => promise.then(response => response.data)

export function retrieveDocumentEvidence(query, topK = 5, mode = 'hybrid') {
  return unwrap(request.post('/rag/document/retrieve', { query, topK, filters: { mode } }, { timeout: RAG_TIMEOUT }))
}

export function retrieveStructuredEvidence(query, sourceIds = ['1'], topK = 5) {
  return unwrap(request.post('/rag/structured/retrieve', {
    query, sourceIds, topK, filters: { allowTextToSql: true }
  }, { timeout: RAG_TIMEOUT }))
}

export function retrieveGraphEvidence(query, sourceIds = ['1'], topK = 5) {
  return unwrap(request.post('/rag/graph/retrieve', {
    query, sourceIds, topK, filters: { allowTextToCypher: true }
  }, { timeout: RAG_TIMEOUT }))
}

export function retrieveKnowledgeBaseEvidence(params) {
  return unwrap(request.post('/rag/kb/retrieve', params, { timeout: RAG_TIMEOUT }))
}

export function getRagCapability(sourceType, sourceId) {
  return unwrap(request.get(`/rag/capabilities/${sourceType}/${sourceId}`))
}

export function refreshRagCapability(sourceType, sourceId) {
  return unwrap(request.post(`/rag/capabilities/${sourceType}/${sourceId}/refresh`))
}

export function suggestRagGovernance(sourceType, sourceId, description = '') {
  return unwrap(request.post(`/rag/capabilities/${sourceType}/${sourceId}/suggest`, { description }, { timeout: RAG_TIMEOUT }))
}

export function listRagQueryExamples(sourceType, sourceId) {
  return unwrap(request.get(`/rag/capabilities/${sourceType}/${sourceId}/examples`))
}

export function addRagQueryExampleCandidates(sourceType, sourceId, candidates) {
  return unwrap(request.post(`/rag/capabilities/${sourceType}/${sourceId}/examples/candidates`, candidates))
}

export function verifyRagQueryExample(id, verified) {
  return unwrap(request.put(`/rag/capabilities/examples/${id}/verified`, null, { params: { value: verified } }))
}

export function deleteRagQueryExample(id) {
  return unwrap(request.delete(`/rag/capabilities/examples/${id}`))
}

export function getKnowledgeBaseBindings(knowledgeBaseId) {
  return unwrap(request.get(`/rag/kb/${knowledgeBaseId}/bindings`))
}

export function listKnowledgeBases() {
  return unwrap(request.get('/rag/kb'))
}

export function updateKnowledgeBaseBindingConfig(knowledgeBaseId, bindingId, config) {
  return unwrap(request.put(`/rag/kb/${knowledgeBaseId}/bindings/${bindingId}/config`, config))
}

export function upsertKnowledgeBaseBindingBySource(knowledgeBaseId, sourceType, sourceId, config) {
  return unwrap(request.put(`/rag/kb/${knowledgeBaseId}/bindings/by-source`, { sourceType, sourceId, config }))
}

export function generateRagAnswer(question, evidences) {
  return unwrap(request.post('/rag/answer', { question, evidences }, { timeout: RAG_TIMEOUT }))
}

export async function sendRagQuestion(params) {
  const { query, topK = 5, strategy = 'hybrid', enabled, onStep } = params
  const results = { document: null, structured: null, graph: null, kb: null }
  const tasks = []
  const usesKnowledgeBase = enabled.structured || enabled.graph

  if (enabled.document) {
    onStep?.('document', 'running')
    tasks.push(retrieveDocumentEvidence(query, topK, strategy)
      .then(data => { results.document = data; onStep?.('document', 'success', data) })
      .catch(error => { onStep?.('document', 'failed', null, error); results.document = { evidences: [], diagnostics: { error: error.message } } }))
  }

  if (usesKnowledgeBase) {
    if (enabled.structured) onStep?.('structured', 'running')
    if (enabled.graph) onStep?.('graph', 'running')
    onStep?.('kb', 'running')
    tasks.push(retrieveKnowledgeBaseEvidence({
      knowledgeBaseId: 1,
      query,
      topK,
      filters: {
        sourceTypes: [enabled.structured && 'STRUCTURED_DATA', enabled.graph && 'GRAPH'].filter(Boolean),
        structured: { allowTextToSql: true },
        graph: { allowTextToCypher: true }
      }
    }).then(data => {
      results.kb = data
      const diagnosticsFor = prefix => Object.fromEntries(Object.entries(data.diagnostics || {}).filter(([key]) => key.startsWith(`${prefix}:`)))
      results.structured = { sourceType: 'STRUCTURED_DATA', evidences: data.evidences?.filter(e => e.sourceType === 'STRUCTURED_DATA') || [], diagnostics: diagnosticsFor('STRUCTURED_DATA') }
      results.graph = { sourceType: 'GRAPH', evidences: data.evidences?.filter(e => e.sourceType === 'GRAPH') || [], diagnostics: diagnosticsFor('GRAPH') }
      if (enabled.structured) {
        const failed = Object.keys(results.structured.diagnostics).length > 0
        onStep?.('structured', failed ? 'failed' : 'success', results.structured, failed ? new Error(JSON.stringify(results.structured.diagnostics)) : null)
      }
      if (enabled.graph) {
        const failed = Object.keys(results.graph.diagnostics).length > 0
        onStep?.('graph', failed ? 'failed' : 'success', results.graph, failed ? new Error(JSON.stringify(results.graph.diagnostics)) : null)
      }
      onStep?.('kb', 'success', data)
    }).catch(error => {
      if (enabled.structured) onStep?.('structured', 'failed', null, error)
      if (enabled.graph) onStep?.('graph', 'failed', null, error)
      onStep?.('kb', 'failed', null, error)
      results.kb = { evidences: [], diagnostics: { error: error.message } }
    }))
  }

  await Promise.all(tasks)
  const evidences = [results.document, results.structured, results.graph]
    .flatMap(result => result?.evidences || [])
  return { ...results, evidences }
}