import re

with open('/Users/amritha/Desktop/work/intern/migratron-ui/src/components/Dashboard.tsx', 'r') as f:
    content = f.read()

# 1. Update Mode
content = content.replace(
    "const [mode, setMode] = useState<'eval' | 'inference'>('eval');",
    "const [mode, setMode] = useState<'demo' | 'eval' | 'inference'>('demo');\n  const [terminalLogs, setTerminalLogs] = useState<{type: string, message?: string, step?: number, tool_name?: string, tool_args?: any}[]>([]);\n  const [agentReasoning, setAgentReasoning] = useState<{step: number, tool_name: string, tool_args: any, result?: string}[]>([]);\n  const [generatedPatch, setGeneratedPatch] = useState<string | null>(null);\n  const [reward, setReward] = useState<number>(0);\n  const [ws, setWs] = useState<WebSocket | null>(null);"
)

# 2. Update mode derived variables
content = content.replace(
    "const repoName = mode === 'eval' ? selectedTask?.repo_name : (isInferenceLoaded ? inferenceRepoUrl.replace('https://github.com/', '').replace('http://github.com/', '') : null);",
    "const repoName = (mode === 'eval' || mode === 'demo') ? selectedTask?.repo_name : (isInferenceLoaded ? inferenceRepoUrl.replace('https://github.com/', '').replace('http://github.com/', '') : null);"
)
content = content.replace(
    "const commitHash = mode === 'eval' ? selectedTask?.commit_hash : (isInferenceLoaded ? inferenceCommitHash : null);",
    "const commitHash = (mode === 'eval' || mode === 'demo') ? selectedTask?.commit_hash : (isInferenceLoaded ? inferenceCommitHash : null);"
)

# 3. Update useEffect for timer
timer_old = """  useEffect(() => {
    let timer: ReturnType<typeof setInterval>;
    if (simulationState === 'running') {
      timer = setInterval(() => {
        setSimulationStep(prev => {
          if (prev >= 6) {
            setSimulationState('completed');
            clearInterval(timer);
            return prev;
          }
          return prev + 1;
        });
      }, 1500);
    }
    return () => clearInterval(timer);
  }, [simulationState]);"""

timer_new = """  useEffect(() => {
    let timer: ReturnType<typeof setInterval>;
    if (simulationState === 'running' && mode === 'demo') {
      timer = setInterval(() => {
        setSimulationStep(prev => {
          if (prev >= 6) {
            setSimulationState('completed');
            clearInterval(timer);
            return prev;
          }
          return prev + 1;
        });
      }, 1500);
    }
    return () => clearInterval(timer);
  }, [simulationState, mode]);
  
  // Cleanup WS on unmount
  useEffect(() => {
    return () => {
      if (ws) {
        ws.close();
      }
    };
  }, [ws]);"""
content = content.replace(timer_old, timer_new)

# 4. Update handleStartEpisode
start_old = """  const handleStartEpisode = () => {
    if (selectedFile && repoName) {
      setSimulationState('running');
      setSimulationStep(0);
      setActiveTab('code');
    }
  };"""
  
start_new = """  const handleStartEpisode = () => {
    if (selectedFile && repoName) {
      setSimulationState('running');
      setSimulationStep(0);
      setActiveTab('code');
      
      if (mode !== 'demo') {
        setTerminalLogs([]);
        setAgentReasoning([]);
        setGeneratedPatch(null);
        setReward(0);
        
        const socket = new WebSocket(`ws://${window.location.host}/ws/simulate`);
        setWs(socket);
        
        socket.onopen = () => {
          socket.send(JSON.stringify({
            repo: repoName,
            commit: commitHash,
            mode: mode,
            task_index: selectedTaskIndex !== null ? selectedTaskIndex : 0
          }));
        };
        
        socket.onmessage = (event) => {
          const data = JSON.parse(event.data);
          
          if (data.type === 'info') {
             setTerminalLogs(prev => [...prev, data]);
          } else if (data.type === 'tool_call') {
             setTerminalLogs(prev => [...prev, data]);
             setSimulationStep(data.step);
             setAgentReasoning(prev => [...prev, {
                 step: data.step,
                 tool_name: data.tool_name,
                 tool_args: data.tool_args
             }]);
          } else if (data.type === 'tool_result') {
             setTerminalLogs(prev => [...prev, data]);
             setAgentReasoning(prev => {
                 const updated = [...prev];
                 const idx = updated.findIndex(r => r.step === data.step);
                 if (idx >= 0) {
                     updated[idx].result = data.result;
                 }
                 return updated;
             });
          } else if (data.type === 'completed') {
             setSimulationState('completed');
             setGeneratedPatch(data.patch);
             setReward(data.reward);
             setTerminalLogs(prev => [...prev, {type: 'info', message: data.success ? '[SUCCESS] Task completed successfully.' : '[FAILED] Task failed.'}]);
             socket.close();
          } else if (data.type === 'error') {
             setTerminalLogs(prev => [...prev, {type: 'info', message: `[ERROR] ${data.message}`}]);
             setSimulationState('completed');
             socket.close();
          }
        };
        
        socket.onerror = (error) => {
          console.error("WebSocket Error:", error);
          setTerminalLogs(prev => [...prev, {type: 'info', message: '[ERROR] WebSocket connection failed.'}]);
          setSimulationState('completed');
        };
      }
    }
  };"""
