<template>
  <div class="modal-overlay">
    <div
      v-if="isOpen"
      ref="modal"
      class="modal fixed right-10 flex flex-col border-1 border-secondary shadow-lg"
      :class="{ fullscreen: isFullscreen }"
      :style="modalStyle"
    >
      <div
        class="flex flex-none items-center justify-between pl-2 p-1 gap-3 rounded-t-md duration-300"
        :style="`background-color: ${getGroupColor(classifyNodeGroup(tracesStore.selectedNode))};`"
        ref="dragHandle"
        @dblclick="toggleFullscreen"
      >
        <slot name="header">
          <h3
            class="font-mono truncate"
            :class="{
              'text-white': ['event', 'llm'].includes(classifyNodeGroup(tracesStore.selectedNode))
            }"
          >
            {{ tracesStore.selectedNode?.name }}
          </h3>
        </slot>
        <div class="flex gap-1 items-center">
          <Tabs v-model="formatModel">
            <TabsList class="gap-1 rounded-md h-8">
              <TabsTrigger
                :data-grp="classifyNodeGroup(tracesStore.selectedNode)"
                class="header-tab"
                value="formatted"
              >
                Formatted
              </TabsTrigger>
              <TabsTrigger
                :data-grp="classifyNodeGroup(tracesStore.selectedNode)"
                class="header-tab"
                value="json"
              >
                JSON
              </TabsTrigger>
            </TabsList>
          </Tabs>
          <Tooltip :delay-duration="500">
            <TooltipTrigger as-child>
              <div
                class="w-8 h-8 rounded-md bg-white flex items-center justify-center cursor-pointer"
                @click="toggleFullscreen"
              >
                <Fullscreen :size="18" />
              </div>
            </TooltipTrigger>
            <TooltipContent side="bottom" class="z-[2000]">Toggle fullscreen</TooltipContent>
          </Tooltip>
          <Tooltip :delay-duration="500">
            <TooltipTrigger as-child>
              <div
                class="w-8 h-8 rounded-md bg-white flex items-center justify-center cursor-pointer"
                @click="close"
              >
                <X :size="18" />
              </div>
            </TooltipTrigger>
            <TooltipContent side="bottom" class="z-[2000]">Close</TooltipContent>
          </Tooltip>
        </div>
      </div>
      <div class="flex">
        <div class="flex-col gap-1 m-1">
          <div class="h-10 flex gap-1 p-1 overflow-auto">
            <Tooltip :delay-duration="500">
              <TooltipTrigger as-child>
                <div
                  class="flex px-3 rounded-md items-center border-1 font-light text-sm cursor-help bg-gray-50 text-gray-600"
                >
                  <Timer :size="16" class="mr-1" />
                  <span class="font-mono">{{ formattedTimes.duration }}s</span>
                  <Separator orientation="vertical" class="mx-3" />
                  <Calendar1 :size="16" class="mr-1" />
                  <span class="font-mono">{{ formattedTimes.beginTime }}</span>
                  <Separator orientation="vertical" class="mx-3" />
                  <ArrowRight :size="16" class="mr-1" />
                  <span class="font-mono">{{ formattedTimes.endTime }}</span>
                </div>
              </TooltipTrigger>
              <TooltipContent side="bottom" class="text-black shadow-md" :arrow-display="true">
                <div class="flex gap-6 text-sm font-light text-white">
                  <div class="flex-column gap-2">
                    <p>Duration (seconds)</p>
                    <p>Start Time</p>
                    <p>End Time</p>
                  </div>
                  <div class="flex-column gap-2 font-mono font-thin align-end">
                    <p>{{ formattedTimes.duration }}s</p>
                    <p>
                      {{ formattedTimes.beginTime }}
                    </p>
                    <p>
                      {{ formattedTimes.endTime }}
                    </p>
                  </div>
                </div>
              </TooltipContent>
            </Tooltip>
            <Tooltip :delay-duration="500">
              <TooltipTrigger as-child>
                <div
                  class="flex px-3 rounded-md items-center border-1 font-light text-sm cursor-help bg-gray-50 text-gray-600"
                >
                  <MessageSquareText :size="16" class="mr-1" />
                  <span class="font-mono">{{ formattedAttributes.promptTokens }}</span>

                  <div class="h-full flex px-3">
                    <Separator orientation="vertical" class="" />
                    <Plus
                      class="absolute -translate-x-[7.5px] translate-y-[7px] bg-white rounded-full border"
                      :size="16"
                    />
                  </div>

                  <BrainCircuit :size="16" class="mr-1" />
                  <span class="font-mono">{{ formattedAttributes.completionTokens }}</span>

                  <div class="h-full flex px-3">
                    <Separator orientation="vertical" class="" />
                    <Equal
                      class="absolute -translate-x-[7.5px] translate-y-[7px] bg-white rounded-full border"
                      :size="16"
                    />
                  </div>

                  <Sigma :size="16" class="mr-1" />
                  <span class="font-mono">{{ formattedAttributes.totalTokens }}</span>
                </div>
              </TooltipTrigger>
              <TooltipContent side="bottom" class="text-black shadow-md" :arrow-display="true">
                <div class="flex gap-6 text-sm font-light text-white">
                  <div class="flex-column gap-2">
                    <p>Prompt tokens</p>
                    <p>Completion tokens</p>
                    <p>Total tokens</p>
                  </div>
                  <div class="flex-column gap-2 font-mono font-thin align-end">
                    <div class="flex justify-end">
                      {{ formattedAttributes.promptTokens }}
                    </div>
                    <div class="flex items-center justify-between">
                      <Plus :size="16" class="mr-2" />
                      {{ formattedAttributes.completionTokens }}
                    </div>
                    <div class="flex items-center justify-between">
                      <Equal :size="16" class="mr-2" />
                      {{ formattedAttributes.totalTokens }}
                    </div>
                  </div>
                </div>
              </TooltipContent>
            </Tooltip>
            <Tooltip :delay-duration="500">
              <TooltipTrigger as-child>
                <div
                  class="flex px-3 rounded-md items-center gap-1 border-1 font-light text-sm cursor-help bg-gray-50 text-gray-600"
                >
                  <DollarSign :size="16" />
                  <span class="ml-1 font-mono mt-[2px]">{{
                    formattedAttributes.estimatedCost
                  }}</span>
                </div>
              </TooltipTrigger>
              <TooltipContent side="bottom" class="text-black shadow-md" :arrow-display="true">
                <div class="flex gap-6 text-sm font-light text-white">
                  <div class="flex-column gap-2">
                    <p>Estimated cost</p>
                  </div>
                  <div class="flex-column gap-2 font-mono font-thin">
                    <p>${{ formattedAttributes.estimatedCost }}</p>
                  </div>
                </div>
              </TooltipContent>
            </Tooltip>
          </div>
          <div class="flex-1 flex gap-2 p-1 items-center overflow-auto">
            <div class="relative w-10 min-w-sm items-center">
              <Input
                id="search"
                type="text"
                placeholder="Search in node"
                class="pl-9 shadow-none"
                v-model="searchTerm"
              />
              <span class="absolute start-0 inset-y-0 flex items-center justify-center pl-3">
                <Search class="text-muted-foreground" :size="16" />
              </span>
              <span
                v-if="searchTerm !== ''"
                class="absolute end-0 inset-y-0 flex items-center justify-center pr-3 cursor-pointer"
              >
                <X @click="searchTerm = ''" class="text-muted-foreground" :size="16" />
              </span>
            </div>
            <div class="flex gap-1">
              <Button
                v-for="(keyTerm, TermIdx) in quickDisplayButtons"
                :key="TermIdx"
                variant="outline"
                @click="searchTerm = keyTerm"
                >{{ keyTerm }}</Button
              >
            </div>
          </div>
        </div>
        <QuickNavigate class="ml-auto m-1" v-if="isFullscreen" />
      </div>
      <div class="h-full min-h-0 flex flex-col">
        <JsonPanel
          v-if="formatModel === FormatMode.Json"
          :is-fullscreen="isFullscreen"
          :display-node="displayNode"
        />
        <FormattedPanel
          v-else-if="formatModel === FormatMode.Formatted"
          :is-fullscreen="isFullscreen"
          :display-node="displayNode"
        />
      </div>
    </div>
  </div>
