import React, { useState, useEffect } from 'react'
import { MessageSquare, AlertCircle, Copy, ChevronDown, ChevronRight } from 'lucide-react'
import ReactMarkdown from 'react-markdown'
import rehypeRaw from 'rehype-raw'
import type { SourceMetadata, SubQuestionSources } from '../types'
import { getPdfViewerUrl } from '../lib/api'
import { processCitations, processCitationsForSubq, extractCitedSources, highlightTerms } from '../utils/citationParser'
import { bulletizeMarkdown } from '../utils/citationParser'

const V2_BASE = `${import.meta.env.VITE_API_BASE_URL ?? '/api/v1'}/v2`

function getHighlightUrl(document_id: string, chunk_index: number, sub_question: string): string {
  return `${V2_BASE}/highlights?document_id=${encodeURIComponent(document_id)}&chunk_index=${chunk_index}&sub_question=${encodeURIComponent(sub_question)}`
}

interface ResponsePanelProps {
  answer: string | null
  sources?: SourceMetadata[]
  subQuestionSources?: SubQuestionSources[] | null
  isLoading?: boolean
  phase?: string
  error?: string | null
  historyId?: number | null
}

const CitationLink = ({ href, children }: { href?: string; children?: React.ReactNode }) => (
  <a
    href={href}
    target="_blank"
    rel="noopener noreferrer"
    className="text-blue-600 hover:text-blue-800 hover:underline"
  >
    {children}
  </a>
)

const HighlightMark = ({ children }: { children?: React.ReactNode }) => (
  <mark className="bg-yellow-200 rounded px-0.5">{children}</mark>
)

