"use client"

import { useState, useMemo, useCallback, memo } from "react"
import { ChevronDown, Code, FileText, Hash, FolderTree } from "lucide-react"
import { CodeHighlighter, CODE_STYLE } from "./code-highlighter"
import { CopyButton } from "./copy-button"
import { InfoIcon } from "./info-icon"

interface CodeContextSectionProps {
  functionName: string | null
  filePath: string | null
  originalCode: string | null
  dependencyCode: string | null
}

interface ParsedFile {
  path: string
  filename: string
  language: string
  code: string
  tokens: number
  lineCount: number
}

function estimateTokens(text: string): number {
  return Math.ceil(text.length / 4)
}

function getFilename(path: string): string {
  return path.split("/").pop() || path
}

function parseMarkdownCodeBlocks(markdown: string): ParsedFile[] {
  const files: ParsedFile[] = []
  const regex = /```(\w+):([^\n]+)\n([\s\S]*?)```/g
  let match

  while ((match = regex.exec(markdown)) !== null) {
    const [, language, path, code] = match
    const trimmedCode = code.trimEnd()
    files.push({
      path,
      filename: getFilename(path),
      language: language || "python",
      code: trimmedCode,
      tokens: estimateTokens(trimmedCode),
      lineCount: trimmedCode.split("\n").length,
    })
  }

  return files
}

