import { createContext, useContext, useState, useMemo, type ReactNode } from 'react';
import type { DisplayMessage } from './ChatContext';
import type { PlaygroundSessionResponse } from '../types';

interface PlaygroundState {
  // Session
  session: PlaygroundSessionResponse | null;
  setSession: React.Dispatch<React.SetStateAction<PlaygroundSessionResponse | null>>;

  // Chat messages
  chatMessages: DisplayMessage[];
  setChatMessages: React.Dispatch<React.SetStateAction<DisplayMessage[]>>;

  // Selected model (needed to know which model the session belongs to)
  selectedModel: string;
  setSelectedModel: React.Dispatch<React.SetStateAction<string>>;

  // Playground mode (top-level section)
  playgroundMode: 'automated' | 'manual' | 'history';
  setPlaygroundMode: React.Dispatch<React.SetStateAction<'automated' | 'manual' | 'history'>>;

  // Automated-testing sweep selection (Sampling vs Configurator). Lifted here
  // so the Testing menu can deep-link straight to the right sweep.
  sweepMode: 'sampling' | 'config';
  setSweepMode: React.Dispatch<React.SetStateAction<'sampling' | 'config'>>;

  // Active tab (within manual mode)
  activeTab: 'chat' | 'tools' | 'inspector';
  setActiveTab: React.Dispatch<React.SetStateAction<'chat' | 'tools' | 'inspector'>>;

  // System prompt
  systemPrompt: string;
  setSystemPrompt: React.Dispatch<React.SetStateAction<string>>;

  // TPS
  lastTPS: number | null;
  setLastTPS: React.Dispatch<React.SetStateAction<number | null>>;

  // Config state (locked after session creation)
  templateMode: 'builtin' | 'custom';
  setTemplateMode: React.Dispatch<React.SetStateAction<'builtin' | 'custom'>>;
  selectedTemplate: string;
  setSelectedTemplate: React.Dispatch<React.SetStateAction<string>>;
  customScript: string;
  setCustomScript: React.Dispatch<React.SetStateAction<string>>;
  contextWindow: number;
  setContextWindow: React.Dispatch<React.SetStateAction<number>>;
  nBatch: number;
  setNBatch: React.Dispatch<React.SetStateAction<number>>;
  nUBatch: number;
  setNUBatch: React.Dispatch<React.SetStateAction<number>>;
  nSeqMax: number;
  setNSeqMax: React.Dispatch<React.SetStateAction<number>>;
  flashAttention: string;
  setFlashAttention: React.Dispatch<React.SetStateAction<string>>;
  cacheType: string;
  setCacheType: React.Dispatch<React.SetStateAction<string>>;
  cacheMode: string;
  setCacheMode: React.Dispatch<React.SetStateAction<string>>;
  moeMode: string;
  setMoeMode: React.Dispatch<React.SetStateAction<string>>;
  moeKeepTopN: number;
  setMoeKeepTopN: React.Dispatch<React.SetStateAction<number>>;
  tensorBuftOverrides: string[];
  setTensorBuftOverrides: React.Dispatch<React.SetStateAction<string[]>>;

  // Tracks which model's catalog config has been applied to avoid
  // re-clobbering user edits on component remount.
  hydratedModelId: string;
  setHydratedModelId: React.Dispatch<React.SetStateAction<string>>;
}

const PlaygroundContext = createContext<PlaygroundState | null>(null);

export function PlaygroundProvider({ children }: { children: ReactNode }) {
  const [session, setSession] = useState<PlaygroundSessionResponse | null>(null);
  const [chatMessages, setChatMessages] = useState<DisplayMessage[]>([]);
  const [selectedModel, setSelectedModel] = useState('');
  const [playgroundMode, setPlaygroundMode] = useState<'automated' | 'manual' | 'history'>('automated');
  const [sweepMode, setSweepMode] = useState<'sampling' | 'config'>('sampling');
  const [activeTab, setActiveTab] = useState<'chat' | 'tools' | 'inspector'>('chat');
  const [systemPrompt, setSystemPrompt] = useState('You are a helpful assistant.');
  const [lastTPS, setLastTPS] = useState<number | null>(null);
  const [templateMode, setTemplateMode] = useState<'builtin' | 'custom'>('builtin');
  const [selectedTemplate, setSelectedTemplate] = useState('');
  const [customScript, setCustomScript] = useState('');
  const [contextWindow, setContextWindow] = useState(8192);
  const [nBatch, setNBatch] = useState(2048);
  const [nUBatch, setNUBatch] = useState(512);
  const [nSeqMax, setNSeqMax] = useState(1);
  const [flashAttention, setFlashAttention] = useState('auto');
  const [cacheType, setCacheType] = useState('');
  const [cacheMode, setCacheMode] = useState('none');
  const [moeMode, setMoeMode] = useState('');
  const [moeKeepTopN, setMoeKeepTopN] = useState(0);
  const [tensorBuftOverrides, setTensorBuftOverrides] = useState<string[]>([]);
  const [hydratedModelId, setHydratedModelId] = useState('');

  const value = useMemo<PlaygroundState>(() => ({
    session, setSession,
    chatMessages, setChatMessages,
    selectedModel, setSelectedModel,
    playgroundMode, setPlaygroundMode,
    sweepMode, setSweepMode,
    activeTab, setActiveTab,
    systemPrompt, setSystemPrompt,
    lastTPS, setLastTPS,
    templateMode, setTemplateMode,
    selectedTemplate, setSelectedTemplate,
    customScript, setCustomScript,
    contextWindow, setContextWindow,
    nBatch, setNBatch,
    nUBatch, setNUBatch,
    nSeqMax, setNSeqMax,
    flashAttention, setFlashAttention,
    cacheType, setCacheType,
    cacheMode, setCacheMode,
    moeMode, setMoeMode,
    moeKeepTopN, setMoeKeepTopN,
    tensorBuftOverrides, setTensorBuftOverrides,
    hydratedModelId, setHydratedModelId,
  }), [
    session, chatMessages, selectedModel, playgroundMode, sweepMode, activeTab, systemPrompt, lastTPS,
    templateMode, selectedTemplate, customScript, contextWindow, nBatch, nUBatch,
    nSeqMax, flashAttention, cacheType, cacheMode, moeMode, moeKeepTopN, tensorBuftOverrides, hydratedModelId,
  ]);

  return (
    <PlaygroundContext.Provider value={value}>
      {children}
    </PlaygroundContext.Provider>
  );
}

export function usePlayground() {
  const context = useContext(PlaygroundContext);
  if (!context) {
    throw new Error('usePlayground must be used within a PlaygroundProvider');
  }
  return context;
}