content = content.replace(start_old, start_new)

# 5. Nav Buttons
nav_old = """            <button
              onClick={() => { setMode('eval'); resetEvalSelection(); resetInference(); }}
              className={`px-3 py-1 rounded cursor-pointer transition-colors ${mode === 'eval' ? 'bg-[#007fd4] text-white' : 'hover:bg-white/10'}`}
            >
              Evaluation Mode
            </button>
            <button
              onClick={() => { setMode('inference'); resetEvalSelection(); resetInference(); }}
              className={`px-3 py-1 rounded cursor-pointer transition-colors ${mode === 'inference' ? 'bg-[#007fd4] text-white' : 'hover:bg-white/10'}`}
            >
              Inference Mode
            </button>"""

nav_new = """            <button
              onClick={() => { setMode('demo'); resetEvalSelection(); resetInference(); }}
              className={`px-3 py-1 rounded cursor-pointer transition-colors ${mode === 'demo' ? 'bg-[#007fd4] text-white' : 'hover:bg-white/10'}`}
            >
              Demo Mode
            </button>
            <button
              onClick={() => { setMode('eval'); resetEvalSelection(); resetInference(); }}
              className={`px-3 py-1 rounded cursor-pointer transition-colors ${mode === 'eval' ? 'bg-[#007fd4] text-white' : 'hover:bg-white/10'}`}
            >
              Evaluation Mode
            </button>
            <button
              onClick={() => { setMode('inference'); resetEvalSelection(); resetInference(); }}
              className={`px-3 py-1 rounded cursor-pointer transition-colors ${mode === 'inference' ? 'bg-[#007fd4] text-white' : 'hover:bg-white/10'}`}
            >
              Inference Mode
            </button>"""
content = content.replace(nav_old, nav_new)

# 6. Sidebar Headers
content = content.replace("repoTree ? 'Repository Files' : (mode === 'eval' ? 'Evaluation Tasks' : 'Inference Setup')", "repoTree ? 'Repository Files' : ((mode === 'eval' || mode === 'demo') ? 'Evaluation Tasks' : 'Inference Setup')")
content = content.replace("!repoTree && mode === 'eval' && (", "!repoTree && (mode === 'eval' || mode === 'demo') && (")
content = content.replace("mode === 'eval' ? resetEvalSelection : resetInference", "(mode === 'eval' || mode === 'demo') ? resetEvalSelection : resetInference")
content = content.replace("activeTab === 'gold' && mode === 'eval'", "activeTab === 'gold' && (mode === 'eval' || mode === 'demo')")
content = content.replace("mode === 'eval' && selectedTask?.gold_patch &&", "(mode === 'eval' || mode === 'demo') && selectedTask?.gold_patch &&")
content = content.replace("mode === 'eval' && selectedTask && (", "(mode === 'eval' || mode === 'demo') && selectedTask && (")

# 7. Render Patch
patch_old = """                  {activeTab === 'patch' && (
                    <div className="text-[#d4d4d4]">
                      {mockGeneratedPatch.split('\\n').map((line, i) => {"""
patch_new = """                  {activeTab === 'patch' && (
                    <div className="text-[#d4d4d4]">
                      {(mode === 'demo' ? mockGeneratedPatch : (generatedPatch || '')).split('\\n').map((line, i) => {"""
content = content.replace(patch_old, patch_new)