</template>

<script setup lang="ts">
import type { Interactable } from '@interactjs/types'
import interact from 'interactjs'
import {
  Fullscreen,
  X,
  BrainCircuit,
  DollarSign,
  Equal,
  MessageSquareText,
  Plus,
  Search,
  Sigma,
  Calendar1,
  Timer,
  ArrowRight
} from 'lucide-vue-next'
import { ref, computed, onMounted, onBeforeUnmount, type Ref, type ComputedRef, watch } from 'vue'

import { Tooltip, TooltipContent, TooltipTrigger } from '@/components/ui/tooltip'
import { classifyNodeGroup, getGroupColor } from '@/utils/nodeGroups'
import { useTracesStore } from '@/stores/traces'
import { Button } from '@/components/ui/button'

import { Input } from '@/components/ui/input'
import QuickNavigate from '@/components/node-explorer/NodeQuickNavigate.vue'
import { localizeNumber } from '@/utils/formatNumber'
import Separator from '@/components/ui/separator/Separator.vue'
import type { TraceNode } from '@/common/interface/trace'
import { FormatMode } from '@/common/views'
import JsonPanel from '@/components/node-explorer/JsonPanel.vue'
import FormattedPanel from '@/components/node-explorer/FormattedPanel.vue'
import { Tabs, TabsList, TabsTrigger } from '@/components/ui/tabs'
import { dateToTimeOnly } from '@/utils/formatDate'
import { isEmptyPanel, FormatPanelType } from '@/utils/emptyNode'

