"use client"

import { memo, useMemo, useState, useRef, useEffect } from "react"
import { Code, FileText, ChevronDown } from "lucide-react"
import { CodeHighlighter, CODE_STYLE_RELAXED } from "./code-highlighter"
import { CopyButton } from "./copy-button"
import { findFunctionInCode, type FunctionLocation } from "./python-parser"

interface FunctionToOptimizeSectionProps {
  functionName: string | null
  filePath: string | null
  originalCode: string | null
}

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

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,
      lineCount: trimmedCode.split("\n").length,
    })
  }

  return files
}

function findTargetFile(files: ParsedFile[], filePath: string | null): ParsedFile | null {
  if (!filePath || files.length === 0) return null

  return files.find(file =>
    filePath.endsWith(file.path) ||
    file.path.endsWith(filePath) ||
    file.path === filePath
  ) || null
}

async function searchAllFiles(
  files: ParsedFile[],
  functionName: string
): Promise<Array<{ file: ParsedFile; location: FunctionLocation } | null>> {
  const searchPromises = files.map(async file => {
    const location = await findFunctionInCode(file.code, functionName)
    return location ? { file, location } : null
  })

  return Promise.all(searchPromises)
}

export const FunctionToOptimizeSection = memo(function FunctionToOptimizeSection({
  functionName,
  filePath,
  originalCode,
}: FunctionToOptimizeSectionProps) {
  const [isExpanded, setIsExpanded] = useState(true)
  const [functionLocation, setFunctionLocation] = useState<FunctionLocation | null>(null)
  const [actualFile, setActualFile] = useState<ParsedFile | null>(null)
  const codeContainerRef = useRef<HTMLDivElement>(null)

  const allFiles = useMemo(() => {
    if (!originalCode) return []
    return parseMarkdownCodeBlocks(originalCode)
  }, [originalCode])

  useEffect(() => {
    if (!functionName || allFiles.length === 0) {
      setFunctionLocation(null)
      setActualFile(null)
      return
    }

    let cancelled = false

    async function searchForFunction(funcName: string) {
      // Try to find the function in the specified file first
      const targetFile = findTargetFile(allFiles, filePath)
      if (targetFile) {
        const location = await findFunctionInCode(targetFile.code, funcName)
        if (!cancelled && location) {
          setFunctionLocation(location)
          setActualFile(targetFile)
          return
        }
      }

      // Search all files in parallel
      const searchResults = await searchAllFiles(allFiles, funcName)
      if (cancelled) return

      const found = searchResults.find(result => result !== null)
      if (found) {
        setFunctionLocation(found.location)
        setActualFile(found.file)
      } else {
        // Use target file or first file as fallback
        const fallbackFile = targetFile || allFiles[0]
        setFunctionLocation(null)
        setActualFile(fallbackFile)
      }
    }

    searchForFunction(functionName)
    return () => { cancelled = true }
  }, [functionName, filePath, allFiles])

  const functionFile = actualFile ?? allFiles[0] ?? null

  const functionLines = useMemo(() => {
    if (!functionLocation) return null
    const lines: number[] = []
    for (let i = functionLocation.startLine; i <= functionLocation.endLine; i++) {
      lines.push(i)
    }
    return lines
  }, [functionLocation])

  useEffect(() => {
    if (!isExpanded || !functionLocation || !codeContainerRef.current) return

    const scrollToFunction = () => {
      if (!codeContainerRef.current) return
      const container = codeContainerRef.current

      const lineHeight = 22.4
      const paddingTop = 16
      const targetLine = functionLocation.startLine - 1
      const scrollPosition = paddingTop + (targetLine * lineHeight) - (container.clientHeight / 3)

      container.scrollTo({
        top: Math.max(0, scrollPosition),
        behavior: 'smooth'
      })
    }

    const timer = setTimeout(scrollToFunction, 300)
    return () => clearTimeout(timer)
  }, [isExpanded, functionLocation])

  if (!functionFile) {
    return null
  }

  const highlightLines = functionLines && functionLines.length > 0 ? functionLines : undefined

  return (
    <div className="bg-white dark:bg-zinc-900 rounded-sm border border-zinc-200 dark:border-zinc-800 overflow-hidden">
      <div className="w-full p-4 bg-zinc-950 flex items-center justify-between">
        <div
          role="button"
          tabIndex={0}
          onClick={() => setIsExpanded(!isExpanded)}
          onKeyDown={e => { if (e.key === "Enter" || e.key === " ") setIsExpanded(!isExpanded) }}
          className="flex items-center gap-3 cursor-pointer hover:opacity-80 flex-1"
        >
          <div className="p-2 bg-zinc-800 rounded-sm">
            <Code className="h-4 w-4 text-zinc-400" />
          </div>
          <div className="text-left">
            <h2 className="text-lg font-semibold text-zinc-50">
              Function to Optimize
            </h2>
            <div className="flex items-center gap-2 mt-1">
              {functionName && (
                <code className="text-sm font-mono text-zinc-300 bg-zinc-800 px-2 py-0.5 rounded-sm">
                  {functionName}
                </code>
              )}
              {functionLocation && (
                <span className="text-xs text-zinc-500">
                  lines {functionLocation.startLine}-{functionLocation.endLine}
                </span>
              )}
            </div>
          </div>
          <ChevronDown className={`h-4 w-4 text-zinc-500 transition-transform duration-200 ${isExpanded ? '' : '-rotate-90'}`} />
        </div>
        <div className="flex items-center gap-2">
          <CopyButton text={functionFile.code} label="function code" />
        </div>
      </div>

      {isExpanded && (
        <>
          <div className="px-4 py-2 bg-zinc-900 border-t border-b border-zinc-800 flex items-center gap-2">
            <FileText className="h-4 w-4 text-zinc-500" />
            <span className="text-sm font-mono font-medium text-zinc-300">
              {functionFile.filename}
            </span>
            {functionFile.path !== functionFile.filename && (
              <span className="text-xs font-mono text-zinc-500">
                ({functionFile.path})
              </span>
            )}
            <span className="text-xs text-zinc-500 ml-auto">
              {functionFile.lineCount} lines
            </span>
          </div>

          <div ref={codeContainerRef} className="max-h-[500px] overflow-y-auto">
            <CodeHighlighter
              language={functionFile.language}
              code={functionFile.code}
              customStyle={CODE_STYLE_RELAXED}
              highlightLines={highlightLines}
            />
          </div>
        </>
      )}
    </div>
  )
})