import { type JSX } from "solid-js"
import {
  For,
  Show,
  Match,
  Switch,
  onMount,
  onCleanup,
  splitProps,
  createMemo,
  createEffect,
  createSignal,
} from "solid-js"
import { DateTime } from "luxon"
import { IconOpenAI, IconGemini, IconAnthropic } from "./icons/custom"
import {
  IconCpuChip,
  IconSparkles,
  IconUserCircle,
  IconChevronDown,
  IconCommandLine,
  IconChevronRight,
  IconPencilSquare,
  IconWrenchScrewdriver,
} from "./icons"
import DiffView from "./DiffView"
import CodeBlock from "./CodeBlock"
import styles from "./share.module.css"
import { type UIMessage } from "ai"
import { createStore, reconcile } from "solid-js/store"

const MIN_DURATION = 2

type Status =
  | "disconnected"
  | "connecting"
  | "connected"
  | "error"
  | "reconnecting"

type SessionMessage = UIMessage<{
  time: {
    created: number
    completed?: number
  }
  assistant?: {
    modelID: string
    providerID: string
    cost: number
    tokens: {
      input: number
      output: number
      reasoning: number
    }
  }
  sessionID: string
  tool: Record<
    string,
    {
      [key: string]: any
      time: {
        start: number
        end: number
      }
    }
  >
}>

type SessionInfo = {
  title: string
  cost?: number
}

function getFileType(path: string) {
  return path.split(".").pop()
}

function formatDuration(ms: number): string {
  const ONE_SECOND = 1000
  const ONE_MINUTE = 60 * ONE_SECOND

  if (ms >= ONE_MINUTE) {
    const minutes = Math.floor(ms / ONE_MINUTE)
    return minutes === 1 ? `1min` : `${minutes}mins`
  }

  if (ms >= ONE_SECOND) {
    const seconds = Math.floor(ms / ONE_SECOND)
    return `${seconds}s`
  }

  return `${ms}ms`
}

// Converts `{a:{b:{c:1}}` to `[['a.b.c', 1]]`
function flattenToolArgs(obj: any, prefix: string = ""): Array<[string, any]> {
  const entries: Array<[string, any]> = []

  for (const [key, value] of Object.entries(obj)) {
    const path = prefix ? `${prefix}.${key}` : key

    if (value !== null && typeof value === "object" && !Array.isArray(value)) {
      entries.push(...flattenToolArgs(value, path))
    } else {
      entries.push([path, value])
    }
  }

  return entries
}

function getStatusText(status: [Status, string?]): string {
  switch (status[0]) {
    case "connected":
      return "Connected"
    case "connecting":
      return "Connecting..."
    case "disconnected":
      return "Disconnected"
    case "reconnecting":
      return "Reconnecting..."
    case "error":
      return status[1] || "Error"
    default:
      return "Unknown"
  }
}

function ProviderIcon(props: { provider: string; size?: number }) {
  const size = props.size || 16
  return (
    <Switch fallback={<IconSparkles width={size} height={size} />}>
      <Match when={props.provider === "openai"}>
        <IconOpenAI width={size} height={size} />
      </Match>
      <Match when={props.provider === "anthropic"}>
        <IconAnthropic width={size} height={size} />
      </Match>
      <Match when={props.provider === "gemini"}>
        <IconGemini width={size} height={size} />
      </Match>
    </Switch>
  )
}

interface ResultsButtonProps extends JSX.HTMLAttributes<HTMLButtonElement> {
  results: boolean
}
function ResultsButton(props: ResultsButtonProps) {
  const [local, rest] = splitProps(props, ["results"])
  return (
    <button
      type="button"
      data-element-button-text
      data-element-button-more
      {...rest}
    >
      <span>{local.results ? "Hide results" : "Show results"}</span>
      <span data-button-icon>
        <Show
          when={local.results}
          fallback={<IconChevronRight width={10} height={10} />}
        >
          <IconChevronDown width={10} height={10} />
        </Show>
      </span>
    </button>
  )
}