# 8. Render Terminal
terminal_old = """                <>
                  <div><span className="text-[#23d18b]">$</span> Starting Migratron env for {repoName}...</div>
                  {simulationStep >= 1 && <div className="text-[#3b8eea]">[INFO] Target file selected: {selectedFile?.path}</div>}
                  {simulationStep >= 2 && <div className="text-[#3b8eea]">[INFO] Initializing agent context...</div>}
                  {simulationStep >= 3 && <div className="text-[#3b8eea]">[INFO] Executing agent reasoning loop...</div>}
                  {simulationStep >= 4 && <div>Agent generated patch for: <span className="text-[#f5f543]">{selectedFile?.path}</span></div>}
                  {simulationStep >= 5 && <div className="text-[#3b8eea]">[INFO] Applying patch and verifying syntax...</div>}
                  {simulationStep >= 6 && <div className="text-[#23d18b]">[SUCCESS] Task completed. Evaluating GRPO Reward...</div>}
                  <div className="animate-pulse mt-2">_</div>
                </>"""
terminal_new = """                <>
                  {mode === 'demo' ? (
                    <>
                      <div><span className="text-[#23d18b]">$</span> Starting Migratron env for {repoName}...</div>
                      {simulationStep >= 1 && <div className="text-[#3b8eea]">[INFO] Target file selected: {selectedFile?.path}</div>}
                      {simulationStep >= 2 && <div className="text-[#3b8eea]">[INFO] Initializing agent context...</div>}
                      {simulationStep >= 3 && <div className="text-[#3b8eea]">[INFO] Executing agent reasoning loop...</div>}
                      {simulationStep >= 4 && <div>Agent generated patch for: <span className="text-[#f5f543]">{selectedFile?.path}</span></div>}
                      {simulationStep >= 5 && <div className="text-[#3b8eea]">[INFO] Applying patch and verifying syntax...</div>}
                      {simulationStep >= 6 && <div className="text-[#23d18b]">[SUCCESS] Task completed. Evaluating GRPO Reward...</div>}
                    </>
                  ) : (
                    terminalLogs.map((log, idx) => (
                      <div key={idx} className={log.message?.includes('[ERROR]') ? 'text-[#f14c4c]' : (log.message?.includes('[SUCCESS]') ? 'text-[#23d18b]' : (log.message?.includes('[INFO]') ? 'text-[#3b8eea]' : ''))}>
                        {log.message || `Agent called ${log.tool_name}`}
                      </div>
                    ))
                  )}
                  {simulationState === 'running' && <div className="animate-pulse mt-2">_</div>}
                </>"""
content = content.replace(terminal_old, terminal_new)

# 9. Render Agent Reasoning
agent_old = """                  {simulationStep >= 2 && (
                    <div className="border border-[#3c3c3c] rounded p-3 bg-[#252526] animate-fade-in">
                      <div className="flex justify-between items-center mb-2">
                        <span className="text-xs font-sans text-[#cccccc] uppercase">Thought Process</span>
                      </div>
                      <p className="text-[13px] text-[#d4d4d4] italic">
                        "The objective is to migrate {repoName}. I need to focus on {selectedFile.path}. I will first read the contents of this file to understand the current implementation."
                      </p>
                    </div>
                  )}

                  {simulationStep >= 3 && (
                    <div className="border border-[#3c3c3c] rounded p-3 bg-[#252526] border-l-2 border-l-[#007fd4] animate-fade-in">
                      <div className="flex justify-between items-center mb-2">
                        <span className="text-xs font-sans text-[#569cd6] uppercase flex items-center gap-1">
                          <Code size={12} /> Tool Execution
                        </span>
                      </div>
                      <div className="text-[12px] text-[#ce9178] bg-[#1e1e1e] p-2 rounded">
                        <div className="text-[#9cdcfe]">call:</div> <span className="text-[#dcdcaa]">view_file</span> {'{'} <br />
                        &nbsp;&nbsp;<span className="text-[#9cdcfe]">AbsolutePath</span>: <span className="text-[#ce9178]">"/work/{selectedFile.path}"</span><br />
                        {'}'}
                      </div>
                    </div>
                  )}

                  {simulationStep >= 4 && (
                    <div className="border border-[#3c3c3c] rounded p-3 bg-[#252526] border-l-2 border-l-[#23d18b] animate-fade-in">
                      <div className="flex justify-between items-center mb-2">
                        <span className="text-xs font-sans text-[#23d18b] uppercase flex items-center gap-1">
                          <CheckCircle2 size={12} /> Tool Response
                        </span>
                      </div>
                      <p className="text-[12px] text-[#d4d4d4]">
                        File contents read. Identifying migration targets and generating patch block...
                      </p>
                    </div>
                  )}"""