function parseAnswerSections(answer: string): string[] {
  const sections = answer.split(/## Sub-question \d+:[^\n]*\n/)
  if (sections.length <= 1) return [bulletizeMarkdown(answer)]
  return sections.filter((s) => s.trim().length > 0).map((s) => bulletizeMarkdown(s.trim()))
}

function SubQuestionSourceCard({ 
  source, 
  index, 
  highlightReady = false,
  subQuestionText = '',
}: { 
  source: SourceMetadata
  index: number
  highlightReady?: boolean
  subQuestionText?: string 
}) {
  return (
    <div
      key={index}
      className="border border-gray-200 rounded p-3 bg-gray-50"
    >
      <div className="flex items-center justify-between">
        <div className="font-medium text-sm">{source.filename}</div>
        {source.page_number !== null && (
          <span className="text-xs text-blue-600 font-medium">
            Page {source.page_number}
          </span>
        )}
      </div>
      <div className="text-sm text-gray-500">{source.upload_date}</div>
      <div className="text-sm text-gray-600 mt-1">{source.content_summary}</div>
      <div className="flex items-center gap-2 mt-1">
        <div className="text-xs text-gray-400">Chunk {source.chunk_index}</div>
        {source.chunk_file_path && (
          <a
            href={
              highlightReady && source.document_id
                ? getHighlightUrl(source.document_id, source.chunk_index, subQuestionText)
                : getPdfViewerUrl(source.chunk_file_path, source.page_number ?? undefined, source.filename)
            }
            target="_blank"
            rel="noopener noreferrer"
            className="text-xs text-blue-600 hover:text-blue-800 hover:underline"
            data-testid="view-chunk-pdf-link"
          >
            View PDF
          </a>
        )}
      </div>
    </div>
  )
}

function SubQuestionSection({
  index,
  subQuestion,
  answerSection,
  allSubQuestionSources,
  highlightReadyKeys,
}: {
  index: number
  subQuestion: SubQuestionSources
  answerSection: string
  allSubQuestionSources: SubQuestionSources[]
  highlightReadyKeys: Set<string>
}) {
  const [expanded, setExpanded] = useState(false)
  // Enrich sources with sub_question_text so buildCitationUrl can route to highlight page.
  // Look up citations across ALL sub-questions' sources because the LLM
  // may cite chunks from other sub-questions' contexts.
  const allSources = allSubQuestionSources.flatMap(sq =>
    sq.sources.map(s => ({ ...s, sub_question_text: sq.sub_question_text }))
  )
  const processedAnswer = processCitations(answerSection, allSources, highlightReadyKeys)
  const highlightedAnswer = highlightTerms(processedAnswer)

  return (
    <div
      id={`subq-${index}`}
      data-testid={`subq-section-${index}`}
      className="border border-gray-200 rounded-lg p-4 space-y-3"
    >
      <div className="flex items-center space-x-2 text-gray-700 font-medium">
        <span className="text-gray-400">Sub-question {index + 1}:</span>
        <span>{subQuestion.sub_question_text}</span>
      </div>

      <div className="prose prose-sm max-w-none text-gray-800">
        <ReactMarkdown
          components={{ a: CitationLink, mark: HighlightMark }}
          rehypePlugins={[rehypeRaw]}
        >
          {highlightedAnswer}
        </ReactMarkdown>
      </div>

      {subQuestion.sources.length > 0 && (
        <div className="mt-2">
          <button
            data-testid="sources-toggle"
            onClick={() => setExpanded(!expanded)}
            className="flex items-center space-x-2 text-xs uppercase text-gray-500 tracking-wide mb-2 hover:text-gray-700 transition-colors duration-200"
          >
            {expanded ? (
              <ChevronDown className="w-4 h-4" />
            ) : (
              <ChevronRight className="w-4 h-4" />
            )}
            <span>Sources ({subQuestion.sources.length})</span>
          </button>
          {expanded && (
            <div data-testid="sources-container" className="grid grid-cols-2 gap-2">
              {subQuestion.sources.map((source, idx) => (
                <SubQuestionSourceCard 
                  key={idx} 
                  source={source} 
                  index={idx}
                  highlightReady={source.document_id ? highlightReadyKeys.has(`${source.document_id}_${source.chunk_index}_${encodeURIComponent(subQuestion.sub_question_text)}`) : false}
                  subQuestionText={subQuestion.sub_question_text}
                />
              ))}
            </div>
          )}
        </div>
      )}
    </div>
  )
}

function SubQuestionSections({
  answer,
  subQuestionSources,
  isLoading,
  historyId,
}: {
  answer: string | null
  subQuestionSources: SubQuestionSources[]
  isLoading?: boolean
  historyId?: number | null
}) {
  const [copied, setCopied] = useState(false)
  const [highlightReadyKeys, setHighlightReadyKeys] = useState<Set<string>>(new Set())
  const [highlightStatus, setHighlightStatus] = useState<'loading' | 'done' | ''>('')

  useEffect(() => {
    if (!answer || isLoading || !subQuestionSources.length) return

    const targets: Array<{
      document_id: string
      chunk_index: number
      sub_question_text: string
      sub_question_index: number
    }> = []
    
    const sections = parseAnswerSections(answer)
    
    subQuestionSources.forEach((sq) => {
      const answerSection = sections[sq.sub_question_index] ?? ''
      const citedSources = extractCitedSources(answerSection, sq.sources)
      citedSources.forEach((source) => {
        if (source.document_id) {
          targets.push({
            document_id: source.document_id,
            chunk_index: source.chunk_index,
            sub_question_text: sq.sub_question_text,
            sub_question_index: sq.sub_question_index,
          })
        }
      })
    })
    
    if (targets.length === 0) return
    
    setHighlightStatus('loading')
    
    console.log(
      `[Highlight] Sending batch request: ${targets.length} targets`,
      targets.map((t: { document_id: string; chunk_index: number }) => `${t.document_id}#${t.chunk_index}`),
    )
    const startTime = performance.now()
    
    const url = historyId
      ? `${V2_BASE}/highlights/batch?history_id=${historyId}`
      : `${V2_BASE}/highlights/batch`
    fetch(url, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ targets }),
    })
      .then((res) => res.json())
      .then((data: { status: string; cached_count?: number; errors?: string[]; highlight_time_ms?: number }) => {
        const elapsed = Math.round(performance.now() - startTime)
        console.log(
          `[Highlight] Response: status=${data.status} cached=${data.cached_count ?? '?'} ` +
          `errors=${data.errors?.length ?? 0} backend_time=${data.highlight_time_ms ?? '?'}ms total=${elapsed}ms`,
        )
        if (data.errors?.length) {
          console.warn('[Highlight] Errors:', data.errors)
        }
        if (data.status === 'failed') {
          console.error('[Highlight] Batch computation failed:', data.errors)
          setHighlightStatus('')
          return
        }
        if (data.status === 'completed' || data.status === 'partial') {
          const keys = new Set<string>()
          targets.forEach((t: { document_id: string; chunk_index: number; sub_question_text: string }) => {
            keys.add(`${t.document_id}_${t.chunk_index}_${encodeURIComponent(t.sub_question_text)}`)
          })
          setHighlightReadyKeys(keys)
          setHighlightStatus('done')
          setTimeout(() => setHighlightStatus(''), 4000)
        }
      })
      .catch((err) => {
        console.error('Highlight batch computation failed:', err)
        setHighlightStatus('')
      })
  }, [answer, isLoading, subQuestionSources, historyId])

  const sections = answer ? parseAnswerSections(answer) : []

  const handleCopyAnswer = async (): Promise<void> => {
    if (answer) {
      await navigator.clipboard.writeText(answer)
      setCopied(true)
      setTimeout(() => setCopied(false), 2000)
    }
  }

  if (isLoading) {
    return (
      <div className="space-y-3 p-4">
        <div className="flex items-center space-x-2 text-sm text-gray-500 mb-4">
          <div className="w-4 h-4 border-2 border-gray-300 border-t-blue-500 rounded-full animate-spin" />
          <span>Generating answers...</span>
        </div>
        <div
          data-testid="skeleton-line"
          className="h-4 bg-gray-200 rounded animate-pulse w-full"
        />
        <div
          data-testid="skeleton-line"
          className="h-4 bg-gray-200 rounded animate-pulse w-11/12"
        />
        <div
          data-testid="skeleton-line"
          className="h-4 bg-gray-200 rounded animate-pulse w-4/5"
        />
        <div
          data-testid="skeleton-line"
          className="h-4 bg-gray-200 rounded animate-pulse w-3/5"
        />
        <div
          data-testid="skeleton-line"
          className="h-4 bg-gray-200 rounded animate-pulse w-2/5"
        />
        <div className="mt-4 space-y-2">
          <div
            data-testid="source-skeleton"
            className="border border-gray-200 rounded p-3 bg-gray-50"
          >
            <div className="h-3 bg-gray-200 rounded w-3/4 mb-2 animate-pulse" />
            <div className="h-2 bg-gray-200 rounded w-1/2 animate-pulse" />
          </div>
          <div
            data-testid="source-skeleton"
            className="border border-gray-200 rounded p-3 bg-gray-50"
          >
            <div className="h-3 bg-gray-200 rounded w-2/3 mb-2 animate-pulse" />
            <div className="h-2 bg-gray-200 rounded w-1/3 animate-pulse" />
          </div>
        </div>
      </div>
    )
  }

  return (
    <div className="p-4 space-y-4">
      {highlightStatus !== '' && (
        <div
          data-testid="highlight-status-toast"
          className={`fixed top-4 left-4 z-50 px-4 py-2 rounded-lg shadow-lg text-sm font-medium transition-opacity duration-300 ${
            highlightStatus === 'done'
              ? 'bg-green-50 border border-green-200 text-green-700'
              : 'bg-amber-50 border border-amber-200 text-amber-700'
          }`}
        >
          {highlightStatus === 'loading' && (
            <span className="flex items-center gap-2">
              <span className="w-3 h-3 border-2 border-amber-400 border-t-transparent rounded-full animate-spin" />
              Preparing highlights...
            </span>
          )}
          {highlightStatus === 'done' && 'Highlights ready'}
        </div>
      )}
      <div className="flex items-start justify-between">
        <div className="flex-1" />
        <button
          data-testid="copy-answer-btn"
          onClick={handleCopyAnswer}
          className="flex items-center space-x-1 text-gray-500 hover:text-gray-700 transition-colors duration-200"
          title="Copy answer to clipboard"
        >
          <Copy className="w-4 h-4" />
          <span className="text-sm">{copied ? 'Copied!' : 'Copy'}</span>
        </button>
      </div>

      <div className="space-y-4">
        {subQuestionSources.map((subQuestion, index) => (
          <SubQuestionSection
            key={index}
            index={index}
            subQuestion={subQuestion}
            answerSection={sections[index] ?? ''}
            allSubQuestionSources={subQuestionSources}
            highlightReadyKeys={highlightReadyKeys}
          />
        ))}
      </div>
    </div>
  )
}