export const CodeContextSection = memo(function CodeContextSection({
  functionName,
  filePath,
  originalCode,
  dependencyCode,
}: CodeContextSectionProps) {
  const [isExpanded, setIsExpanded] = useState(false)
  const [expandedSections, setExpandedSections] = useState<Set<string>>(new Set())
  const [expandedFiles, setExpandedFiles] = useState<Set<string>>(new Set())

  const { rwFiles, roFiles, metrics } = useMemo(() => {
    const rwFiles = originalCode ? parseMarkdownCodeBlocks(originalCode) : []
    const roFiles = dependencyCode ? parseMarkdownCodeBlocks(dependencyCode) : []
    const rwTokens = rwFiles.reduce((sum, f) => sum + f.tokens, 0)
    const roTokens = roFiles.reduce((sum, f) => sum + f.tokens, 0)
    const totalFiles = rwFiles.length + roFiles.length
    const totalTokens = rwTokens + roTokens
    const rwChars = rwFiles.reduce((sum, f) => sum + f.code.length, 0)
    const roChars = roFiles.reduce((sum, f) => sum + f.code.length, 0)

    return {
      rwFiles,
      roFiles,
      metrics: { rwTokens, roTokens, totalFiles, totalTokens, rwChars, roChars },
    }
  }, [originalCode, dependencyCode])

  const toggleSection = useCallback((section: string): void => {
    setExpandedSections(prev => {
      const next = new Set(prev)
      if (next.has(section)) {
        next.delete(section)
      } else {
        next.add(section)
      }
      return next
    })
  }, [])

  const toggleFile = useCallback((fileKey: string): void => {
    setExpandedFiles(prev => {
      const next = new Set(prev)
      if (next.has(fileKey)) {
        next.delete(fileKey)
      } else {
        next.add(fileKey)
      }
      return next
    })
  }, [])

  if (!originalCode && !dependencyCode) {
    return null
  }

  return (
    <div className="bg-white dark:bg-zinc-900 rounded-sm border border-zinc-200 dark:border-zinc-800">
      <div
        role="button"
        tabIndex={0}
        onClick={() => setIsExpanded(!isExpanded)}
        onKeyDown={e => { if (e.key === "Enter" || e.key === " ") setIsExpanded(!isExpanded) }}
        className="w-full p-6 flex items-center justify-between hover:bg-zinc-50 dark:hover:bg-zinc-800 transition-colors rounded-sm cursor-pointer"
      >
        <div className="flex items-center gap-2">
          <Code className="h-4 w-4 text-zinc-400" />
          <h2 className="text-xl font-bold text-zinc-900 dark:text-white">Code Context</h2>
          <InfoIcon
            content="The code provided to the LLM for optimization. Read-writable code can be modified, read-only dependencies are for reference only."
            side="right"
          />
        </div>
        <div className="flex items-center gap-4">
          <div className="flex items-center gap-4 text-sm text-zinc-500 dark:text-zinc-400">
            <span className="flex items-center gap-1">
              <Hash className="h-4 w-4" />
              {metrics.totalTokens.toLocaleString()} tokens
            </span>
            <span className="flex items-center gap-1">
              <FolderTree className="h-4 w-4" />
              {metrics.totalFiles} files
            </span>
          </div>
          <ChevronDown className={`h-4 w-4 text-zinc-400 transition-transform duration-200 ${isExpanded ? '' : '-rotate-90'}`} />
        </div>
      </div>

      {isExpanded && (
        <div className="px-6 pb-6 pt-0 border-t border-zinc-200 dark:border-zinc-800 space-y-4">
          {(functionName || filePath) && (
            <div className="grid grid-cols-1 sm:grid-cols-2 gap-4 py-4">
              {functionName && (
                <div className="flex flex-col">
                  <span className="text-xs text-zinc-500 dark:text-zinc-400 mb-1">Function</span>
                  <code className="text-sm font-mono text-zinc-900 dark:text-white bg-zinc-100 dark:bg-zinc-800 px-2 py-1 rounded-sm w-fit">
                    {functionName}
                  </code>
                </div>
              )}
              {filePath && (
                <div className="flex flex-col">
                  <span className="text-xs text-zinc-500 dark:text-zinc-400 mb-1">File</span>
                  <code className="text-sm font-mono text-zinc-900 dark:text-white bg-zinc-100 dark:bg-zinc-800 px-2 py-1 rounded-sm w-fit truncate max-w-full" title={filePath}>
                    {filePath}
                  </code>
                </div>
              )}
            </div>
          )}

          {rwFiles.length > 0 && roFiles.length > 0 && (
            <div className="p-4 bg-zinc-50 dark:bg-zinc-950 rounded-sm border border-zinc-200 dark:border-zinc-800">
              <div className="flex items-center gap-2 mb-3">
                <Hash className="h-4 w-4 text-zinc-400" />
                <span className="text-sm font-medium text-zinc-700 dark:text-zinc-300">
                  Token Distribution (estimated)
                </span>
              </div>
              <TokenDistributionBar
                rwTokens={metrics.rwTokens}
                roTokens={metrics.roTokens}
                totalTokens={metrics.totalTokens || 1}
              />
              <div className="flex gap-4 text-xs">
                <div className="flex items-center gap-1.5">
                  <div className="w-3 h-3 rounded-sm bg-emerald-500" />
                  <span className="text-zinc-600 dark:text-zinc-400">
                    Read-Writable: {metrics.rwTokens.toLocaleString()} ({rwFiles.length} files)
                  </span>
                </div>
                <div className="flex items-center gap-1.5">
                  <div className="w-3 h-3 rounded-sm bg-slate-500" />
                  <span className="text-zinc-600 dark:text-zinc-400">
                    Read-Only: {metrics.roTokens.toLocaleString()} ({roFiles.length} files)
                  </span>
                </div>
              </div>
            </div>
          )}

          {rwFiles.length > 0 && (
            <CodeGroupSection
              title="Read-Writable Code"
              subtitle="Function to optimize + first-degree helpers"
              accentColor="emerald"
              tokenCount={metrics.rwTokens}
              charCount={metrics.rwChars}
              files={rwFiles}
              isExpanded={expandedSections.has("rw")}
              onToggle={() => toggleSection("rw")}
              expandedFiles={expandedFiles}
              onToggleFile={toggleFile}
              sectionKey="rw"
            />
          )}

          {roFiles.length > 0 && (
            <CodeGroupSection
              title="Read-Only Dependencies"
              subtitle="Transitive dependencies for context"
              accentColor="slate"
              tokenCount={metrics.roTokens}
              charCount={metrics.roChars}
              files={roFiles}
              isExpanded={expandedSections.has("ro")}
              onToggle={() => toggleSection("ro")}
              expandedFiles={expandedFiles}
              onToggleFile={toggleFile}
              sectionKey="ro"
            />
          )}
        </div>
      )}
    </div>
  )
})

interface CodeGroupSectionProps {
  title: string
  subtitle: string
  accentColor: "emerald" | "slate"
  tokenCount: number
  charCount: number
  files: ParsedFile[]
  isExpanded: boolean
  onToggle: () => void
  expandedFiles: Set<string>
  onToggleFile: (fileKey: string) => void
  sectionKey: string
}

function getAccentColorClasses(accentColor: "emerald" | "slate"): { border: string; bg: string; icon: string } {
  switch (accentColor) {
    case "emerald":
      return {
        border: "border-zinc-200 dark:border-zinc-800",
        bg: "bg-zinc-50 dark:bg-zinc-900",
        icon: "text-emerald-500",
      }
    case "slate":
      return {
        border: "border-zinc-200 dark:border-zinc-800",
        bg: "bg-zinc-50 dark:bg-zinc-900",
        icon: "text-zinc-500",
      }
  }
}