interface TextPartProps extends JSX.HTMLAttributes<HTMLDivElement> {
  text: string
  expand?: boolean
  highlight?: boolean
}
function TextPart(props: TextPartProps) {
  const [local, rest] = splitProps(props, ["text", "expand", "highlight"])
  const [expanded, setExpanded] = createSignal(false)
  const [overflowed, setOverflowed] = createSignal(false)
  let preEl: HTMLPreElement | undefined

  function checkOverflow() {
    if (preEl && !local.expand) {
      setOverflowed(preEl.scrollHeight > preEl.clientHeight + 1)
    }
  }

  onMount(() => {
    checkOverflow()
    window.addEventListener("resize", checkOverflow)
  })

  createEffect(() => {
    local.text
    setTimeout(checkOverflow, 0)
  })

  onCleanup(() => {
    window.removeEventListener("resize", checkOverflow)
  })

  return (
    <div
      class={styles["message-text"]}
      data-highlight={local.highlight}
      data-expanded={expanded() || local.expand === true}
      {...rest}
    >
      <pre ref={(el) => (preEl = el)}>{local.text}</pre>
      {((!local.expand && overflowed()) || expanded()) && (
        <button
          type="button"
          data-element-button-text
          onClick={() => setExpanded((e) => !e)}
        >
          {expanded() ? "Show less" : "Show more"}
        </button>
      )}
    </div>
  )
}

interface TerminalPartProps extends JSX.HTMLAttributes<HTMLDivElement> {
  text: string
  desc?: string
  expand?: boolean
}
function TerminalPart(props: TerminalPartProps) {
  const [local, rest] = splitProps(props, ["text", "desc", "expand"])
  const [expanded, setExpanded] = createSignal(false)
  const [overflowed, setOverflowed] = createSignal(false)
  let preEl: HTMLElement | undefined

  function checkOverflow() {
    if (!preEl) return

    const code = preEl.getElementsByTagName("code")[0]

    if (code && !local.expand) {
      console.log(preEl.clientHeight, code.offsetHeight)
      setOverflowed(preEl.clientHeight < code.offsetHeight)
    }
  }

  onMount(() => {
    window.addEventListener("resize", checkOverflow)
  })

  onCleanup(() => {
    window.removeEventListener("resize", checkOverflow)
  })

  return (
    <div
      class={styles["message-terminal"]}
      data-expanded={expanded() || local.expand === true}
      {...rest}
    >
      <div data-section="body">
        <div data-section="header"><span>{local.desc}</span></div>
        <div data-section="content">
          <CodeBlock
            lang="ansi"
            onRendered={checkOverflow}
            ref={(el) => (preEl = el)}
            code={`\x1b[90m>\x1b[0m ${local.text}`}
          />
        </div>
      </div>
      {((!local.expand && overflowed()) || expanded()) && (
        <button
          type="button"
          data-element-button-text
          onClick={() => setExpanded((e) => !e)}
        >
          {expanded() ? "Show less" : "Show more"}
        </button>
      )}
    </div>
  )
}

function ToolFooter(props: { time: number }) {
  return (
    props.time > MIN_DURATION
      ? <span data-part-footer title={`${props.time}ms`}>
        {formatDuration(props.time)}
      </span>
      : <div data-part-footer="spacer"></div>
  )
}