function FlatResponse({
  answer,
  sources,
  isLoading,
  phase,
  error,
}: {
  answer: string | null
  sources?: SourceMetadata[]
  isLoading?: boolean
  phase?: string
  error?: string | null
}) {
  const [sourcesExpanded, setSourcesExpanded] = useState(true)
  const [copied, setCopied] = useState(false)
  const safeSources = sources ?? []

  const processedAnswer = answer ? processCitations(answer, safeSources) : answer ?? ''
  const highlightedAnswer = processedAnswer ? highlightTerms(processedAnswer) : processedAnswer

  const handleCopyAnswer = async (): Promise<void> => {
    if (answer) {
      await navigator.clipboard.writeText(answer)
      setCopied(true)
      setTimeout(() => setCopied(false), 2000)
    }
  }

  if (answer === null && !isLoading && !error) {
    return (
      <div className="h-full flex items-center justify-center bg-white/50">
        <div className="text-center space-y-2">
          <MessageSquare className="mx-auto w-12 h-12 text-gray-600" />
          <div className="text-gray-700 font-semibold">Ask a question to see the answer here.</div>
        </div>
      </div>
    )
  }

  const phaseMessages: Record<string, string> = {
    retrieving: 'Searching documents...',
    filtering: 'Filtering relevant passages...',
    generating: 'Generating answer...',
  }

  if (isLoading) {
    return (
      <div className="space-y-3 p-4">
        <div className="flex items-center space-x-2 text-sm text-gray-500 mb-4">
          <div className="w-4 h-4 border-2 border-gray-300 border-t-blue-500 rounded-full animate-spin" />
          <span>{phaseMessages[phase ?? ''] || 'Processing...'}</span>
        </div>
        <div
          data-testid="skeleton-line"
          className="h-4 bg-gray-200 rounded animate-pulse w-full"
        />
        <div
          data-testid="skeleton-line"
          className="h-4 bg-gray-200 rounded animate-pulse w-11/12"
        />
        <div
          data-testid="skeleton-line"
          className="h-4 bg-gray-200 rounded animate-pulse w-4/5"
        />
        <div
          data-testid="skeleton-line"
          className="h-4 bg-gray-200 rounded animate-pulse w-3/5"
        />
        <div
          data-testid="skeleton-line"
          className="h-4 bg-gray-200 rounded animate-pulse w-2/5"
        />
        <div className="mt-4 space-y-2">
          <div
            data-testid="source-skeleton"
            className="border border-gray-200 rounded p-3 bg-gray-50"
          >
            <div className="h-3 bg-gray-200 rounded w-3/4 mb-2 animate-pulse" />
            <div className="h-2 bg-gray-200 rounded w-1/2 animate-pulse" />
          </div>
          <div
            data-testid="source-skeleton"
            className="border border-gray-200 rounded p-3 bg-gray-50"
          >
            <div className="h-3 bg-gray-200 rounded w-2/3 mb-2 animate-pulse" />
            <div className="h-2 bg-gray-200 rounded w-1/3 animate-pulse" />
          </div>
        </div>
      </div>
    )
  }

  if (error) {
    return (
      <div className="bg-red-50 border border-red-200 text-red-700 rounded p-4 m-4">
        <div className="flex items-center space-x-2">
          <AlertCircle className="w-5 h-5" />
          <span>{error}</span>
        </div>
      </div>
    )
  }

  return (
    <div className="p-4 space-y-4">
      <div className="space-y-2">
      <div className="flex items-start justify-between">
          <div className="flex-1" />
          <button
            data-testid="copy-answer-btn"
            onClick={handleCopyAnswer}
            className="flex items-center space-x-1 text-gray-500 hover:text-gray-700 transition-colors duration-200"
            title="Copy answer to clipboard"
          >
            <Copy className="w-4 h-4" />
            <span className="text-sm">{copied ? 'Copied!' : 'Copy'}</span>
          </button>
        </div>
        <div className="prose prose-sm max-w-none text-gray-800">
          <ReactMarkdown
            components={{ a: CitationLink, mark: HighlightMark }}
            rehypePlugins={[rehypeRaw]}
          >
            {highlightedAnswer}
          </ReactMarkdown>
        </div>
      </div>

      {safeSources.length > 0 && (
        <div className="mt-6">
          <button
            data-testid="sources-toggle"
            onClick={() => setSourcesExpanded(!sourcesExpanded)}
            className="flex items-center space-x-2 text-xs uppercase text-gray-500 tracking-wide mb-3 hover:text-gray-700 transition-colors duration-200"
          >
            {sourcesExpanded ? (
              <ChevronDown className="w-4 h-4" />
            ) : (
              <ChevronRight className="w-4 h-4" />
            )}
            <span>Sources ({safeSources.length})</span>
          </button>
          {sourcesExpanded && (
            <div data-testid="sources-container" className="grid grid-cols-2 gap-2">
              {safeSources.map((source, index) => (
                <div
                  key={index}
                  className="border border-gray-200 rounded p-3 bg-gray-50"
                >
                  <div className="flex items-center justify-between">
                    <div className="font-medium text-sm">{source.filename}</div>
                    {source.page_number !== null && (
                      <span className="text-xs text-blue-600 font-medium">
                        Page {source.page_number}
                      </span>
                    )}
                  </div>
                  <div className="text-sm text-gray-500">{source.upload_date}</div>
                  <div className="text-sm text-gray-600 mt-1">{source.content_summary}</div>
                  <div className="flex items-center justify-between mt-1">
                    <div className="text-xs text-gray-400">Chunk {source.chunk_index}</div>
                    {source.chunk_file_path && (
                      <a
                        href={getPdfViewerUrl(source.chunk_file_path, source.page_number ?? undefined, source.filename)}
                        target="_blank"
                        rel="noopener noreferrer"
                        className="text-xs text-blue-600 hover:text-blue-800 hover:underline"
                        data-testid="view-chunk-pdf-link"
                      >
                        View PDF
                      </a>
                    )}
                  </div>
                </div>
              ))}
            </div>
          )}
        </div>
      )}
    </div>
  )
}

export const ResponsePanel: React.FC<ResponsePanelProps> = ({
  answer,
  sources,
  subQuestionSources,
  isLoading,
  phase,
  error,
  historyId,
}) => {
  if (subQuestionSources && subQuestionSources.length > 0) {
    return (
      <SubQuestionSections
        answer={answer}
        subQuestionSources={subQuestionSources}
        isLoading={isLoading}
        historyId={historyId}
      />
    )
  }

  return (
    <FlatResponse
      answer={answer}
      sources={sources}
      isLoading={isLoading}
      phase={phase}
      error={error}
    />
  )
}