agent_new = """                  {mode === 'demo' ? (
                    <>
                      {simulationStep >= 2 && (
                        <div className="border border-[#3c3c3c] rounded p-3 bg-[#252526] animate-fade-in">
                          <div className="flex justify-between items-center mb-2">
                            <span className="text-xs font-sans text-[#cccccc] uppercase">Thought Process</span>
                          </div>
                          <p className="text-[13px] text-[#d4d4d4] italic">
                            "The objective is to migrate {repoName}. I need to focus on {selectedFile.path}. I will first read the contents of this file to understand the current implementation."
                          </p>
                        </div>
                      )}

                      {simulationStep >= 3 && (
                        <div className="border border-[#3c3c3c] rounded p-3 bg-[#252526] border-l-2 border-l-[#007fd4] animate-fade-in">
                          <div className="flex justify-between items-center mb-2">
                            <span className="text-xs font-sans text-[#569cd6] uppercase flex items-center gap-1">
                              <Code size={12} /> Tool Execution
                            </span>
                          </div>
                          <div className="text-[12px] text-[#ce9178] bg-[#1e1e1e] p-2 rounded">
                            <div className="text-[#9cdcfe]">call:</div> <span className="text-[#dcdcaa]">view_file</span> {'{'} <br />
                            &nbsp;&nbsp;<span className="text-[#9cdcfe]">AbsolutePath</span>: <span className="text-[#ce9178]">"/work/{selectedFile.path}"</span><br />
                            {'}'}
                          </div>
                        </div>
                      )}

                      {simulationStep >= 4 && (
                        <div className="border border-[#3c3c3c] rounded p-3 bg-[#252526] border-l-2 border-l-[#23d18b] animate-fade-in">
                          <div className="flex justify-between items-center mb-2">
                            <span className="text-xs font-sans text-[#23d18b] uppercase flex items-center gap-1">
                              <CheckCircle2 size={12} /> Tool Response
                            </span>
                          </div>
                          <p className="text-[12px] text-[#d4d4d4]">
                            File contents read. Identifying migration targets and generating patch block...
                          </p>
                        </div>
                      )}
                    </>
                  ) : (
                    agentReasoning.map((r, idx) => (
                      <div key={idx} className="space-y-3 mb-3">
                        <div className="border border-[#3c3c3c] rounded p-3 bg-[#252526] border-l-2 border-l-[#007fd4] animate-fade-in">
                          <div className="flex justify-between items-center mb-2">
                            <span className="text-xs font-sans text-[#569cd6] uppercase flex items-center gap-1">
                              <Code size={12} /> Tool Execution (Step {r.step})
                            </span>
                          </div>
                          <div className="text-[12px] text-[#ce9178] bg-[#1e1e1e] p-2 rounded whitespace-pre-wrap">
                            <div className="text-[#9cdcfe]">call:</div> <span className="text-[#dcdcaa]">{r.tool_name}</span> {'{'}<br />
                            {Object.entries(r.tool_args || {}).map(([k, v]) => (
                                <React.Fragment key={k}>
                                  &nbsp;&nbsp;<span className="text-[#9cdcfe]">{k}</span>: <span className="text-[#ce9178]">{JSON.stringify(v)}</span><br />
                                </React.Fragment>
                            ))}
                            {'}'}
                          </div>
                        </div>
                        {r.result && (
                          <div className="border border-[#3c3c3c] rounded p-3 bg-[#252526] border-l-2 border-l-[#23d18b] animate-fade-in">
                            <div className="flex justify-between items-center mb-2">
                              <span className="text-xs font-sans text-[#23d18b] uppercase flex items-center gap-1">
                                <CheckCircle2 size={12} /> Tool Response
                              </span>
                            </div>
                            <div className="text-[12px] text-[#d4d4d4] max-h-32 overflow-y-auto whitespace-pre-wrap font-mono">
                              {r.result}
                            </div>
                          </div>
                        )}
                      </div>
                    ))
                  )}"""
content = content.replace(agent_old, agent_new)

# 10. Reward
reward_old = "{simulationState === 'completed' ? '+2.80' : '0.00'}"
reward_new = "{mode === 'demo' ? (simulationState === 'completed' ? '+2.80' : '0.00') : (reward > 0 ? `+${reward.toFixed(2)}` : reward.toFixed(2))}"
content = content.replace(reward_old, reward_new)

with open('/Users/amritha/Desktop/work/intern/migratron-ui/src/components/Dashboard.tsx', 'w') as f:
    f.write(content)

print("Patcher finished successfully.")