const tracesStore = useTracesStore()

const isOpen = ref(true)
const isFullscreen = ref(false)
const modal = ref<HTMLElement | null>(null)
const dragHandle = ref<HTMLElement | null>(null)
const position = ref({ x: 0, y: 0 })
const searchTerm: Ref<string> = ref('')
const quickDisplayButtons: string[] = ['text', 'description', 'token']

const formatSelection = ref<FormatMode | ''>(FormatMode.Formatted)

const formatModel = computed<FormatMode | ''>({
  // This is to avoid nothing being selected.
  get: () => formatSelection.value,
  set: (v) => {
    if (v) formatSelection.value = v
  }
})

const modalStyle = computed(() => ({
  transform: `translate(${position.value.x}px, ${position.value.y}px)`
}))

onMounted(() => {
  initializeDrag()
  setPanelMode()
})

onBeforeUnmount(() => {
  if (modal.value) {
    interact(modal.value).unset()
  }
})

function setPanelMode() {
  const shouldFallbackToJsonPanel =
    isEmptyPanel(displayNode.value?.inputs, FormatPanelType.Inputs) &&
    isEmptyPanel(displayNode.value?.outputs, FormatPanelType.Outputs)
  formatSelection.value = shouldFallbackToJsonPanel ? FormatMode.Json : FormatMode.Formatted
}

watch(
  () => tracesStore.selectedNodeId,
  () => {
    setPanelMode()
  }
)

const close = () => {
  isOpen.value = false
  tracesStore.selectedNodeId = null
}

const initializeDrag = () => {
  if (!modal.value || !dragHandle.value) return

  interact(modal.value).draggable({
    listeners: {
      move(event) {
        position.value.x += event.dx
        position.value.y += event.dy
      }
    },
    modifiers: [
      interact.modifiers.restrictRect({
        restriction: 'parent',
        endOnly: true
      })
    ],
    inertia: true,
    allowFrom: dragHandle.value
  }) as Interactable
}

const toggleFullscreen = () => {
  isFullscreen.value = !isFullscreen.value
  if (isFullscreen.value) {
    position.value = { x: 0, y: 0 } // Reset position in fullscreen
  }
}

type AnyObj = Record<string, any>