export default function Share(props: { api: string }) {
  let params = new URLSearchParams(document.location.search)
  const id = params.get("id")

  const [store, setStore] = createStore<{
    info?: SessionInfo
    messages: Record<string, SessionMessage>
  }>({
    messages: {},
  })
  const messages = createMemo(() =>
    Object.values(store.messages).toSorted((a, b) => a.id?.localeCompare(b.id)),
  )
  const [connectionStatus, setConnectionStatus] = createSignal<
    [Status, string?]
  >(["disconnected", "Disconnected"])

  onMount(() => {
    const apiUrl = props.api

    if (!id) {
      setConnectionStatus(["error", "id not found"])
      return
    }

    if (!apiUrl) {
      console.error("API URL not found in environment variables")
      setConnectionStatus(["error", "API URL not found"])
      return
    }

    let reconnectTimer: number | undefined
    let socket: WebSocket | null = null

    // Function to create and set up WebSocket with auto-reconnect
    const setupWebSocket = () => {
      // Close any existing connection
      if (socket) {
        socket.close()
      }

      setConnectionStatus(["connecting"])

      // Always use secure WebSocket protocol (wss)
      const wsBaseUrl = apiUrl.replace(/^https?:\/\//, "wss://")
      const wsUrl = `${wsBaseUrl}/share_poll?id=${id}`
      console.log("Connecting to WebSocket URL:", wsUrl)

      // Create WebSocket connection
      socket = new WebSocket(wsUrl)

      // Handle connection opening
      socket.onopen = () => {
        setConnectionStatus(["connected"])
        console.log("WebSocket connection established")
      }

      // Handle incoming messages
      socket.onmessage = (event) => {
        console.log("WebSocket message received")
        try {
          const data = JSON.parse(event.data)
          const [root, type, ...splits] = data.key.split("/")
          if (root !== "session") return
          if (type === "info") {
            setStore("info", reconcile(data.content))
            return
          }
          if (type === "message") {
            const [, messageID] = splits
            setStore("messages", messageID, reconcile(data.content))
          }
        } catch (error) {
          console.error("Error parsing WebSocket message:", error)
        }
      }

      // Handle errors
      socket.onerror = (error) => {
        console.error("WebSocket error:", error)
        setConnectionStatus(["error", "Connection failed"])
      }

      // Handle connection close and reconnection
      socket.onclose = (event) => {
        console.log(`WebSocket closed: ${event.code} ${event.reason}`)
        setConnectionStatus(["reconnecting"])

        // Try to reconnect after 2 seconds
        clearTimeout(reconnectTimer)
        reconnectTimer = window.setTimeout(
          setupWebSocket,
          2000,
        ) as unknown as number
      }
    }

    // Initial connection
    setupWebSocket()

    // Clean up on component unmount
    onCleanup(() => {
      console.log("Cleaning up WebSocket connection")
      if (socket) {
        socket.close()
      }
      clearTimeout(reconnectTimer)
    })
  })

  const data = createMemo(() => {
    const result = {
      created: undefined as number | undefined,
      system: [] as string[],
      messages: [] as SessionMessage[],
      models: [] as string[][],
      cost: 0,
      tokens: {
        input: 0,
        output: 0,
        reasoning: 0,
      },
    }
    for (let i = 0; i < messages().length; i++) {
      const msg = messages()[i]

      const system = i === 0 && msg.role === "system"
      const assistant = msg.metadata?.assistant

      if (system) {
        for (const part of msg.parts) {
          if (part.type === "text") {
            result.system.push(part.text)
          }
        }
        result.created = msg.metadata?.time.created
        continue
      }

      result.messages.push(msg)

      if (assistant) {
        result.cost += assistant.cost
        result.tokens.input += assistant.tokens.input
        result.tokens.output += assistant.tokens.output
        result.tokens.reasoning += assistant.tokens.reasoning

        result.models.push([
          assistant.providerID,
          assistant.modelID,
        ])
      }
    }
    return result
  })
  const [showingSystemPrompt, showSystemPrompt] = createSignal(false)

  return (
    <main class={`${styles.root} not-content`}>
      <div class={styles.header}>
        <div data-section="title">
          <h1>{store.info?.title}</h1>
          <div>
            <div data-section="date">
              {data().created ? (
                <span
                  title={DateTime.fromMillis(
                    data().created || 0,
                  ).toLocaleString(DateTime.DATETIME_FULL_WITH_SECONDS)}
                >
                  {DateTime.fromMillis(
                    data().created || 0,
                  ).toLocaleString(DateTime.DATE_MED)}
                </span>
              ) : (
                <span data-element-label data-placeholder>
                  Started at &mdash;
                </span>
              )}
            </div>
            <p data-section="status">
              <span data-status={connectionStatus()[0]}>&#9679;</span>
              <span data-element-label>{getStatusText(connectionStatus())}</span>
            </p>
          </div>
        </div>
        <div data-section="row">
          <ul data-section="stats">
            <li>
              <span data-element-label>Cost</span>
              {data().cost !== undefined ? (
                <span>${data().cost.toFixed(2)}</span>
              ) : (
                <span data-placeholder>&mdash;</span>
              )}
            </li>
            <li>
              <span data-element-label>Input Tokens</span>
              {data().tokens.input ? (
                <span>{data().tokens.input}</span>
              ) : (
                <span data-placeholder>&mdash;</span>
              )}
            </li>
            <li>
              <span data-element-label>Output Tokens</span>
              {data().tokens.output ? (
                <span>{data().tokens.output}</span>
              ) : (
                <span data-placeholder>&mdash;</span>
              )}
            </li>
            <li>
              <span data-element-label>Reasoning Tokens</span>
              {data().tokens.reasoning ? (
                <span>{data().tokens.reasoning}</span>
              ) : (
                <span data-placeholder>&mdash;</span>
              )}
            </li>
          </ul>
          <ul data-section="stats" data-section-models>
            {data().models.length > 0 ? (
              <For each={Array.from(data().models)}>
                {([provider, model]) => (
                  <li>
                    <div data-stat-model-icon title={provider}>
                      <ProviderIcon provider={provider} />
                    </div>
                    <span data-stat-model>{model}</span>
                  </li>
                )}
              </For>
            ) : (
              <li>
                <span data-element-label>Models</span>
                <span data-placeholder>&mdash;</span>
              </li>
            )}
          </ul>
          <div data-section="system-prompt">
            <div data-section="icon">
              <IconCpuChip width={16} height={16} />
            </div>
            <div data-section="content">
              <button
                type="button"
                data-element-button-text
                data-element-button-more
                onClick={() => showSystemPrompt((e) => !e)}
              >
                <span>
                  {
                    showingSystemPrompt()
                      ? "Hide system prompt"
                      : "Show system prompt"
                  }
                </span>
                <span data-button-icon>
                  <Show
                    when={showingSystemPrompt()}
                    fallback={<IconChevronRight width={12} height={12} />}
                  >
                    <IconChevronDown width={12} height={12} />
                  </Show>
                </span>
              </button>
              <Show when={showingSystemPrompt()}>
                <TextPart
                  expand
                  data-size="sm"
                  text={data().system.join("\n").trim()}
                />
              </Show>
            </div>
          </div>
        </div>
      </div>

      <div>
        <Show
          when={data().messages.length > 0}
          fallback={<p>Waiting for messages...</p>}
        >
          <div class={styles.parts}>
            <For each={data().messages}>
              {(msg, msgIndex) => (
                <For each={msg.parts}>
                  {(part, partIndex) => {
                    if (
                      part.type === "step-start" &&
                      (partIndex() > 0 || !msg.metadata?.assistant)
                    )
                      return null

                    const [results, showResults] = createSignal(false)
                    const isLastPart = createMemo(
                      () =>
                        data().messages.length === msgIndex() + 1 &&
                        msg.parts.length === partIndex() + 1,
                    )
                    return (
                      <Switch>
                        {/* User text */}
                        <Match
                          when={
                            msg.role === "user" && part.type === "text" && part
                          }
                        >
                          {(part) => (
                            <div data-section="part" data-part-type="user-text">
                              <div data-section="decoration">
                                <div title="Message">
                                  <IconUserCircle width={18} height={18} />
                                </div>
                                <div></div>
                              </div>
                              <div data-section="content">
                                <TextPart
                                  highlight
                                  text={part().text}
                                  expand={isLastPart()}
                                />
                              </div>
                            </div>
                          )}
                        </Match>
                        {/* AI text */}
                        <Match
                          when={
                            msg.role === "assistant" &&
                            part.type === "text" &&
                            part
                          }
                        >
                          {(part) => (
                            <div data-section="part" data-part-type="ai-text">
                              <div data-section="decoration">
                                <div title="AI response">
                                  <IconSparkles width={18} height={18} />
                                </div>
                                <div></div>
                              </div>
                              <div data-section="content">
                                <TextPart
                                  text={part().text}
                                  expand={isLastPart()}
                                />
                              </div>
                            </div>
                          )}
                        </Match>
                        {/* AI model */}
                        <Match
                          when={
                            msg.role === "assistant" &&
                            part.type === "step-start" &&
                            msg.metadata?.assistant
                          }
                        >
                          {(assistant) => (
                            <div data-section="part" data-part-type="ai-model">
                              <div data-section="decoration">
                                <div>
                                  <ProviderIcon
                                    size={18}
                                    provider={assistant().providerID}
                                  />
                                </div>
                                <div></div>
                              </div>
                              <div data-section="content">
                                <div data-part-tool-body>
                                  <span
                                    data-size="md"
                                    data-part-title
                                    data-element-label
                                  >
                                    {assistant().providerID}
                                  </span>
                                  <span data-part-model>
                                    {assistant().modelID}
                                  </span>
                                </div>
                              </div>
                            </div>
                          )}
                        </Match>
                        {/* System text */}
                        <Match
                          when={
                            msg.role === "system" &&
                            part.type === "text" &&
                            part
                          }
                        >
                          {(part) => (
                            <div
                              data-section="part"
                              data-part-type="system-text"
                            >
                              <div data-section="decoration">
                                <div title="System message">
                                  <IconCpuChip width={18} height={18} />
                                </div>
                                <div></div>
                              </div>
                              <div data-section="content">
                                <div data-part-tool-body>
                                  <span data-element-label data-part-title>
                                    System
                                  </span>
                                  <TextPart
                                    data-size="sm"
                                    text={part().text}
                                    data-color="dimmed"
                                  />
                                </div>
                              </div>
                            </div>
                          )}
                        </Match>
                        {/* Edit tool */}
                        <Match
                          when={
                            msg.role === "assistant" &&
                            part.type === "tool-invocation" &&
                            part.toolInvocation.toolName === "opencode_edit" &&
                            part
                          }
                        >
                          {(part) => {
                            const metadata = createMemo(() => msg.metadata?.tool[part().toolInvocation.toolCallId])
                            const args = part().toolInvocation.args
                            const filePath = args.filePath

                            const duration = createMemo(() =>
                              DateTime.fromMillis(metadata()?.time.end || 0).diff(
                                DateTime.fromMillis(metadata()?.time.start || 0),
                              ).toMillis(),
                            )

                            return (
                              <div
                                data-section="part"
                                data-part-type="tool-edit"
                              >
                                <div data-section="decoration">
                                  <div title="Edit file">
                                    <IconPencilSquare width={18} height={18} />
                                  </div>
                                  <div></div>
                                </div>
                                <div data-section="content">
                                  <div data-part-tool-body>
                                    <span data-part-title data-size="md">
                                      <span data-element-label>Edit</span>
                                      <b>{filePath}</b>
                                    </span>
                                    <div data-part-tool-edit>
                                      <DiffView
                                        class={styles["diff-code-block"]}
                                        changes={metadata()?.changes || []}
                                        lang={getFileType(filePath)}
                                      />
                                    </div>
                                  </div>
                                  <ToolFooter time={duration()} />
                                </div>
                              </div>
                            )
                          }}
                        </Match>
                        {/* Bash tool */}
                        <Match
                          when={
                            msg.role === "assistant" &&
                            part.type === "tool-invocation" &&
                            part.toolInvocation.toolName === "opencode_bash" &&
                            part
                          }
                        >
                          {(part) => {
                            const metadata = createMemo(() => msg.metadata?.tool[part().toolInvocation.toolCallId])

                            const command = part().toolInvocation.args.command
                            const desc = part().toolInvocation.args.description
                            const stdout = metadata()?.stdout
                            const result = stdout || (part().toolInvocation.state === "result" && part().toolInvocation.result)

                            const duration = createMemo(() =>
                              DateTime.fromMillis(metadata()?.time.end || 0).diff(
                                DateTime.fromMillis(metadata()?.time.start || 0),
                              ).toMillis(),
                            )

                            return (
                              <div
                                data-section="part"
                                data-part-type="tool-bash"
                              >
                                <div data-section="decoration">
                                  <div title="Bash command">
                                    <IconCommandLine width={18} height={18} />
                                  </div>
                                  <div></div>
                                </div>
                                <div data-section="content">
                                  <div data-part-tool-body>
                                    <TerminalPart
                                      desc={desc}
                                      data-size="sm"
                                      text={command + (result ? `\n${result}` : "")}
                                    />
                                  </div>
                                  <ToolFooter time={duration()} />
                                </div>
                              </div>
                            )
                          }}
                        </Match>
                        {/* Tool call */}
                        <Match
                          when={
                            msg.role === "assistant" &&
                            part.type === "tool-invocation" &&
                            part
                          }
                        >
                          {(part) => {
                            const metadata = createMemo(() => msg.metadata?.tool[part().toolInvocation.toolCallId])

                            const duration = createMemo(() =>
                              DateTime.fromMillis(metadata()?.time.end || 0).diff(
                                DateTime.fromMillis(metadata()?.time.start || 0),
                              ).toMillis(),
                            )

                            return (
                              <div
                                data-section="part"
                                data-part-type="tool-fallback"
                              >
                                <div data-section="decoration">
                                  <div title="Tool call">
                                    <IconWrenchScrewdriver
                                      width={18}
                                      height={18}
                                    />
                                  </div>
                                  <div></div>
                                </div>
                                <div data-section="content">
                                  <div data-part-tool-body>
                                    <span data-part-title data-size="md">
                                      {part().toolInvocation.toolName}
                                    </span>
                                    <div data-part-tool-args>
                                      <For
                                        each={flattenToolArgs(
                                          part().toolInvocation.args,
                                        )}
                                      >
                                        {([name, value]) => (
                                          <>
                                            <div></div>
                                            <div>{name}</div>
                                            <div>{value}</div>
                                          </>
                                        )}
                                      </For>
                                    </div>
                                    <Switch>
                                      <Match
                                        when={
                                          part().toolInvocation.state ===
                                          "result" &&
                                          part().toolInvocation.result
                                        }
                                      >
                                        <div data-part-tool-result>
                                          <ResultsButton
                                            results={results()}
                                            onClick={() => showResults((e) => !e)}
                                          />
                                          <Show when={results()}>
                                            <TextPart
                                              expand
                                              data-size="sm"
                                              data-color="dimmed"
                                              text={part().toolInvocation.result}
                                            />
                                          </Show>
                                        </div>
                                      </Match>
                                      <Match
                                        when={
                                          part().toolInvocation.state === "call"
                                        }
                                      >
                                        <TextPart
                                          data-size="sm"
                                          data-color="dimmed"
                                          text="Calling..."
                                        />
                                      </Match>
                                    </Switch>
                                  </div>
                                  <ToolFooter time={duration()} />
                                </div>
                              </div>
                            )
                          }}
                        </Match>
                        {/* Fallback */}
                        <Match when={true}>
                          <div data-section="part" data-part-type="fallback">
                            <div data-section="decoration">
                              <div>
                                <Switch
                                  fallback={
                                    <IconWrenchScrewdriver
                                      width={16}
                                      height={16}
                                    />
                                  }
                                >
                                  <Match
                                    when={
                                      msg.role === "assistant" &&
                                      part.type !== "tool-invocation"
                                    }
                                  >
                                    <IconSparkles width={18} height={18} />
                                  </Match>
                                  <Match when={msg.role === "system"}>
                                    <IconCpuChip width={18} height={18} />
                                  </Match>
                                  <Match when={msg.role === "user"}>
                                    <IconUserCircle width={18} height={18} />
                                  </Match>
                                </Switch>
                              </div>
                              <div></div>
                            </div>
                            <div data-section="content">
                              <div data-part-tool-body>
                                <span data-element-label data-part-title>
                                  {part.type}
                                </span>
                                <TextPart
                                  text={JSON.stringify(part, null, 2)}
                                />
                              </div>
                            </div>
                          </div>
                        </Match>
                      </Switch>
                    )
                  }}
                </For>
              )}
            </For>
          </div>
        </Show>
      </div>

      <div style={{ margin: "2rem 0" }}>
        <div
          style={{
            border: "1px solid #ccc",
            padding: "1rem",
            "overflow-y": "auto",
          }}
        >
          <Show
            when={data().messages.length > 0}
            fallback={<p>Waiting for messages...</p>}
          >
            <ul style={{ "list-style-type": "none", padding: 0 }}>
              <For each={data().messages}>
                {(msg) => (
                  <li
                    style={{
                      padding: "0.75rem",
                      margin: "0.75rem 0",
                      "box-shadow": "0 1px 3px rgba(0,0,0,0.1)",
                    }}
                  >
                    <div>
                      <strong>Key:</strong> {msg.id}
                    </div>
                    <pre>{JSON.stringify(msg, null, 2)}</pre>
                  </li>
                )}
              </For>
            </ul>
          </Show>
        </div>
      </div>
    </main>
  )
}