fix: highlighting

lua/cp/commands/init.lua

@@ -52,6 +52,20 @@ local function parse_command(args)
       else
         return { type = 'action', action = 'interact' }
       end
+    elseif first == 'run' then
+      local test_arg = args[2]
+      if test_arg then
+        local test_index = tonumber(test_arg)
+        if not test_index then
+          return { type = 'error', message = 'Test index must be a number' }
+        end
+        if test_index < 1 or test_index ~= math.floor(test_index) then
+          return { type = 'error', message = 'Test index must be >= 1' }
+        end
+        return { type = 'action', action = 'run', test_index = test_index }
+      else
+        return { type = 'action', action = 'run' }
+      end
     else
       return { type = 'action', action = first }
     end
@@ -109,7 +123,7 @@ function M.handle_command(opts)
     if cmd.action == 'interact' then
       ui.toggle_interactive(cmd.interactor_cmd)
     elseif cmd.action == 'run' then
-      ui.run_io_view()
+      ui.run_io_view(cmd.test_index)
     elseif cmd.action == 'panel' then
       ui.toggle_panel()
     elseif cmd.action == 'debug' then

lua/cp/config.lua

@@ -18,7 +18,6 @@
 ---@field overrides? table<string, CpPlatformOverrides>
 
 ---@class PanelConfig
----@field ansi boolean
 ---@field diff_mode "none"|"vim"|"git"
 ---@field max_output_lines integer
 
@@ -36,6 +35,7 @@
 ---@field setup_io_output? fun(bufnr: integer, state: cp.State)
 
 ---@class CpUI
+---@field ansi boolean
 ---@field panel PanelConfig
 ---@field diff DiffConfig
 ---@field picker string|nil