function searchSelectedNode(term: string) {
  if (term == null) return undefined
  const needle = String(term).toLowerCase()
  const normalise = (v: unknown) => String(v).toLowerCase()
  const seen = new WeakSet<object>()
  const visit = (node: any): any | undefined => {
    if (node === null || typeof node !== 'object') {
      return normalise(node).includes(needle) ? node : undefined
    }

    if (seen.has(node)) return undefined
    seen.add(node)

    if (Array.isArray(node)) {
      const out: any[] = []
      for (const item of node) {
        const child = visit(item)
        if (child !== undefined) out.push(child)
      }
      return out.length ? out : undefined
    }

    const out: AnyObj = {}
    for (const [key, value] of Object.entries(node)) {
      const keyMatch = key.toLowerCase().includes(needle)
      const child = visit(value)

      const valueMatch =
        value !== null && typeof value === 'object'
          ? child !== undefined
          : normalise(value).includes(needle)

      if (keyMatch || valueMatch) {
        out[key] = child !== undefined ? child : value
        if (
          // If the term is in the `text` we still want to be able to display the `role`
          key === 'text' &&
          Object.prototype.hasOwnProperty.call(node, 'role') &&
          out.role === undefined
        ) {
          out.role = (node as AnyObj).role
        } else if (
          // If the term is in the `role` we still want to be able to display the `text`
          key === 'role' &&
          Object.prototype.hasOwnProperty.call(node, 'text') &&
          out.text === undefined
        ) {
          out.text = (node as AnyObj).text
        }
      } else if (child !== undefined) {
        out[key] = child
      }
    }

    return Object.keys(out).length ? out : undefined
  }
  return visit(tracesStore.selectedNode)
}

const displayNode: ComputedRef<TraceNode | undefined> = computed(() => {
  if (!searchTerm.value) {
    return tracesStore.selectedNode
  } else {
    return searchSelectedNode(searchTerm.value)
  }
})

const formattedAttributes = computed(() => {
  const usageMetadata = tracesStore.selectedNode?.usageMetadata
  const completionResponse = tracesStore.selectedNode?.attributes.completionResponse
  return {
    promptTokens: localizeNumber(
      usageMetadata?.promptTokens ?? completionResponse?.promptTokens ?? 0
    ),
    completionTokens: localizeNumber(
      usageMetadata?.completionTokens ?? completionResponse?.completionTokens ?? 0
    ),
    totalTokens: localizeNumber(usageMetadata?.totalTokens ?? completionResponse?.totalTokens ?? 0),
    estimatedCost: localizeNumber(
      usageMetadata?.estimatedCost ?? completionResponse?.estimatedCost ?? 0
    )
  }
})

const formattedTimes = computed(() => {
  return {
    duration: ((tracesStore.selectedNode?.duration || 0) / 1000).toFixed(2),
    beginTime: dateToTimeOnly(tracesStore.selectedNode?.begin),
    endTime: dateToTimeOnly(tracesStore.selectedNode?.end)
  }
})
</script>

<style scoped>
.modal-overlay {
  position: fixed;
  top: 0;
  left: 0;
  right: 0;
  bottom: 0;
  background: transparent;
  display: flex;
  justify-content: center;
  align-items: center;
  /* z-index needs to be lower than 1000 so that Select displays correctly */
  z-index: 10;
  pointer-events: none;
}

.modal {
  background: white;
  border-radius: 8px;
  width: 800px;
  max-width: 100vw;
  height: 750px;
  max-height: 100vh;
  pointer-events: auto;
  /* level 3 drop shadow  https://design-dss.dataiku.com/383a36eba/p/94353b-surface/b/803d30 */
  box-shadow:
    0px 1px 25px 0px rgba(0, 0, 0, 0.06),
    0px 2px 12px -3px rgba(0, 0, 0, 0.18);
}

.fullscreen {
  min-width: 100%;
  max-width: 100%;
  height: 100%;
  border-radius: 0;
  margin: 0;
  right: 0;
}

.header-tab {
  background: var(--grp-c);
}
.header-tab[data-grp='agent'] {
  --grp-c: var(--color-agent);
}
.header-tab[data-grp='guard'] {
  --grp-c: var(--color-guard);
}
.header-tab[data-grp='tool'] {
  --grp-c: var(--color-tool);
}
.header-tab[data-grp='llm'] {
  --grp-c: var(--color-llm);
}
.header-tab[data-grp='event'] {
  --grp-c: var(--color-event);
}
.header-tab[data-grp='span'] {
  --grp-c: var(--color-span);
}

.header-tab[data-state='active'] {
  color: white;
}

.header-tab[data-state='inactive'] {
  background-color: white;
  color: var(--muted-foreground);
}

.header-tab:focus-visible {
  outline: none;
  box-shadow: 0 0 0 2px var(--ring);
}
</style>
