import { CalendarDate, getLocalTimeZone, today } from '@internationalized/date'
import { defineStore } from 'pinia'
import { computed, ref } from 'vue'

import type { NodesFilters } from '@/common/interface/filters'
import type { TracesFilters } from '@/common/interface/filters'
import type { Trace } from '@/common/interface/trace'
import { filterTraceNode, getDateRange, getDurationRange } from '@/utils/traceFilter'
import { findNodeById } from '@/utils/nodeFilter'
import { isEmptyNode } from '@/utils/emptyNode'

const initNodeFiltersGroups: string[] = ['agent', 'guard', 'llm', 'tool', 'event']

export const useTracesStore = defineStore('traces', () => {
  // Traces
  const datasetTraces = ref<Trace[]>([])
  const localTrace = ref<Trace | null>(null)

  // Filters
  const nodesFilters = ref<NodesFilters>({ name: '', group: initNodeFiltersGroups })

  const tracesFilters = ref<TracesFilters>({
    name: '',
    duration: [0, 0],
    date: {
      start: today(getLocalTimeZone()),
      end: today(getLocalTimeZone())
    }
  })

  // UI state
  const selectedTraceId = ref<string | null>(null)
  const selectedNodeId = ref<string | null>(null)
  const isLoadingList = ref<boolean>(true)
  const isLoadingTrace = ref<boolean>(true)
  // TODO: `isLoadingTrace` is currently not used, we might need to add a loading state

  const traces = computed(() => {
    return [localTrace.value, ...datasetTraces.value].filter(
      (trace): trace is Trace => trace !== null
    )
  })

  const filteredTraces = computed(() => {
    return datasetTraces.value.filter((trace) => {
      const date = new Date(trace.begin!)
      const calendarDate = new CalendarDate(date.getFullYear(), date.getMonth() + 1, date.getDate())
      const inDateRange =
        trace.begin &&
        calendarDate.compare(tracesFilters.value.date.start!) >= 0 &&
        calendarDate.compare(tracesFilters.value.date.end!) <= 0

      const inDurationRange =
        trace.duration &&
        trace.duration >= tracesFilters.value.duration[0] &&
        trace.duration <= tracesFilters.value.duration[1]

      const containsName = trace.name?.includes(tracesFilters.value.name)

      return inDurationRange && inDateRange && containsName
    })
  })

  const selectedTrace = computed(() => {
    let trace = traces.value.find((trace) => trace.id == selectedTraceId.value)

    if (!trace || !trace.parentNode) return null

    trace = {
      ...trace,
      parentNode: filterTraceNode(trace.parentNode, nodesFilters.value)
    }

    return trace
  })

  const selectedNode = computed(() => {
    if (!selectedTrace.value?.parentNode || !selectedNodeId.value) return null

    return findNodeById(selectedTrace.value.parentNode, selectedNodeId.value)
  })

  function updateDatasetTrace(trace: Trace) {
    datasetTraces.value = datasetTraces.value.map((currTrace) => {
      if (currTrace.id == trace.id) {
        return trace
      }

      return currTrace
    })
  }

  function setNodesFilters(fieldsToUpdate: Partial<NodesFilters>) {
    nodesFilters.value = { ...nodesFilters.value, ...fieldsToUpdate }
  }

  function resetNodeFilters() {
    setNodesFilters({ name: '', group: initNodeFiltersGroups })
  }

  function setTracesFilters(fieldsToUpdate: Partial<TracesFilters>) {
    tracesFilters.value = { ...tracesFilters.value, ...fieldsToUpdate }
  }

  function setupFilters() {
    const [minDate, maxDate] = getDateRange(datasetTraces.value)
    const [minDuration, maxDuration] = getDurationRange(datasetTraces.value)

    setTracesFilters({
      name: '',
      date: {
        start: new CalendarDate(minDate.getFullYear(), minDate.getMonth() + 1, minDate.getDate()),
        end: new CalendarDate(maxDate.getFullYear(), maxDate.getMonth() + 1, maxDate.getDate())
      },
      duration: [minDuration, maxDuration]
    })
  }

  return {
    datasetTraces,
    localTrace,
    selectedTraceId,
    selectedNodeId,
    updateDatasetTrace,
    nodesFilters,
    setNodesFilters,
    resetNodeFilters,
    tracesFilters,
    setTracesFilters,
    setupFilters,
    isLoadingList,
    isLoadingTrace,
    traces,
    filteredTraces,
    selectedTrace,
    selectedNode,
    isEmptyNode
  }
})