@@ -115,7 +115,8 @@ M.defaults = {
   scrapers = constants.PLATFORMS,
   filename = nil,
   ui = {
-    panel = { ansi = true, diff_mode = 'none', max_output_lines = 50 },
+    ansi = true,
+    panel = { diff_mode = 'none', max_output_lines = 50 },
     diff = {
       git = {
         args = { 'diff', '--no-index', '--word-diff=plain', '--word-diff-regex=.', '--no-prefix' },
@@ -243,7 +244,7 @@ function M.setup(user_config)
   })
 
   vim.validate({
-    ansi = { cfg.ui.panel.ansi, 'boolean' },
+    ansi = { cfg.ui.ansi, 'boolean' },
     diff_mode = {
       cfg.ui.panel.diff_mode,
       function(v)

lua/cp/runner/run.lua

@@ -122,7 +122,7 @@ local function run_single_test_case(test_case)
   local out = r.stdout or ''
   local highlights = {}
   if out ~= '' then
-    if config.ui.panel.ansi then
+    if config.ui.ansi then
       local parsed = ansi.parse_ansi_text(out)
       out = table.concat(parsed.lines, '\n')
       highlights = parsed.highlights
@@ -224,14 +224,22 @@ function M.run_test_case(index)
   return true
 end
 
+---@param indices? integer[]
 ---@return RanTestCase[]
-function M.run_all_test_cases()
+function M.run_all_test_cases(indices)
-  local results = {}
+  local to_run = indices
-  for i = 1, #panel_state.test_cases do
+  if not to_run then
-    M.run_test_case(i)
+    to_run = {}
-    results[i] = panel_state.test_cases[i]
+    for i = 1, #panel_state.test_cases do
+      to_run[i] = i
+    end
   end
-  return results
+
+  for _, i in ipairs(to_run) do
+    M.run_test_case(i)
+  end
+
+  return panel_state.test_cases
 end
 
 ---@return PanelState
@@ -247,7 +255,7 @@ function M.handle_compilation_failure(output)
   local txt
   local hl = {}
 
-  if config.ui.panel.ansi then
+  if config.ui.ansi then
     local p = ansi.parse_ansi_text(output or '')
     txt = table.concat(p.lines, '\n')
     hl = p.highlights

lua/cp/setup.lua

@@ -60,6 +60,18 @@ local function start_tests(platform, contest_id, problems)
         ev.memory_mb or 0,
         ev.interactive
       )
+
+      local io_state = state.get_io_view_state()
+      if io_state then
+        local test_cases = cache.get_test_cases(platform, contest_id, state.get_problem_id())
+        local input_lines = {}
+        for _, tc in ipairs(test_cases) do
+          for _, line in ipairs(vim.split(tc.input, '\n', { plain = true, trimempty = false })) do
+            table.insert(input_lines, line)
+          end
+        end
+        require('cp.utils').update_buffer_content(io_state.input_buf, input_lines, nil, nil)
+      end
     end)
   end
 end

lua/cp/ui/ansi.lua

@@ -321,6 +321,25 @@ function M.setup_highlight_groups()
   vim.api.nvim_set_hl(0, 'CpAnsiBold', { bold = true })
   vim.api.nvim_set_hl(0, 'CpAnsiItalic', { italic = true })
   vim.api.nvim_set_hl(0, 'CpAnsiBoldItalic', { bold = true, italic = true })
+
+  for _, combo in ipairs(combinations) do
+    for color_name, _ in pairs(color_map) do
+      local parts = { 'CpAnsi' }
+      if combo.bold then
+        table.insert(parts, 'Bold')
+      end
+      if combo.italic then
+        table.insert(parts, 'Italic')
+      end
+      table.insert(parts, color_name)
+      local hl_name = table.concat(parts)
+      dyn_hl_cache[hl_name] = true
+    end
+  end
+
+  dyn_hl_cache['CpAnsiBold'] = true
+  dyn_hl_cache['CpAnsiItalic'] = true
+  dyn_hl_cache['CpAnsiBoldItalic'] = true
 end
 
 ---@param text string

lua/cp/ui/panel.lua

@@ -15,6 +15,22 @@ local utils = require('cp.utils')
 local current_diff_layout = nil
 local current_mode = nil
 
+local function populate_input()
+  local io_state = state.get_io_view_state()
+
+  local test_cases =
+    cache.get_test_cases(state.get_platform(), state.get_contest_id(), state.get_problem_id())
+
+  local input_lines = {}
+  for _, tc in ipairs(test_cases) do
+    for _, line in ipairs(vim.split(tc.input, '\n', { plain = true, trimempty = false })) do
+      table.insert(input_lines, line)
+    end
+  end
+
+  utils.update_buffer_content(io_state.input_buf, input_lines)
+end
+
 function M.disable()
   local active_panel = state.get_active_panel()
   if not active_panel then
@@ -269,22 +285,15 @@ function M.ensure_io_view()
     end
   end
 
-  local test_cases = cache.get_test_cases(platform, contest_id, problem_id)
+  utils.update_buffer_content(input_buf, {}, nil, nil)
-  local input_lines = {}
-  for i, tc in ipairs(test_cases) do
-    table.insert(input_lines, string.format('--- Test %d ---', i))
-    for _, line in ipairs(vim.split(tc.input, '\n', { plain = true, trimempty = false })) do
-      table.insert(input_lines, line)
-    end
-  end
-
-  utils.update_buffer_content(input_buf, input_lines, nil, nil)
   utils.update_buffer_content(output_buf, {}, nil, nil)
 
   vim.api.nvim_set_current_win(solution_win)
+
+  populate_input()
 end
 
-function M.run_io_view()
+function M.run_io_view(test_index)
   local platform, contest_id = state.get_platform(), state.get_contest_id()
   if not platform then
     logger.log(
@@ -296,8 +305,7 @@ function M.run_io_view()
 
   if not contest_id then
     logger.log(
-      ("No contest '%s' configured for platform '%s'."):format(
+      ("No contest configured for platform '%s'."):format(
-        contest_id,
         constants.PLATFORM_DISPLAY_NAMES[platform]
       ),
       vim.log.levels.ERROR
@@ -326,18 +334,31 @@ function M.run_io_view()
   M.ensure_io_view()
 
   local run = require('cp.runner.run')
-  local run_render = require('cp.runner.run_render')
   if not run.load_test_cases() then
-    logger.log('no test cases found', vim.log.levels.WARN)
+    logger.log('No test cases available', vim.log.levels.ERROR)
     return
   end
 
-  local execute = require('cp.runner.execute')
+  local test_state = run.get_panel_state()
-  local compile_result = execute.compile_problem()
+  local test_indices = {}
-  if compile_result.success then
+
-    run.run_all_test_cases()
+  if test_index then
+    if test_index < 1 or test_index > #test_state.test_cases then
+      logger.log(
+        string.format(
+          'Test %d does not exist (only %d tests available)',
+          test_index,
+          #test_state.test_cases
+        ),
+        vim.log.levels.ERROR
+      )
+      return
+    end
+    test_indices = { test_index }
   else
-    run.handle_compilation_failure(compile_result.output)
+    for i = 1, #test_state.test_cases do
+      test_indices[i] = i
+    end
   end
 
   local io_state = state.get_io_view_state()
@@ -345,30 +366,91 @@ function M.run_io_view()
     return
   end
 
-  local test_state = run.get_panel_state()
+  local config = config_module.get_config()
-  run_render.setup_highlights()
 
-  local verdict_lines = {}
+  if config.ui.ansi then
-  local verdict_highlights = {}
+    require('cp.ui.ansi').setup_highlight_groups()
-  for i, tc in ipairs(test_state.test_cases) do
-    local status = run_render.get_status_info(tc)
-    local time = tc.time_ms and string.format('%.2f', tc.time_ms) or '—'
-    local mem = tc.rss_mb and string.format('%.0f', tc.rss_mb) or '—'
-    local line = string.format('Test %d: %s (%sms, %sMB)', i, status.text, time, mem)
-    table.insert(verdict_lines, line)
-    local status_pos = line:find(status.text, 1, true)
-    if status_pos then
-      table.insert(verdict_highlights, {
-        line = i - 1,
-        col_start = status_pos - 1,
-        col_end = status_pos - 1 + #status.text,
-        highlight_group = status.highlight_group,
-      })
-    end
   end
 
-  local verdict_ns = vim.api.nvim_create_namespace('cp_io_view_verdict')
+  local execute = require('cp.runner.execute')
-  utils.update_buffer_content(io_state.output_buf, verdict_lines, verdict_highlights, verdict_ns)
+  local compile_result = execute.compile_problem()
+  if not compile_result.success then
+    local ansi = require('cp.ui.ansi')
+    local output = compile_result.output or ''
+    local lines, highlights
+
+    if config.ui.ansi then
+      local parsed = ansi.parse_ansi_text(output)
+      lines = parsed.lines
+      highlights = parsed.highlights
+    else
+      lines = vim.split(output:gsub('\027%[[%d;]*[a-zA-Z]', ''), '\n')
+      highlights = {}
+    end
+
+    local ns = vim.api.nvim_create_namespace('cp_io_view_compile_error')
+    utils.update_buffer_content(io_state.output_buf, lines, highlights, ns)
+    return
+  end
+
+  run.run_all_test_cases(test_indices)
+
+  local input_lines = {}
+  for _, idx in ipairs(test_indices) do
+    local tc = test_state.test_cases[idx]
+    for _, line in ipairs(vim.split(tc.input, '\n', { plain = true, trimempty = false })) do
+      table.insert(input_lines, line)
+    end
+  end
+  utils.update_buffer_content(io_state.input_buf, input_lines, nil, nil)
+
+  local run_render = require('cp.runner.run_render')
+  run_render.setup_highlights()
+
+  if #test_indices == 1 then
+    local idx = test_indices[1]
+    local tc = test_state.test_cases[idx]
+    local status = run_render.get_status_info(tc)
+
+    local output_lines = {}
+    if tc.actual then
+      for _, line in ipairs(vim.split(tc.actual, '\n', { plain = true, trimempty = false })) do
+        table.insert(output_lines, line)
+      end
+    end
+
+    table.insert(output_lines, '')
+    local time = tc.time_ms and string.format('%.2fms', tc.time_ms) or '—'
+    local code = tc.code and tostring(tc.code) or '—'
+    table.insert(output_lines, string.format('--- %s: %s | Exit: %s ---', status.text, time, code))
+
+    local highlights = tc.actual_highlights or {}
+    local ns = vim.api.nvim_create_namespace('cp_io_view_output')
+    utils.update_buffer_content(io_state.output_buf, output_lines, highlights, ns)
+  else
+    local verdict_lines = {}
+    local verdict_highlights = {}
+    for _, idx in ipairs(test_indices) do
+      local tc = test_state.test_cases[idx]
+      local status = run_render.get_status_info(tc)
+      local time = tc.time_ms and string.format('%.2f', tc.time_ms) or '—'
+      local mem = tc.rss_mb and string.format('%.0f', tc.rss_mb) or '—'
+      local line = string.format('Test %d: %s (%sms, %sMB)', idx, status.text, time, mem)
+      table.insert(verdict_lines, line)
+      local status_pos = line:find(status.text, 1, true)
+      if status_pos then
+        table.insert(verdict_highlights, {
+          line = #verdict_lines - 1,
+          col_start = status_pos - 1,
+          col_end = status_pos - 1 + #status.text,
+          highlight_group = status.highlight_group,
+        })
+      end
+    end
+
+    local verdict_ns = vim.api.nvim_create_namespace('cp_io_view_verdict')
+    utils.update_buffer_content(io_state.output_buf, verdict_lines, verdict_highlights, verdict_ns)
+  end
 end
 
 ---@param panel_opts? PanelOpts
@@ -386,6 +468,8 @@ function M.toggle_panel(panel_opts)
       state.set_saved_session(nil)
     end
     state.set_active_panel(nil)
+    M.ensure_io_view()
+    populate_input()
     return
   end
 
@@ -436,6 +520,17 @@ function M.toggle_panel(panel_opts)
     return
   end
 
+  local io_state = state.get_io_view_state()
+  if io_state then
+    if vim.api.nvim_win_is_valid(io_state.output_win) then
+      vim.api.nvim_win_close(io_state.output_win, true)
+    end
+    if vim.api.nvim_win_is_valid(io_state.input_win) then
+      vim.api.nvim_win_close(io_state.input_win, true)
+    end
+    state.set_io_view_state(nil)
+  end
+
   local session_file = vim.fn.tempname()
   state.set_saved_session(session_file)
   vim.cmd(('mksession! %s'):format(session_file))
@@ -537,7 +632,7 @@ function M.toggle_panel(panel_opts)
   refresh_panel()
 
   vim.schedule(function()
-    if config.ui.panel.ansi then
+    if config.ui.ansi then
       require('cp.ui.ansi').setup_highlight_groups()
     end
     if current_diff_layout then