const CodeGroupSection = memo(function CodeGroupSection({
  title,
  subtitle,
  accentColor,
  tokenCount,
  charCount,
  files,
  isExpanded,
  onToggle,
  expandedFiles,
  onToggleFile,
  sectionKey,
}: CodeGroupSectionProps) {
  const { border: borderColor, bg: bgColor, icon: iconColor } = getAccentColorClasses(accentColor)

  return (
    <div className={`rounded-sm border ${borderColor} overflow-hidden`}>
      <div
        role="button"
        tabIndex={0}
        onClick={onToggle}
        onKeyDown={e => { if (e.key === "Enter" || e.key === " ") onToggle() }}
        className={`w-full p-4 flex items-center justify-between hover:bg-zinc-100 dark:hover:bg-zinc-800 transition-colors cursor-pointer ${bgColor}`}
      >
        <div className="flex items-center gap-3">
          <FileText className={`h-4 w-4 ${iconColor}`} />
          <div className="text-left">
            <span className="text-sm font-medium text-zinc-900 dark:text-white">{title}</span>
            <span className="text-xs text-zinc-500 dark:text-zinc-400 ml-2">{subtitle}</span>
          </div>
        </div>
        <div className="flex items-center gap-3">
          <span className="text-xs text-zinc-500 dark:text-zinc-400">
            {tokenCount.toLocaleString()} tokens · {charCount.toLocaleString()} chars · {files.length} files
          </span>
          <ChevronDown className={`h-4 w-4 text-zinc-400 transition-transform duration-200 ${isExpanded ? '' : '-rotate-90'}`} />
        </div>
      </div>

      {isExpanded && (
        <div className="border-t border-zinc-200 dark:border-zinc-800 divide-y divide-zinc-200 dark:divide-zinc-800">
          {files.map((file, index) => {
            const fileKey = `${sectionKey}-${index}`
            const isFileExpanded = expandedFiles.has(fileKey)
            return (
              <div key={fileKey}>
                <div className="px-4 py-2 flex items-center justify-between bg-zinc-100 dark:bg-zinc-950">
                  <div
                    role="button"
                    tabIndex={0}
                    onClick={() => onToggleFile(fileKey)}
                    onKeyDown={e => { if (e.key === "Enter" || e.key === " ") onToggleFile(fileKey) }}
                    className="flex items-center gap-2 cursor-pointer hover:opacity-80 flex-1"
                  >
                    <FileText className="h-3.5 w-3.5 text-zinc-400" />
                    <span className="text-sm font-mono font-medium text-zinc-900 dark:text-white">
                      {file.filename}
                    </span>
                    <span className="text-xs text-zinc-500 dark:text-zinc-400" title={file.path}>
                      {file.path !== file.filename && `(${file.path})`}
                    </span>
                    <ChevronDown className={`h-3.5 w-3.5 text-zinc-400 transition-transform duration-200 ${isFileExpanded ? '' : '-rotate-90'}`} />
                  </div>
                  <div className="flex items-center gap-2">
                    <span className="text-xs text-zinc-500 dark:text-zinc-400">
                      {file.lineCount} lines
                    </span>
                    <CopyButton text={file.code} label={file.filename} size="sm" />
                  </div>
                </div>

                {isFileExpanded && (
                  <div className="max-h-96 overflow-y-auto">
                    <CodeHighlighter
                      language={file.language}
                      code={file.code}
                      customStyle={CODE_STYLE}
                    />
                  </div>
                )}
              </div>
            )
          })}
        </div>
      )}
    </div>
  )
})

interface TokenDistributionBarProps {
  rwTokens: number
  roTokens: number
  totalTokens: number
}

const TokenDistributionBar = memo(function TokenDistributionBar({
  rwTokens,
  roTokens,
  totalTokens,
}: TokenDistributionBarProps) {
  const rwPercent = Math.round((rwTokens / totalTokens) * 100)
  const roPercent = Math.round((roTokens / totalTokens) * 100)

  return (
    <div className="flex h-6 rounded-sm overflow-hidden border border-zinc-200 dark:border-zinc-800 mb-2">
      <div
        className="bg-emerald-500 flex items-center justify-center text-white text-xs font-semibold"
        style={{ width: `${rwPercent}%` }}
      >
        {rwPercent > 15 && `${rwPercent}%`}
      </div>
      <div
        className="bg-slate-500 flex items-center justify-center text-white text-xs font-semibold"
        style={{ width: `${roPercent}%` }}
      >
        {roPercent > 15 && `${roPercent}%`}
      </div>
    </div>
  )
})