import { useState, useEffect, useMemo, useRef } from 'react'
import VoiceViewer from './VoiceViewer'
import './TrajectoryVisualizer.css'

const SUBMISSIONS_BASE = import.meta.env.VITE_SUBMISSIONS_BASE_URL
  || `${import.meta.env.BASE_URL}submissions`

const NO_CACHE = { cache: 'no-cache' }

const S3_BUCKET = 'sierra-tau-bench-public'
const S3_SUBMISSIONS_PREFIX = 'submissions'

const TrajectoryVisualizer = () => {
  // --- Top-level selector state ---
  const [viewMode, setViewMode] = useState('trajectories') // 'trajectories' or 'tasks'
  const [submissions, setSubmissions] = useState([])
  const [submissionsLoading, setSubmissionsLoading] = useState(true)

  // Trajectory mode selectors
  const [selectedModelDir, setSelectedModelDir] = useState('')
  const [selectedDomain, setSelectedDomain] = useState('')
  const [selectedTrialIdx, setSelectedTrialIdx] = useState(0)
  const [selectedTaskId, setSelectedTaskId] = useState(null)

  // Loaded trajectory data for current model+domain
  const [trajectoryData, setTrajectoryData] = useState(null)
  const [loading, setLoading] = useState(false)
  const [error, setError] = useState(null)

  // Task mode state
  const [taskData, setTaskData] = useState(null)
  const [selectedDomainTask, setSelectedDomainTask] = useState(null)
  const [selectedTaskDetail, setSelectedTaskDetail] = useState(null)

  // Configuration modal
  const [showConfigModal, setShowConfigModal] = useState(false)
  const [modalClosing, setModalClosing] = useState(false)

  // CLI download popover
  const [showCliDownload, setShowCliDownload] = useState(false)
  const [cliCopied, setCliCopied] = useState(false)

  // Document viewer state
  const [showDocModal, setShowDocModal] = useState(false)
  const [docModalClosing, setDocModalClosing] = useState(false)
  const [selectedDoc, setSelectedDoc] = useState(null)
  const [docLoading, setDocLoading] = useState(false)

  const handleCloseModal = () => {
    setModalClosing(true)
    setTimeout(() => {
      setShowConfigModal(false)
      setModalClosing(false)
    }, 200) // Match the CSS slideDown duration (0.2s)
  }

  const handleCloseDocModal = () => {
    setDocModalClosing(true)
    setTimeout(() => {
      setShowDocModal(false)
      setDocModalClosing(false)
      setSelectedDoc(null)
    }, 200)
  }

  const handleDocClick = async (docId) => {
    try {
      setDocLoading(true)
      setShowDocModal(true)
      setSelectedDoc(null)
      const domain = selectedDomainTask || 'banking_knowledge'
      const res = await fetch(`${import.meta.env.BASE_URL}task-data/domains/${domain}/documents/${docId}.json`)
      if (!res.ok) throw new Error(`Failed to load document: ${res.statusText}`)
      const doc = await res.json()
      setSelectedDoc(doc)
    } catch (err) {
      setSelectedDoc({ id: docId, title: 'Error', content: `Failed to load document: ${err.message}` })
    } finally {
      setDocLoading(false)
    }
  }

  // Available domains
  const domains = [
    { id: 'airline', label: 'Airline', icon: '✈️', color: '#3b82f6' },
    { id: 'retail', label: 'Retail', icon: '🛍️', color: '#8b5cf6' },
    { id: 'telecom', label: 'Telecom', icon: '📱', color: '#059669' },
    { id: 'banking_knowledge', label: 'Banking', icon: '🏦', color: '#d97706' }
  ]

  // --- Parse URL params for deep linking ---
  const getUrlParams = () => {
    const hash = window.location.hash || ''
    const qIdx = hash.indexOf('?')
    if (qIdx === -1) return {}
    const params = new URLSearchParams(hash.slice(qIdx + 1))
    const trialRaw = params.get('trial')
    return {
      model: params.get('model'),
      domain: params.get('domain'),
      task: params.get('task'),
      trial: trialRaw != null ? Number(trialRaw) : null,
      view: params.get('view'),
    }
  }

  // Ref to hold URL params that need deferred restoration (task/trial depend on data loading)
  const pendingUrlParams = useRef(null)
  // Suppress URL updates while we're restoring from URL
  const restoringFromUrl = useRef(false)

  // --- Sync state → URL hash (replaceState to avoid history clutter) ---
  useEffect(() => {
    if (restoringFromUrl.current) return
    if (submissionsLoading) return

    const params = new URLSearchParams()
    if (selectedModelDir) params.set('model', selectedModelDir)
    if (selectedDomain) params.set('domain', selectedDomain)
    if (viewMode === 'tasks') params.set('view', 'tasks')
    if (selectedTaskId != null) params.set('task', String(selectedTaskId))
    if (selectedTrialIdx > 0) params.set('trial', String(selectedTrialIdx))

    const qs = params.toString()
    const newHash = qs ? `#trajectory-visualizer?${qs}` : '#trajectory-visualizer'
    if (window.location.hash !== newHash) {
      window.history.replaceState(null, '', newHash)
    }
  }, [selectedModelDir, selectedDomain, selectedTaskId, selectedTrialIdx, viewMode, submissionsLoading])

  // --- Load submissions on mount ---
  useEffect(() => {
    const loadSubmissions = async () => {
      try {
        setSubmissionsLoading(true)
        const res = await fetch(`${SUBMISSIONS_BASE}/manifest.json`, NO_CACHE)
        if (!res.ok) throw new Error('Failed to load manifest')
        const manifest = await res.json()
        const textDirs = manifest.submissions || []
        const voiceDirs = manifest.voice_submissions || []

        const loaded = []

        const loadDir = async (dir, modality) => {
          try {
            const r = await fetch(`${SUBMISSIONS_BASE}/${dir}/submission.json`, NO_CACHE)
            if (!r.ok) return
            const sub = await r.json()
            if (sub.trajectories_available && sub.trajectory_files) {
              loaded.push({
                dir,
                model_name: sub.model_name,
                model_organization: sub.model_organization || '',
                reasoning_effort: sub.reasoning_effort || null,
                submission_type: sub.submission_type || 'standard',
                user_simulator: sub.methodology?.user_simulator || null,
                trajectory_files: sub.trajectory_files,
                availableDomains: Object.keys(sub.trajectory_files),
                modality,
                voice_config: sub.voice_config || null,
              })
            }
          } catch { /* skip */ }
        }

        for (const dir of textDirs) await loadDir(dir, 'text')
        for (const dir of voiceDirs) await loadDir(dir, 'voice')

        loaded.sort((a, b) => {
          if (a.modality !== b.modality) return a.modality === 'text' ? -1 : 1
          return a.model_name.localeCompare(b.model_name)
        })
        setSubmissions(loaded)

        // Restore state from URL params
        const urlParams = getUrlParams()
        if (urlParams.model) {
          restoringFromUrl.current = true
          const match = loaded.find(s => s.dir === urlParams.model)
          if (match) {
            setSelectedModelDir(match.dir)
            const dom = urlParams.domain && match.availableDomains.includes(urlParams.domain)
              ? urlParams.domain
              : match.availableDomains[0]
            setSelectedDomain(dom || '')
            if (urlParams.view === 'tasks') {
              setViewMode('tasks')
            } else {
              setViewMode('trajectories')
            }
            // task and trial need trajectory data to be loaded first
            if (urlParams.task != null || urlParams.trial != null) {
              pendingUrlParams.current = { task: urlParams.task, trial: urlParams.trial }
            } else {
              restoringFromUrl.current = false
            }
          } else {
            restoringFromUrl.current = false
            if (loaded.length > 0) {
              setSelectedModelDir(loaded[0].dir)
              setSelectedDomain(loaded[0].availableDomains[0] || '')
            }
          }
        } else if (loaded.length > 0) {
          setSelectedModelDir(loaded[0].dir)
          setSelectedDomain(loaded[0].availableDomains[0] || '')
        }
      } catch (err) {
        setError(err.message)
      } finally {
        setSubmissionsLoading(false)
      }
    }
    loadSubmissions()
  }, [])

  // --- Load trajectory data when model+domain changes ---
  useEffect(() => {
    if (viewMode !== 'trajectories' || !selectedModelDir || !selectedDomain) return

    const sub = submissions.find(s => s.dir === selectedModelDir)
    if (!sub) return
    const fileName = sub.trajectory_files[selectedDomain]
    if (!fileName) return

    const loadData = async () => {
      try {
        setLoading(true)
        setError(null)
        setTrajectoryData(null)
        setSelectedTaskId(null)
        setSelectedTrialIdx(0)

        // Voice trajectories are directories containing results.json;
        // text trajectories are flat JSON files.
        const url = sub.modality === 'voice'
          ? `${SUBMISSIONS_BASE}/${sub.dir}/trajectories/${fileName}/results.json`
          : `${SUBMISSIONS_BASE}/${sub.dir}/trajectories/${fileName}`
        const res = await fetch(url, NO_CACHE)
        if (!res.ok) throw new Error(`Failed to load trajectory: ${res.statusText}`)
        const data = await res.json()

        // For dir-format voice results, simulations are stored in separate
        // files and results.json only contains a simulation_index. Synthesize
        // the simulations array from the index so existing UI logic works.
        if ((!data.simulations || data.simulations.length === 0) && data.simulation_index) {
          data.simulations = data.simulation_index.map(entry => ({
            id: entry.id,
            task_id: entry.task_id,
            trial: entry.trial,
            reward_info: entry.reward != null ? { reward: entry.reward } : null,
            termination_reason: entry.termination_reason || null,
            agent_cost: entry.agent_cost ?? null,
            duration: entry.duration ?? null,
          }))
        }

        setTrajectoryData(data)

        // Apply pending URL params (task/trial) now that data is loaded.
        // These state updates are batched with setTrajectoryData above.
        if (pendingUrlParams.current) {
          const { task, trial } = pendingUrlParams.current
          pendingUrlParams.current = null
          if (task != null) {
            const taskExists = data.tasks?.some(t => String(t.id) === String(task))
              || data.simulations?.some(s => String(s.task_id) === String(task))
            if (taskExists) {
              // Use the actual ID from the data to preserve its native type
              const actualTask = data.tasks?.find(t => String(t.id) === String(task))
              const actualSim = !actualTask && data.simulations?.find(s => String(s.task_id) === String(task))
              setSelectedTaskId(actualTask ? actualTask.id : actualSim ? actualSim.task_id : task)
              if (trial != null && trial > 0) {
                setSelectedTrialIdx(trial)
              }
            }
          }
          restoringFromUrl.current = false
        }
      } catch (err) {
        setError(err.message)
        if (pendingUrlParams.current) {
          pendingUrlParams.current = null
          restoringFromUrl.current = false
        }
      } finally {
        setLoading(false)
      }
    }
    loadData()
  }, [selectedModelDir, selectedDomain, viewMode, submissions])

  // --- Derived data ---
  const currentSubmission = useMemo(
    () => submissions.find(s => s.dir === selectedModelDir),
    [submissions, selectedModelDir]
  )

  // All unique tasks from trajectory data
  const tasks = useMemo(() => {
    if (!trajectoryData?.tasks) return []
    return trajectoryData.tasks
  }, [trajectoryData])

  // Number of distinct trials available
  const numTrials = useMemo(() => {
    if (!trajectoryData?.simulations) return 0
    const trials = new Set(trajectoryData.simulations.map(s => s.trial))
    return trials.size
  }, [trajectoryData])

  // Simulations for the selected task, sorted by trial
  const taskSimulations = useMemo(() => {
    if (!trajectoryData?.simulations || selectedTaskId === null) return []
    return trajectoryData.simulations
      .filter(s => s.task_id === selectedTaskId)
      .sort((a, b) => a.trial - b.trial)
  }, [trajectoryData, selectedTaskId])

  // The current simulation (for the selected trial)
  const currentSimulation = useMemo(() => {
    if (!taskSimulations.length) return null
    return taskSimulations[selectedTrialIdx] || taskSimulations[0]
  }, [taskSimulations, selectedTrialIdx])

  // Task metadata for the selected task
  const currentTask = useMemo(() => {
    if (selectedTaskId === null || !tasks.length) return null
    return tasks.find(t => t.id === selectedTaskId) || null
  }, [tasks, selectedTaskId])

  // Compute pass^k scores from trajectory data
  const passKScores = useMemo(() => {
    if (!trajectoryData?.simulations || !tasks.length) return null
    // Derive numTrials from actual data rather than info metadata (which can be wrong)
    const trialSet = new Set(trajectoryData.simulations.map(s => s.trial))
    const numTrials = trialSet.size
    if (numTrials === 0) return null

    // Helper: binomial coefficient C(n, k)
    const comb = (n, k) => {
      if (k < 0 || k > n) return 0
      if (k === 0 || k === n) return 1
      let result = 1
      for (let i = 0; i < Math.min(k, n - k); i++) {
        result = result * (n - i) / (i + 1)
      }
      return Math.round(result)
    }

    // Group simulations by task_id and count successes (reward ≈ 1.0)
    const taskSuccesses = {}
    const taskTrialCounts = {}
    for (const sim of trajectoryData.simulations) {
      const tid = sim.task_id
      if (!taskTrialCounts[tid]) { taskTrialCounts[tid] = 0; taskSuccesses[tid] = 0 }
      taskTrialCounts[tid]++
      const reward = sim.reward_info?.reward ?? 0
      if (Math.abs(reward - 1.0) < 1e-6) taskSuccesses[tid]++
    }

    const taskIds = Object.keys(taskTrialCounts)
    const n = taskIds.length
    if (n === 0) return null

    // Compute pass^k for k = 1..numTrials
    const maxK = Math.min(numTrials, Math.min(...Object.values(taskTrialCounts)))
    const scores = {}
    for (let k = 1; k <= maxK; k++) {
      let sum = 0
      for (const tid of taskIds) {
        sum += comb(taskSuccesses[tid], k) / comb(taskTrialCounts[tid], k)
      }
      scores[k] = (sum / n) * 100 // as percentage
    }
    return scores
  }, [trajectoryData, tasks])

  const isVoice = currentSubmission?.modality === 'voice'

  // Build task info for the voice viewer from the current task's user scenario
  const voiceTaskInfo = useMemo(() => {
    if (!currentTask?.user_scenario?.instructions) return {}
    const instr = currentTask.user_scenario.instructions
    if (typeof instr === 'string') return { reason: instr }
    return {
      reason: instr.reason_for_call || '',
      knownInfo: instr.known_info || '',
      unknownInfo: instr.unknown_info || '',
      taskInstructions: instr.task_instructions || '',
    }
  }, [currentTask])

  const handleDownload = () => {
    if (!trajectoryData || !currentSubmission) return
    const blob = new Blob([JSON.stringify(trajectoryData, null, 2)], { type: 'application/json' })
    const url = URL.createObjectURL(blob)
    const a = document.createElement('a')
    a.href = url
    a.download = `${currentSubmission.model_name}_${selectedDomain}_trajectories.json`
    document.body.appendChild(a)
    a.click()
    document.body.removeChild(a)
    URL.revokeObjectURL(url)
  }

  const getVoiceS3Command = () => {
    if (!currentSubmission || !selectedDomain) return ''
    const fileName = currentSubmission.trajectory_files[selectedDomain]
    if (!fileName) return ''
    const s3Path = `s3://${S3_BUCKET}/${S3_SUBMISSIONS_PREFIX}/${currentSubmission.dir}/trajectories/${fileName}/`
    const localDir = `./${currentSubmission.model_name}_${selectedDomain}/`
    return `aws s3 sync ${s3Path} ${localDir} --no-sign-request`
  }

  const handleCopyCliCommand = () => {
    const cmd = getVoiceS3Command()
    navigator.clipboard.writeText(cmd).then(() => {
      setCliCopied(true)
      setTimeout(() => setCliCopied(false), 2000)
    })
  }

  // --- Handlers ---
  const handleModelChange = (dir) => {
    setSelectedModelDir(dir)
    setSelectedTaskId(null)
    setSelectedTrialIdx(0)
    setShowCliDownload(false)
    const sub = submissions.find(s => s.dir === dir)
    if (sub && sub.availableDomains.length > 0) {
      // Keep current domain if available, else pick first
      if (!sub.availableDomains.includes(selectedDomain)) {
        setSelectedDomain(sub.availableDomains[0])
      }
    }
  }

  const handleDomainChange = (domain) => {
    setSelectedDomain(domain)
    setSelectedTaskId(null)
    setSelectedTrialIdx(0)
    setShowCliDownload(false)
  }

  // --- Task mode ---
  const loadTaskData = async (domain) => {
    try {
      setLoading(true)
      setError(null)
      const [tasksRes, policyRes] = await Promise.all([
        fetch(`${import.meta.env.BASE_URL}task-data/domains/${domain}/tasks.json`),
        fetch(`${import.meta.env.BASE_URL}task-data/domains/${domain}/policy.md`)
      ])
      if (!tasksRes.ok) throw new Error(`Failed to load tasks: ${tasksRes.statusText}`)
      const tasksJson = await tasksRes.json()
      const policy = policyRes.ok ? await policyRes.text() : null
      setTaskData({ tasks: tasksJson, policy, domain })
      setSelectedDomainTask(domain)
      setSelectedTaskDetail(null)
    } catch (err) {
      setError(err.message)
    } finally {
      setLoading(false)
    }
  }

  // --- Helpers ---
  // Get a meaningful task description, preferring notes over generic "Task: task_XXX" purpose
  const getTaskDescription = (task) => {
    const purpose = task?.description?.purpose
    const notes = task?.description?.notes
    if (notes) return notes
    if (purpose && !/^Task:\s*task_\d+$/i.test(purpose)) return purpose
    return null
  }

  // Extract user scenario text from either string or object format
  const getUserScenarioText = (task) => {
    const instructions = task?.user_scenario?.instructions
    if (!instructions) return null
    if (typeof instructions === 'string') return instructions
    return instructions.reason_for_call || instructions.task_instructions || null
  }

  const getCleanTaskId = (taskId) => {
    if (!taskId && taskId !== 0) return 'Unknown'
    if (typeof taskId === 'number' || /^\d+$/.test(String(taskId))) return String(taskId)
    // Handle "task_001" format — extract just the number
    const taskNumMatch = String(taskId).match(/^task_0*(\d+)$/)
    if (taskNumMatch) return taskNumMatch[1]
    if (String(taskId).length < 15) return String(taskId)
    const bracketMatch = String(taskId).match(/\[([^\]]+)\]/)
    if (bracketMatch) {
      return bracketMatch[1].replace(/_/g, ' ').replace(/\b\w/g, l => l.toUpperCase())
    }
    const cleaned = String(taskId).split(/[[\|]/)[0].replace(/_/g, ' ')
    return cleaned.charAt(0).toUpperCase() + cleaned.slice(1)
  }

  const formatMessage = (message) => ({
    role: message.role,
    content: message.content,
    tool_calls: message.tool_calls,
    turn: message.turn_idx,
    timestamp: new Date(message.timestamp).toLocaleString(),
    cost: message.cost || 0,
    tokens: message.usage ? `${message.usage.prompt_tokens || 0}/${message.usage.completion_tokens || 0}` : 'N/A'
  })

  const getDisplayMessages = (sim) => {
    if (!sim?.messages) return []
    return sim.messages.slice(0, 60).map(formatMessage)
  }

  // --- Render ---
  return (
    <div className="trajectory-visualizer">
      <div className="visualizer-header">
        <h2>{isVoice ? 'τ-voice Visualizer' : 'τ-bench Visualizer'}</h2>
        <p className="visualizer-description">
          {isVoice
            ? 'Explore τ-voice results: view task outcomes for audio-native voice agent evaluations, or examine the underlying task definitions.'
            : 'Explore τ-bench dataset: view conversation trajectories showing AI agent interactions with users, or examine the underlying task definitions across airline, retail, telecom, and banking domains.'
          }
        </p>

        {/* View Mode Toggle */}
        <div className="view-toggle">
          <button
            className={`toggle-btn ${viewMode === 'trajectories' ? 'active' : ''}`}
            onClick={() => {
              setViewMode('trajectories')
              setTaskData(null)
              setSelectedTaskDetail(null)
              setSelectedDomainTask(null)
            }}
          >
            🔄 Trajectories
          </button>
          <button
            className={`toggle-btn ${viewMode === 'tasks' ? 'active' : ''}`}
            onClick={() => {
              setViewMode('tasks')
              setSelectedTaskId(null)
              setTrajectoryData(null)
            }}
          >
            📋 Tasks
          </button>
        </div>
      </div>

      {viewMode === 'trajectories' && (
        <>
          {/* ===== SELECTOR BAR ===== */}
          <div className="selector-bar">
            {/* Model dropdown */}
            <div className="selector-group">
              <label className="selector-label">Model</label>
              <select
                className="selector-dropdown"
                value={selectedModelDir}
                onChange={e => handleModelChange(e.target.value)}
                disabled={submissionsLoading}
              >
                {submissionsLoading && <option value="">Loading...</option>}
                {submissions.some(s => s.modality === 'text') && (
                  <optgroup label="τ-bench (Text)">
                    {submissions.filter(s => s.modality === 'text').map(s => (
                      <option key={s.dir} value={s.dir}>
                        {s.model_name}{s.reasoning_effort ? ` [${s.reasoning_effort.charAt(0).toUpperCase() + s.reasoning_effort.slice(1)}]` : ''} ({s.model_organization})
                      </option>
                    ))}
                  </optgroup>
                )}
                {submissions.some(s => s.modality === 'voice') && (
                  <optgroup label="τ-voice (Voice)">
                    {submissions.filter(s => s.modality === 'voice').map(s => (
                      <option key={s.dir} value={s.dir}>
                        {s.model_name} ({s.model_organization}){s.submission_type === 'custom' ? ` [Custom]${s.user_simulator ? ` — user sim: ${s.user_simulator}` : ''}` : ''}
                      </option>
                    ))}
                  </optgroup>
                )}
              </select>
            </div>

            {/* Domain dropdown */}
            <div className="selector-group">
              <label className="selector-label">Domain</label>
              <div className="domain-pills">
                {domains.map(d => {
                  const available = currentSubmission?.availableDomains?.includes(d.id)
                  return (
                    <button
                      key={d.id}
                      className={`domain-pill ${selectedDomain === d.id ? 'active' : ''} ${!available ? 'disabled' : ''}`}
                      style={selectedDomain === d.id ? { background: d.color, borderColor: d.color } : {}}
                      onClick={() => available && handleDomainChange(d.id)}
                      disabled={!available}
                    >
                      {d.icon} {d.label}
                    </button>
                  )
                })}
              </div>
            </div>

            {/* Trial selector */}
            {numTrials > 0 && selectedTaskId !== null && (
              <div className="selector-group">
                <label className="selector-label">Trial</label>
                <div className="trial-pills">
                  {Array.from({ length: numTrials }, (_, i) => (
                    <button
                      key={i}
                      className={`trial-pill ${selectedTrialIdx === i ? 'active' : ''}`}
                      onClick={() => setSelectedTrialIdx(i)}
                    >
                      {i + 1}
                    </button>
                  ))}
                </div>
              </div>
            )}

            {/* Config button */}
            {trajectoryData?.info && (
              <div className="selector-group selector-group-end">
                <button className="config-button" onClick={() => setShowConfigModal(true)}>
                  ⚙️ Config
                </button>
              </div>
            )}
          </div>

          {/* ===== MAIN CONTENT ===== */}
          <div className="viz-content">
            {loading && (
              <div className="loading-state">
                <div className="loading-spinner"></div>
                <p>Loading trajectory data...</p>
                <p className="loading-note">Large files may take a moment</p>
              </div>
            )}

            {error && !loading && (
              <div className="error-state">
                <p>⚠️ {error}</p>
              </div>
            )}

            {!loading && !error && !trajectoryData && !submissionsLoading && (
              <div className="empty-state">
                <h3>Select a Model &amp; Domain</h3>
                <p>Choose a model and domain above to explore trajectory data.</p>
              </div>
            )}

            {/* Task list view — no task selected yet */}
            {!loading && !error && trajectoryData && selectedTaskId === null && (
              <div className="task-list-view">
                <div className="task-list-header">
                  <div className="task-list-header-left">
                    <h3>
                      {isVoice && <span className="voice-badge">🎙️ Voice</span>}
                      {currentSubmission?.model_name}{currentSubmission?.reasoning_effort ? ` [${currentSubmission.reasoning_effort.charAt(0).toUpperCase() + currentSubmission.reasoning_effort.slice(1)}]` : ''} — {selectedDomain.charAt(0).toUpperCase() + selectedDomain.slice(1)}
                    </h3>
                    <p className="task-list-subtitle">
                      {tasks.length} tasks · {numTrials} trial{numTrials !== 1 ? 's' : ''} each
                      {isVoice ? '' : ' · Select a task to view conversations'}
                    </p>
                  </div>
                  <div className="task-list-header-right">
                    {trajectoryData && !isVoice && (
                      <button className="download-btn" onClick={handleDownload} title="Download raw trajectory data as JSON">
                        ⬇ Download JSON
                      </button>
                    )}
                    {trajectoryData && isVoice && (
                      <div className="cli-download-wrapper">
                        <button
                          className="download-btn"
                          onClick={() => setShowCliDownload(!showCliDownload)}
                          title="Voice trajectory data is too large for browser download. Use AWS CLI instead."
                        >
                          ⬇ Download via CLI
                        </button>
                        {showCliDownload && (
                          <div className="cli-download-popover">
                            <p className="cli-download-note">Voice trajectories are large and must be downloaded via CLI:</p>
                            <pre className="cli-download-command">{getVoiceS3Command()}</pre>
                            <div className="cli-download-actions">
                              <button className="cli-copy-btn" onClick={handleCopyCliCommand}>
                                {cliCopied ? '✓ Copied' : 'Copy command'}
                              </button>
                              <button className="cli-close-btn" onClick={() => setShowCliDownload(false)}>
                                Close
                              </button>
                            </div>
                          </div>
                        )}
                      </div>
                    )}
                    {passKScores && (
                      <div className="pass-k-scores">
                        {Object.entries(passKScores).map(([k, score]) => (
                          <div key={k} className="pass-k-item">
                            <span className="pass-k-label">pass^{k}</span>
                            <span className={`pass-k-value ${score >= 50 ? 'good' : score >= 25 ? 'mid' : 'low'}`}>
                              {score.toFixed(1)}%
                            </span>
                          </div>
                        ))}
                      </div>
                    )}
                  </div>
                </div>

                <div className="task-grid">
                  {tasks.map((task, idx) => {
                    const sims = trajectoryData.simulations?.filter(s => s.task_id === task.id) || []
                    const avgReward = sims.length > 0
                      ? (sims.reduce((sum, s) => sum + (s.reward_info?.reward || 0), 0) / sims.length)
                      : null
                    const desc = getTaskDescription(task)
                    const scenario = getUserScenarioText(task)

                    return (
                      <div
                        key={task.id ?? idx}
                        className="task-card"
                        onClick={() => { setSelectedTaskId(task.id); setSelectedTrialIdx(0) }}
                      >
                        <div className="task-header">
                          <span className="task-id">Task {getCleanTaskId(task.id)}</span>
                          {avgReward !== null && (
                            <span className={`task-reward ${avgReward >= 0.5 ? 'good' : 'bad'}`}>
                              {avgReward.toFixed(2)}
                            </span>
                          )}
                        </div>
                        <div className="task-description">
                          <p>{desc ? desc.slice(0, 120) + (desc.length > 120 ? '…' : '') : (task.description?.purpose || 'No description')}</p>
                        </div>
                        <div className="task-stats">
                          <span>{sims.length} trial{sims.length !== 1 ? 's' : ''}</span>
                          <span>{scenario ? '📞 ' + scenario.slice(0, 50) + (scenario.length > 50 ? '…' : '') : ''}</span>
                        </div>
                      </div>
                    )
                  })}
                </div>
              </div>
            )}

            {/* Conversation view — task selected */}
            {!loading && !error && trajectoryData && selectedTaskId !== null && currentSimulation && (
              <div className="conversation-view">
                <div className="conversation-header">
                  <div className="conversation-meta">
                    <button className="back-button" onClick={() => setSelectedTaskId(null)}>
                      ← All Tasks
                    </button>
                    <h3>Task {getCleanTaskId(selectedTaskId)}</h3>
                    <div className="trial-indicator">
                      Trial {selectedTrialIdx + 1} of {numTrials}
                    </div>
                  </div>

                  {/* Trial selector pills inline */}
                  {numTrials > 1 && (
                    <div className="trial-selector-inline">
                      {Array.from({ length: numTrials }, (_, i) => {
                        const sim = taskSimulations[i]
                        const reward = sim?.reward_info?.reward
                        return (
                          <button
                            key={i}
                            className={`trial-pill-inline ${selectedTrialIdx === i ? 'active' : ''} ${reward !== undefined ? (reward >= 0.5 ? 'success' : 'fail') : ''}`}
                            onClick={() => setSelectedTrialIdx(i)}
                          >
                            Trial {i + 1}
                            {reward !== undefined && <span className="trial-reward">{reward.toFixed(1)}</span>}
                          </button>
                        )
                      })}
                    </div>
                  )}

                  {/* Task context */}
                  {currentTask && (
                    <div className="task-context">
                      <h4>Task Context</h4>
                      <p><strong>Purpose:</strong> {currentTask.description?.purpose}</p>
                      {typeof currentTask.user_scenario?.instructions === 'string' ? (
                        <p><strong>User Scenario:</strong> {currentTask.user_scenario.instructions.slice(0, 300)}{currentTask.user_scenario.instructions.length > 300 ? '…' : ''}</p>
                      ) : (
                        <>
                          <p><strong>User Scenario:</strong> {currentTask.user_scenario?.instructions?.reason_for_call}</p>
                          <p><strong>Known Info:</strong> {currentTask.user_scenario?.instructions?.known_info}</p>
                        </>
                      )}
                    </div>
                  )}

                  {/* Simulation results */}
                  <div className="simulation-results">
                    <h4>Results</h4>
                    <div className="results-grid">
                      <div className="result-item">
                        <span className="result-label">Reward</span>
                        <span className={`result-value ${(currentSimulation.reward_info?.reward ?? 0) >= 0.5 ? 'good' : 'bad'}`}>
                          {currentSimulation.reward_info?.reward?.toFixed(2) ?? 'N/A'}
                        </span>
                      </div>
                      <div className="result-item">
                        <span className="result-label">Termination</span>
                        <span className="result-value">{currentSimulation.termination_reason || '—'}</span>
                      </div>
                      <div className="result-item">
                        <span className="result-label">Agent Cost</span>
                        <span className="result-value">${currentSimulation.agent_cost?.toFixed(4) ?? '—'}</span>
                      </div>
                      <div className="result-item">
                        <span className="result-label">Duration</span>
                        <span className="result-value">{currentSimulation.duration ? `${Math.round(currentSimulation.duration)}s` : '—'}</span>
                      </div>
                    </div>

                    {currentSimulation.reward_info?.nl_assertions?.length > 0 && (
                      <div className="assertions">
                        <h5>Evaluation Assertions</h5>
                        <div className="assertion-list">
                          {currentSimulation.reward_info.nl_assertions.map((a, i) => (
                            <div key={i} className={`assertion ${a.met ? 'passed' : 'failed'}`}>
                              <span className="assertion-status">{a.met ? '✅' : '❌'}</span>
                              <span className="assertion-text">{a.nl_assertion}</span>
                              {a.justification && <p className="assertion-justification">{a.justification}</p>}
                            </div>
                          ))}
                        </div>
                      </div>
                    )}
                  </div>
                </div>

                {/* Voice: timeline + conversation table; Text: message bubbles */}
                {isVoice ? (
                  <VoiceViewer
                    key={`${currentSimulation.id}-${selectedTrialIdx}`}
                    submissionDir={currentSubmission.dir}
                    trajectoryDir={currentSubmission.trajectory_files[selectedDomain]}
                    simulationId={currentSimulation.id}
                    taskId={selectedTaskId}
                    voiceConfig={currentSubmission.voice_config}
                    taskInfo={voiceTaskInfo}
                  />
                ) : (
                  <div className="conversation-messages">
                    {getDisplayMessages(currentSimulation).map((msg, i) => (
                      <div key={i} className={`message ${msg.role}`}>
                        <div className="message-header">
                          <span className="message-role">
                            {msg.role === 'assistant' ? '🤖 Agent' : msg.role === 'tool' ? '🔧 Tool' : '👤 User'}
                          </span>
                          <span className="message-turn">Turn {msg.turn}</span>
                          {msg.cost > 0 && <span className="message-cost">${msg.cost.toFixed(4)}</span>}
                          <span className="message-tokens">{msg.tokens} tokens</span>
                        </div>
                        <div className="message-content">{msg.content}</div>
                        {msg.tool_calls && (
                          <div className="message-tools">
                            <strong>Tool Calls:</strong>
                            <pre>{JSON.stringify(msg.tool_calls, null, 2)}</pre>
                          </div>
                        )}
                      </div>
                    ))}
                    {currentSimulation.messages?.length > 60 && (
                      <div className="message-truncated">
                        <p>… and {currentSimulation.messages.length - 60} more messages (showing first 60)</p>
                      </div>
                    )}
                  </div>
                )}
              </div>
            )}
          </div>
        </>
      )}

      {/* ===== TASK MODE ===== */}
      {viewMode === 'tasks' && (
        <div className="task-mode-container">
          <div className="selector-bar">
            <div className="selector-group">
              <label className="selector-label">Domain</label>
              <div className="domain-pills">
                {domains.map(d => (
                  <button
                    key={d.id}
                    className={`domain-pill ${selectedDomainTask === d.id ? 'active' : ''}`}
                    style={selectedDomainTask === d.id ? { background: d.color, borderColor: d.color } : {}}
                    onClick={() => loadTaskData(d.id)}
                  >
                    {d.icon} {d.label}
                  </button>
                ))}
              </div>
            </div>
          </div>

          <div className="viz-content">
            {loading && (
              <div className="loading-state">
                <div className="loading-spinner"></div>
                <p>Loading task data...</p>
              </div>
            )}

            {error && !loading && (
              <div className="error-state"><p>⚠️ {error}</p></div>
            )}

            {!loading && !error && !taskData && (
              <div className="empty-state">
                <h3>Select a Domain</h3>
                <p>Choose a domain above to explore task definitions and agent policies.</p>
              </div>
            )}

            {/* Task overview list */}
            {!loading && taskData && !selectedTaskDetail && (
              <div className="task-list-view">
                <div className="task-list-header">
                  <h3>{domains.find(d => d.id === taskData.domain)?.label || taskData.domain} Tasks</h3>
                  <p className="task-list-subtitle">{taskData.tasks?.length || 0} task definitions</p>
                </div>
                <div className="task-grid">
                  {taskData.tasks?.map((task, idx) => {
                    const scenarioText = getUserScenarioText(task)
                    const description = getTaskDescription(task)
                    return (
                      <div key={task.id || idx} className="task-card" onClick={() => setSelectedTaskDetail(task)}>
                        <div className="task-header">
                          <span className="task-id">Task {getCleanTaskId(task.id)}</span>
                          {task.required_documents?.length > 0 && (
                            <span className="task-badge docs-badge">📄 {task.required_documents.length} docs</span>
                          )}
                        </div>
                        <div className="task-description">
                          {description ? (
                            <p>{description.slice(0, 180)}{description.length > 180 ? '…' : ''}</p>
                          ) : (
                            <p><strong>Scenario:</strong> {scenarioText ? scenarioText.slice(0, 120) + (scenarioText.length > 120 ? '…' : '') : '—'}</p>
                          )}
                        </div>
                        <div className="task-stats">
                          <span>{task.evaluation_criteria?.actions?.length || 0} actions</span>
                          <span>{task.evaluation_criteria?.nl_assertions?.length || 0} assertions</span>
                          {task.evaluation_criteria?.communicate_info?.length > 0 && (
                            <span>{task.evaluation_criteria.communicate_info.length} info items</span>
                          )}
                          {task.user_tools?.length > 0 && (
                            <span>🔧 {task.user_tools.length} user tools</span>
                          )}
                        </div>
                      </div>
                    )
                  })}
                </div>
              </div>
            )}

            {/* Task detail */}
            {!loading && taskData && selectedTaskDetail && (
              <div className="task-detail-view">
                <div className="task-detail-header">
                  <button className="back-button" onClick={() => setSelectedTaskDetail(null)}>← Back to Tasks</button>
                  <h3>Task {getCleanTaskId(selectedTaskDetail.id)}</h3>
                </div>

                <div className="task-detail-content">
                  <div className="task-section">
                    <h4>Task Description</h4>
                    <div className="task-info">
                      {selectedTaskDetail.description?.notes ? (
                        <p>{selectedTaskDetail.description.notes}</p>
                      ) : (
                        <p>{selectedTaskDetail.description?.purpose || '—'}</p>
                      )}
                    </div>
                  </div>

                  <div className="task-section">
                    <h4>User Scenario</h4>
                    <div className="task-info">
                      {typeof selectedTaskDetail.user_scenario?.instructions === 'string' ? (
                        <>
                          {selectedTaskDetail.user_scenario?.persona && (
                            <p><strong>Persona:</strong> {selectedTaskDetail.user_scenario.persona}</p>
                          )}
                          <pre className="instructions-text">{selectedTaskDetail.user_scenario.instructions}</pre>
                        </>
                      ) : (
                        <>
                          <p><strong>Domain:</strong> {selectedTaskDetail.user_scenario?.instructions?.domain}</p>
                          <p><strong>Reason for Call:</strong> {selectedTaskDetail.user_scenario?.instructions?.reason_for_call}</p>
                          <p><strong>Known Information:</strong> {selectedTaskDetail.user_scenario?.instructions?.known_info}</p>
                          {selectedTaskDetail.user_scenario?.instructions?.unknown_info && (
                            <p><strong>Unknown Information:</strong> {selectedTaskDetail.user_scenario.instructions.unknown_info}</p>
                          )}
                          {selectedTaskDetail.user_scenario?.instructions?.task_instructions && (
                            <div>
                              <p><strong>Task Instructions:</strong></p>
                              <pre className="instructions-text">{selectedTaskDetail.user_scenario.instructions.task_instructions}</pre>
                            </div>
                          )}
                        </>
                      )}
                    </div>
                  </div>

                  {selectedTaskDetail.user_tools?.length > 0 && (
                    <div className="task-section">
                      <h4>User Tools</h4>
                      <div className="user-tools-list">
                        {selectedTaskDetail.user_tools.map((tool, i) => (
                          <span key={i} className="user-tool-badge">{tool}</span>
                        ))}
                      </div>
                    </div>
                  )}

                  {selectedTaskDetail.required_documents?.length > 0 && (
                    <div className="task-section">
                      <h4>Required Documents ({selectedTaskDetail.required_documents.length})</h4>
                      <p className="doc-hint">Click a document to view its contents</p>
                      <div className="required-docs-list">
                        {selectedTaskDetail.required_documents.map((doc, i) => (
                          <button key={i} className="doc-badge clickable" onClick={() => handleDocClick(doc)}>
                            <span className="doc-icon">📄</span>
                            <span className="doc-name">{doc.replace(/^doc_/, '').replace(/_/g, ' ')}</span>
                            <span className="doc-arrow">→</span>
                          </button>
                        ))}
                      </div>
                    </div>
                  )}

                  {selectedTaskDetail.evaluation_criteria && (
                    <div className="task-section">
                      <h4>Evaluation Criteria</h4>
                      {selectedTaskDetail.evaluation_criteria.reward_basis?.length > 0 && (
                        <div className="criteria-subsection">
                          <h5>Reward Basis</h5>
                          <div className="reward-basis-list">
                            {selectedTaskDetail.evaluation_criteria.reward_basis.map((basis, i) => (
                              <span key={i} className="reward-basis-badge">{basis}</span>
                            ))}
                          </div>
                        </div>
                      )}
                      {selectedTaskDetail.evaluation_criteria.actions && (
                        <div className="criteria-subsection">
                          <h5>Expected Actions ({selectedTaskDetail.evaluation_criteria.actions.length})</h5>
                          <div className="actions-list">
                            {selectedTaskDetail.evaluation_criteria.actions.length > 0 ? (
                              selectedTaskDetail.evaluation_criteria.actions.map((action, i) => (
                                <div key={i} className="action-item">
                                  <p><strong>Action:</strong> {action.name}</p>
                                  {action.requestor && <p><strong>Requestor:</strong> {action.requestor}</p>}
                                  {action.arguments && <pre className="action-args">{JSON.stringify(action.arguments, null, 2)}</pre>}
                                </div>
                              ))
                            ) : (
                              <div className="no-actions-message"><p>Agent should not take any action</p></div>
                            )}
                          </div>
                        </div>
                      )}
                      {selectedTaskDetail.evaluation_criteria.communicate_info?.length > 0 && (
                        <div className="criteria-subsection">
                          <h5>Required Communication ({selectedTaskDetail.evaluation_criteria.communicate_info.length})</h5>
                          <div className="communicate-info-list">
                            {selectedTaskDetail.evaluation_criteria.communicate_info.map((info, i) => (
                              <div key={i} className="communicate-info-item">
                                <p>{typeof info === 'string' ? info : JSON.stringify(info)}</p>
                              </div>
                            ))}
                          </div>
                        </div>
                      )}
                      {selectedTaskDetail.evaluation_criteria.nl_assertions?.length > 0 && (
                        <div className="criteria-subsection">
                          <h5>Natural Language Assertions ({selectedTaskDetail.evaluation_criteria.nl_assertions.length}) <span className="experimental-badge-container"><span className="experimental-badge">experimental</span><div className="experimental-tooltip">These assertions are experimental and not used to compute benchmark scores</div></span></h5>
                          <div className="assertions-list">
                            {selectedTaskDetail.evaluation_criteria.nl_assertions.map((a, i) => (
                              <div key={i} className="assertion-item"><p>{a}</p></div>
                            ))}
                          </div>
                        </div>
                      )}
                      {selectedTaskDetail.evaluation_criteria.env_assertions?.length > 0 && (
                        <div className="criteria-subsection">
                          <h5>Environment Assertions ({selectedTaskDetail.evaluation_criteria.env_assertions.length})</h5>
                          <div className="env-assertions-list">
                            {selectedTaskDetail.evaluation_criteria.env_assertions.map((a, i) => (
                              <div key={i} className="env-assertion-item">
                                <p><strong>Function:</strong> {a.func_name}</p>
                                <pre className="assertion-args">{JSON.stringify(a.arguments, null, 2)}</pre>
                              </div>
                            ))}
                          </div>
                        </div>
                      )}
                    </div>
                  )}

                  {selectedTaskDetail.initial_state?.initialization_actions && (
                    <div className="task-section">
                      <h4>Initial State</h4>
                      <pre className="initial-actions">{JSON.stringify(selectedTaskDetail.initial_state.initialization_actions, null, 2)}</pre>
                    </div>
                  )}

                  {taskData?.policy && (
                    <div className="task-section">
                      <h4>Domain Policy</h4>
                      <pre className="policy-text">{taskData.policy}</pre>
                    </div>
                  )}
                </div>
              </div>
            )}
          </div>
        </div>
      )}

      {/* ===== DOCUMENT VIEWER MODAL ===== */}
      {showDocModal && (
        <div className="modal-overlay" onClick={handleCloseDocModal}>
          <div className={`modal-content doc-modal ${docModalClosing ? 'closing' : ''}`} onClick={e => e.stopPropagation()}>
            <div className="modal-header">
              <h3>{selectedDoc?.title || 'Loading document...'}</h3>
              <button className="modal-close" onClick={handleCloseDocModal}>✕</button>
            </div>
            <div className="doc-modal-body">
              {docLoading && (
                <div className="doc-loading">
                  <div className="loading-spinner"></div>
                  <p>Loading document...</p>
                </div>
              )}
              {!docLoading && selectedDoc && (
                <>
                  <div className="doc-id-label">
                    <code>{selectedDoc.id}</code>
                  </div>
                  <div className="doc-content-rendered">
                    {selectedDoc.content.split('\n').map((line, i) => {
                      // Simple markdown rendering for headings, lists, and text
                      if (line.startsWith('## ')) {
                        return <h3 key={i} className="doc-heading">{line.slice(3)}</h3>
                      }
                      if (line.startsWith('### ')) {
                        return <h4 key={i} className="doc-subheading">{line.slice(4)}</h4>
                      }
                      if (/^\d+\.\s/.test(line)) {
                        return <p key={i} className="doc-ordered-item">{line}</p>
                      }
                      if (line.startsWith('- ')) {
                        return <p key={i} className="doc-list-item">{line}</p>
                      }
                      if (line.startsWith('   - ')) {
                        return <p key={i} className="doc-list-item nested">{line.trim()}</p>
                      }
                      if (line.trim() === '') {
                        return <br key={i} />
                      }
                      return <p key={i} className="doc-paragraph">{line}</p>
                    })}
                  </div>
                </>
              )}
            </div>
          </div>
        </div>
      )}

      {/* ===== CONFIG MODAL ===== */}
      {showConfigModal && trajectoryData && (
        <div className="modal-overlay" onClick={handleCloseModal}>
          <div className={`modal-content config-modal ${modalClosing ? 'closing' : ''}`} onClick={e => e.stopPropagation()}>
            <div className="modal-header">
              <h3>Configuration</h3>
              <button className="modal-close" onClick={handleCloseModal}>✕</button>
            </div>
            <div className="config-table-body">
              <table className="config-table">
                <tbody>
                  {trajectoryData.info?.agent_info && (
                    <>
                      <tr className="config-section-header"><td colSpan="2">Agent</td></tr>
                      <tr><td>Implementation</td><td><code>{trajectoryData.info.agent_info.implementation}</code></td></tr>
                      <tr><td>Model</td><td><code>{trajectoryData.info.agent_info.llm}</code></td></tr>
                      {trajectoryData.info.agent_info.llm_args && Object.entries(trajectoryData.info.agent_info.llm_args).map(([k, v]) => (
                        <tr key={k}><td className="config-indent">{k}</td><td><code>{JSON.stringify(v)}</code></td></tr>
                      ))}
                    </>
                  )}
                  {trajectoryData.info?.user_info && (
                    <>
                      <tr className="config-section-header"><td colSpan="2">User Simulator</td></tr>
                      <tr><td>Implementation</td><td><code>{trajectoryData.info.user_info.implementation}</code></td></tr>
                      <tr><td>Model</td><td><code>{trajectoryData.info.user_info.llm}</code></td></tr>
                    </>
                  )}
                  {trajectoryData.info && (
                    <>
                      <tr className="config-section-header"><td colSpan="2">Evaluation</td></tr>
                      {trajectoryData.info.num_trials && <tr><td>Trials</td><td><code>{trajectoryData.info.num_trials}</code></td></tr>}
                      {trajectoryData.info.max_steps && <tr><td>Max Steps</td><td><code>{trajectoryData.info.max_steps}</code></td></tr>}
                      {trajectoryData.info.max_errors && <tr><td>Max Errors</td><td><code>{trajectoryData.info.max_errors}</code></td></tr>}
                      {trajectoryData.info.seed !== undefined && <tr><td>Seed</td><td><code>{trajectoryData.info.seed}</code></td></tr>}
                    </>
                  )}
                </tbody>
              </table>
            </div>
          </div>
        </div>
      )}
    </div>
  )
}

export default TrajectoryVisualizer
