From bd30fb626c88fa0c31ea60f5775e1377885ddce9 Mon Sep 17 00:00:00 2001
From: Barrett Ruth <barrett.ruth@ramp.com>
Date: Fri, 24 Oct 2025 01:11:19 -0400
Subject: [PATCH 1/5] feat: start lang refactor

---
 doc/cp.nvim.txt              | 109 +++++++++++++++++++++++++++++------
 lua/cp/commands/init.lua     |  41 +++++++++----
 lua/cp/commands/picker.lua   |   5 +-
 lua/cp/config.lua            |  44 ++++++++++++++
 lua/cp/pickers/fzf_lua.lua   |  15 +++--
 lua/cp/pickers/telescope.lua |  18 +++---
 lua/cp/setup.lua             |  42 +++++++++++++-
 7 files changed, 229 insertions(+), 45 deletions(-)

diff --git a/doc/cp.nvim.txt b/doc/cp.nvim.txt
index 27cf18e..c8b17ce 100644
--- a/doc/cp.nvim.txt
+++ b/doc/cp.nvim.txt
@@ -25,18 +25,13 @@ COMMANDS                                                         *cp-commands*
         cp.nvim uses a single :CP command with intelligent argument parsing:
 
         Setup Commands ~
-            :CP {platform} {contest_id}
+            :CP {platform} {contest_id} [--lang {language}]
                                 Full setup: set platform and load contest metadata.
                                 Scrapes test cases and creates source file.
-                                Example: >
+                                --lang: Use specific language (default: platform default)
+                                Examples: >
                                     :CP codeforces 1933
-<
-            :CP {platform} {contest_id}
-                                Contest setup: set platform, load contest metadata,
-                                and scrape all test cases in the contest. 
-                                Opens the first problem after completion.
-                                Example: >
-                                    :CP atcoder abc324
+                                    :CP codeforces 1933 --lang python
 <
         View Commands ~
             :CP run [--debug] [n]
@@ -59,8 +54,14 @@ COMMANDS                                                         *cp-commands*
                                     :CP panel --debug 3  " Test 3, debug build
 <
 
-            :CP pick            Launch configured picker for interactive
+            :CP pick [--lang {language}]
+                                Launch configured picker for interactive
                                 platform/contest selection.
+                                --lang: Pre-select language for chosen contest.
+                                Example: >
+                                    :CP pick
+                                    :CP pick --lang python
+<
 
             :CP interact [script]
                                 Open an interactive terminal for the current problem.
@@ -70,15 +71,36 @@ COMMANDS                                                         *cp-commands*
                                 file. Only valid for interactive problems.
 
         Navigation Commands ~
-            :CP next            Navigate to next problem in current contest.
+            :CP next [--lang {language}]
+                                Navigate to next problem in current contest.
                                 Stops at last problem (no wrapping).
-
-
-            :CP prev            Navigate to previous problem in current contest.
+                                --lang: Use specific language for next problem.
+                                By default, preserves current file's language if
+                                enabled for the new problem, otherwise uses platform
+                                default.
+                                Examples: >
+                                    :CP next
+                                    :CP next --lang python
+<
+            :CP prev [--lang {language}]
+                                Navigate to previous problem in current contest.
                                 Stops at first problem (no wrapping).
-
-            :CP {problem_id}    Jump to problem {problem_id} in a contest.
+                                --lang: Use specific language for previous problem.
+                                By default, preserves current file's language if
+                                enabled for the new problem, otherwise uses platform
+                                default.
+                                Examples: >
+                                    :CP prev
+                                    :CP prev --lang cpp
+<
+            :CP {problem_id} [--lang {language}]
+                                Jump to problem {problem_id} in a contest.
                                 Requires that a contest has already been set up.
+                                --lang: Use specific language for this problem.
+                                Examples: >
+                                    :CP B
+                                    :CP C --lang python
+<
 
         State Restoration ~
             :CP                 Restore state from current file.
@@ -357,6 +379,49 @@ run CSES problems with Rust using the single schema:
         }
 <
 
+==============================================================================
+LANGUAGE SELECTION                                           *cp-lang-selection*
+
+cp.nvim supports multiple languages per problem. Each platform enables specific
+languages and has a default. You can override the language for any setup or
+navigation command using the --lang flag.
+
+Language Selection Behavior ~
+
+When setting up or navigating to a problem:
+
+1. Explicit --lang flag takes highest priority
+2. If no --lang flag, tries to preserve current file's language
+   (only if that language is enabled for the new problem)
+3. Falls back to platform's default language
+
+Multiple Solution Files ~
+
+Different languages create different solution files. For example:
+    1848a.cc        (C++ solution)
+    1848a.py        (Python solution)
+
+Both files can exist simultaneously with their own state. Switching between
+languages means switching between different files.
+
+Examples ~
+>
+    :CP codeforces 1848             " Use platform default (likely C++)
+    :CP codeforces 1848 --lang python   " Use Python explicitly
+
+    " In 1848a.cc (C++ file):
+    :CP next                        " Next problem tries to use C++
+    :CP next --lang python          " Next problem uses Python
+
+    " In 1848a.py (Python file):
+    :CP next                        " Next problem tries to use Python
+    :CP next --lang cpp             " Next problem switches to C++
+<
+Language Validation ~
+
+If you request a language that isn't enabled for a platform, cp.nvim will show
+a helpful error message listing available languages for that platform.
+
 ==============================================================================
 WORKFLOW                                                         *cp-workflow*
 
@@ -374,6 +439,7 @@ https://atcoder.jp/contests/{contest_id}/tasks/{contest_id}_{problem_id}
 
 Usage examples: >
     :CP atcoder abc324            " Set up atcoder.jp/contests/abc324
+    :CP atcoder abc324 --lang python  " Set up with Python instead of default
 
 Codeforces ~
                                                                *cp-codeforces*
@@ -381,6 +447,7 @@ URL format: https://codeforces.com/contest/{contest_id}/problem/{problem_id}
 
 Usage examples: >
     :CP codeforces 1934           " Set up codeforces.com/contest/1934
+    :CP codeforces 1934 --lang cpp    " Set up with C++
 
 CSES ~
                                                                      *cp-cses*
@@ -404,7 +471,7 @@ Example: Setting up and solving AtCoder contest ABC324
 
 3. Code your solution, then test: >
     :CP run
-<   View test verdicts in I/O splits. For detailed analysis:
+<   View test verdicts in I/O splits. For detailed analysis: >
     :CP panel
 <   Navigate tests with <c-n>/<c-p>, exit with q
 
@@ -414,12 +481,16 @@ Example: Setting up and solving AtCoder contest ABC324
 
 5. Continue solving problems with :CP next/:CP prev navigation
 
-6. Switch to another file (e.g. previous contest): >
+6. Try a different language for a problem: >
+    :CP C --lang python
+<   Opens problem C with Python instead of C++
+
+7. Switch to another file (e.g. previous contest): >
     :e ~/contests/abc323/a.cpp
     :CP
 <   Automatically restores abc323 contest context
 
-7. Submit solutions on AtCoder website
+8. Submit solutions on AtCoder website
 
 ==============================================================================
 I/O VIEW                                                          *cp-io-view*
diff --git a/lua/cp/commands/init.lua b/lua/cp/commands/init.lua
index 70c07b0..fa4f658 100644
--- a/lua/cp/commands/init.lua
+++ b/lua/cp/commands/init.lua
@@ -18,6 +18,7 @@ local actions = constants.ACTIONS
 ---@field interactor_cmd? string
 ---@field test_index? integer
 ---@field debug? boolean
+---@field language? string
 
 --- Turn raw args into normalized structure to later dispatch
 ---@param args string[] The raw command-line mode args
@@ -79,7 +80,16 @@ local function parse_command(args)
 
       return { type = 'action', action = first, test_index = test_index, debug = debug }
     else
-      return { type = 'action', action = first }
+      local language = nil
+      if #args >= 3 and args[2] == '--lang' then
+        language = args[3]
+      elseif #args >= 2 and args[2] ~= nil and args[2]:sub(1, 2) ~= '--' then
+        return {
+          type = 'error',
+          message = ("Unknown argument '%s' for action '%s'"):format(args[2], first),
+        }
+      end
+      return { type = 'action', action = first, language = language }
     end
   end
 
@@ -95,13 +105,18 @@ local function parse_command(args)
         platform = first,
         contest = args[2],
       }
-    elseif #args == 3 then
+    elseif #args == 4 and args[3] == '--lang' then
       return {
-        type = 'error',
-        message = 'Setup contests with :CP <platform> <contest_id>.',
+        type = 'contest_setup',
+        platform = first,
+        contest = args[2],
+        language = args[4],
       }
     else
-      return { type = 'error', message = 'Too many arguments' }
+      return {
+        type = 'error',
+        message = 'Invalid arguments. Usage: :CP <platform> <contest> [--lang <language>]',
+      }
     end
   end
 
@@ -110,6 +125,12 @@ local function parse_command(args)
       type = 'problem_jump',
       problem_id = first,
     }
+  elseif #args == 3 and args[2] == '--lang' then
+    return {
+      type = 'problem_jump',
+      problem_id = first,
+      language = args[3],
+    }
   end
 
   return { type = 'error', message = 'Unknown command or no contest context.' }
@@ -139,12 +160,12 @@ function M.handle_command(opts)
     elseif cmd.action == 'panel' then
       ui.toggle_panel({ debug = cmd.debug, test_index = cmd.test_index })
     elseif cmd.action == 'next' then
-      setup.navigate_problem(1)
+      setup.navigate_problem(1, cmd.language)
     elseif cmd.action == 'prev' then
-      setup.navigate_problem(-1)
+      setup.navigate_problem(-1, cmd.language)
     elseif cmd.action == 'pick' then
       local picker = require('cp.commands.picker')
-      picker.handle_pick_action()
+      picker.handle_pick_action(cmd.language)
     end
   elseif cmd.type == 'problem_jump' then
     local platform = state.get_platform()
@@ -173,13 +194,13 @@ function M.handle_command(opts)
     end
 
     local setup = require('cp.setup')
-    setup.setup_contest(platform, contest_id, problem_id)
+    setup.setup_contest(platform, contest_id, problem_id, cmd.language)
   elseif cmd.type == 'cache' then
     local cache_commands = require('cp.commands.cache')
     cache_commands.handle_cache_command(cmd)
   elseif cmd.type == 'contest_setup' then
     local setup = require('cp.setup')
-    setup.setup_contest(cmd.platform, cmd.contest, nil)
+    setup.setup_contest(cmd.platform, cmd.contest, nil, cmd.language)
     return
   end
 end
diff --git a/lua/cp/commands/picker.lua b/lua/cp/commands/picker.lua
index a733b58..08cbcd9 100644
--- a/lua/cp/commands/picker.lua
+++ b/lua/cp/commands/picker.lua
@@ -4,8 +4,9 @@ local config_module = require('cp.config')
 local logger = require('cp.log')
 
 --- Dispatch `:CP pick` to appropriate picker
+---@param language? string
 ---@return nil
-function M.handle_pick_action()
+function M.handle_pick_action(language)
   local config = config_module.get_config()
 
   if not (config.ui and config.ui.picker) then
@@ -53,7 +54,7 @@ function M.handle_pick_action()
     picker = fzf_picker
   end
 
-  picker.pick()
+  picker.pick(language)
 end
 
 return M
diff --git a/lua/cp/config.lua b/lua/cp/config.lua
index 1f926b5..4ae06d1 100644
--- a/lua/cp/config.lua
+++ b/lua/cp/config.lua
@@ -351,6 +351,50 @@ function M.get_config()
   return current_config or M.defaults
 end
 
+---Validate and get effective language config for a platform
+---@param platform_id string
+---@param language_id string
+---@return { valid: boolean, effective?: CpLanguage, extension?: string, error?: string }
+function M.get_language_for_platform(platform_id, language_id)
+  local cfg = M.get_config()
+
+  if not cfg.platforms[platform_id] then
+    return { valid = false, error = string.format("Unknown platform '%s'", platform_id) }
+  end
+
+  if not cfg.languages[language_id] then
+    return { valid = false, error = string.format("Unknown language '%s'", language_id) }
+  end
+
+  local platform = cfg.platforms[platform_id]
+  if not vim.tbl_contains(platform.enabled_languages, language_id) then
+    local available = table.concat(platform.enabled_languages, ', ')
+    return {
+      valid = false,
+      error = string.format(
+        "Language '%s' not enabled for %s. Available: %s",
+        language_id,
+        platform_id,
+        available
+      ),
+    }
+  end
+
+  local effective = cfg.runtime.effective[platform_id][language_id]
+  if not effective then
+    return {
+      valid = false,
+      error = string.format('No effective config for %s/%s', platform_id, language_id),
+    }
+  end
+
+  return {
+    valid = true,
+    effective = effective,
+    extension = effective.extension,
+  }
+end
+
 ---@param contest_id string
 ---@param problem_id? string
 ---@return string
diff --git a/lua/cp/pickers/fzf_lua.lua b/lua/cp/pickers/fzf_lua.lua
index 6735f54..d1af5d2 100644
--- a/lua/cp/pickers/fzf_lua.lua
+++ b/lua/cp/pickers/fzf_lua.lua
@@ -2,7 +2,7 @@ local picker_utils = require('cp.pickers')
 
 local M = {}
 
-local function contest_picker(platform, refresh)
+local function contest_picker(platform, refresh, language)
   local constants = require('cp.constants')
   local platform_display_name = constants.PLATFORM_DISPLAY_NAMES[platform]
   local fzf = require('fzf-lua')
@@ -42,19 +42,24 @@ local function contest_picker(platform, refresh)
 
         if contest then
           local cp = require('cp')
-          cp.handle_command({ fargs = { platform, contest.id } })
+          local fargs = { platform, contest.id }
+          if language then
+            table.insert(fargs, '--lang')
+            table.insert(fargs, language)
+          end
+          cp.handle_command({ fargs = fargs })
         end
       end,
       ['ctrl-r'] = function()
         local cache = require('cp.cache')
         cache.clear_contest_list(platform)
-        contest_picker(platform, true)
+        contest_picker(platform, true, language)
       end,
     },
   })
 end
 
-function M.pick()
+function M.pick(language)
   local fzf = require('fzf-lua')
   local platforms = picker_utils.get_platforms()
   local entries = vim.tbl_map(function(platform)
@@ -79,7 +84,7 @@ function M.pick()
         end
 
         if platform then
-          contest_picker(platform.id)
+          contest_picker(platform.id, false, language)
         end
       end,
     },
diff --git a/lua/cp/pickers/telescope.lua b/lua/cp/pickers/telescope.lua
index 9b3c0db..0a7a676 100644
--- a/lua/cp/pickers/telescope.lua
+++ b/lua/cp/pickers/telescope.lua
@@ -8,7 +8,7 @@ local picker_utils = require('cp.pickers')
 
 local M = {}
 
-local function contest_picker(opts, platform, refresh)
+local function contest_picker(opts, platform, refresh, language)
   local constants = require('cp.constants')
   local platform_display_name = constants.PLATFORM_DISPLAY_NAMES[platform]
   local contests = picker_utils.get_platform_contests(platform, refresh)
@@ -43,13 +43,18 @@ local function contest_picker(opts, platform, refresh)
 
           if selection then
             local cp = require('cp')
-            cp.handle_command({ fargs = { platform, selection.value.id } })
+            local fargs = { platform, selection.value.id }
+            if language then
+              table.insert(fargs, '--lang')
+              table.insert(fargs, language)
+            end
+            cp.handle_command({ fargs = fargs })
           end
         end)
 
         map('i', '<c-r>', function()
           actions.close(prompt_bufnr)
-          contest_picker(opts, platform, true)
+          contest_picker(opts, platform, true, language)
         end)
 
         return true
@@ -58,9 +63,8 @@ local function contest_picker(opts, platform, refresh)
     :find()
 end
 
-function M.pick(opts)
-  opts = opts or {}
-
+function M.pick(language)
+  local opts = {}
   local platforms = picker_utils.get_platforms()
 
   pickers
@@ -83,7 +87,7 @@ function M.pick(opts)
           actions.close(prompt_bufnr)
 
           if selection then
-            contest_picker(opts, selection.value.id)
+            contest_picker(opts, selection.value.id, false, language)
           end
         end)
         return true
diff --git a/lua/cp/setup.lua b/lua/cp/setup.lua
index 50d603d..6129614 100644
--- a/lua/cp/setup.lua
+++ b/lua/cp/setup.lua
@@ -8,6 +8,18 @@ local logger = require('cp.log')
 local scraper = require('cp.scraper')
 local state = require('cp.state')
 
+---Get the language of the current file from cache
+---@return string|nil
+local function get_current_file_language()
+  local current_file = vim.fn.expand('%:p')
+  if current_file == '' then
+    return nil
+  end
+  cache.load()
+  local file_state = cache.get_file_state(current_file)
+  return file_state and file_state.language or nil
+end
+
 ---@class TestCaseLite
 ---@field input string
 ---@field expected string
@@ -86,6 +98,14 @@ function M.setup_contest(platform, contest_id, problem_id, language)
   state.set_platform(platform)
   state.set_contest_id(contest_id)
 
+  if language then
+    local lang_result = config_module.get_language_for_platform(platform, language)
+    if not lang_result.valid then
+      logger.log(lang_result.error, vim.log.levels.ERROR)
+      return
+    end
+  end
+
   local is_new_contest = old_platform ~= platform and old_contest_id ~= contest_id
 
   cache.load()
@@ -173,6 +193,15 @@ function M.setup_problem(problem_id, language)
   local config = config_module.get_config()
   local lang = language
     or (config.platforms[platform] and config.platforms[platform].default_language)
+
+  if language then
+    local lang_result = config_module.get_language_for_platform(platform, language)
+    if not lang_result.valid then
+      logger.log(lang_result.error, vim.log.levels.ERROR)
+      return
+    end
+  end
+
   local source_file = state.get_source_file(lang)
   if not source_file then
     return
@@ -235,7 +264,8 @@ function M.setup_problem(problem_id, language)
 end
 
 ---@param direction integer
-function M.navigate_problem(direction)
+---@param language? string
+function M.navigate_problem(direction, language)
   if direction == 0 then
     return
   end
@@ -274,7 +304,15 @@ function M.navigate_problem(direction)
     require('cp.ui.views').disable()
   end
 
-  M.setup_contest(platform, contest_id, problems[new_index].id)
+  local lang = language or get_current_file_language()
+  if lang then
+    local lang_result = config_module.get_language_for_platform(platform, lang)
+    if not lang_result.valid then
+      lang = nil
+    end
+  end
+
+  M.setup_contest(platform, contest_id, problems[new_index].id, lang)
 end
 
 return M

From 48d4c6f1133793cafdd4ce897fba100734347541 Mon Sep 17 00:00:00 2001
From: Barrett Ruth <barrett.ruth@ramp.com>
Date: Fri, 24 Oct 2025 01:21:16 -0400
Subject: [PATCH 2/5] feat: language

---
 lua/cp/runner/execute.lua |  2 +-
 lua/cp/runner/run.lua     |  2 +-
 lua/cp/setup.lua          |  2 ++
 lua/cp/state.lua          | 15 ++++++++++++++-
 4 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/lua/cp/runner/execute.lua b/lua/cp/runner/execute.lua
index a871d6f..8f004e5 100644
--- a/lua/cp/runner/execute.lua
+++ b/lua/cp/runner/execute.lua
@@ -164,7 +164,7 @@ function M.compile_problem(debug)
   local state = require('cp.state')
   local config = require('cp.config').get_config()
   local platform = state.get_platform()
-  local language = config.platforms[platform].default_language
+  local language = state.get_language() or config.platforms[platform].default_language
   local eff = config.runtime.effective[platform][language]
 
   local compile_config = (debug and eff.commands.debug) or eff.commands.build
diff --git a/lua/cp/runner/run.lua b/lua/cp/runner/run.lua
index c8fcd0c..eaec30b 100644
--- a/lua/cp/runner/run.lua
+++ b/lua/cp/runner/run.lua
@@ -109,7 +109,7 @@ local function run_single_test_case(test_case, debug)
   local substitutions = { source = source_file, binary = binary_file }
 
   local platform_config = config.platforms[state.get_platform() or '']
-  local language = platform_config.default_language
+  local language = state.get_language() or platform_config.default_language
   local eff = config.runtime.effective[state.get_platform() or ''][language]
   local run_template = eff and eff.commands and eff.commands.run or {}
   local cmd = build_command(run_template, substitutions)
diff --git a/lua/cp/setup.lua b/lua/cp/setup.lua
index 6129614..20f852a 100644
--- a/lua/cp/setup.lua
+++ b/lua/cp/setup.lua
@@ -202,6 +202,8 @@ function M.setup_problem(problem_id, language)
     end
   end
 
+  state.set_language(lang)
+
   local source_file = state.get_source_file(lang)
   if not source_file then
     return
diff --git a/lua/cp/state.lua b/lua/cp/state.lua
index caf5044..621b184 100644
--- a/lua/cp/state.lua
+++ b/lua/cp/state.lua
@@ -20,6 +20,8 @@
 ---@field set_contest_id fun(contest_id: string)
 ---@field get_problem_id fun(): string?
 ---@field set_problem_id fun(problem_id: string)
+---@field get_language fun(): string?
+---@field set_language fun(language: string)
 ---@field get_active_panel fun(): string?
 ---@field set_active_panel fun(panel: string?)
 ---@field get_base_name fun(): string?
@@ -42,6 +44,7 @@ local state = {
   platform = nil,
   contest_id = nil,
   problem_id = nil,
+  language = nil,
   test_cases = nil,
   saved_session = nil,
   active_panel = nil,
@@ -80,6 +83,16 @@ function M.set_problem_id(problem_id)
   state.problem_id = problem_id
 end
 
+---@return string?
+function M.get_language()
+  return state.language
+end
+
+---@param language string
+function M.set_language(language)
+  state.language = language
+end
+
 ---@return string?
 function M.get_base_name()
   local platform, contest_id, problem_id = M.get_platform(), M.get_contest_id(), M.get_problem_id()
@@ -112,7 +125,7 @@ function M.get_source_file(language)
     return nil
   end
 
-  local target_language = language or platform_cfg.default_language
+  local target_language = language or state.language or platform_cfg.default_language
   local eff = config.runtime.effective[plat] and config.runtime.effective[plat][target_language]
     or nil
   if not eff or not eff.extension then

From c48cf384a48c4ca454ffa493411f21e29c668f3d Mon Sep 17 00:00:00 2001
From: Barrett Ruth <barrett.ruth@ramp.com>
Date: Fri, 24 Oct 2025 01:21:54 -0400
Subject: [PATCH 3/5] feat: error on invalid language

---
 lua/cp/setup.lua | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/lua/cp/setup.lua b/lua/cp/setup.lua
index 20f852a..e6f9d21 100644
--- a/lua/cp/setup.lua
+++ b/lua/cp/setup.lua
@@ -306,8 +306,16 @@ function M.navigate_problem(direction, language)
     require('cp.ui.views').disable()
   end
 
+  if language then
+    local lang_result = config_module.get_language_for_platform(platform, language)
+    if not lang_result.valid then
+      logger.log(lang_result.error, vim.log.levels.ERROR)
+      return
+    end
+  end
+
   local lang = language or get_current_file_language()
-  if lang then
+  if lang and not language then
     local lang_result = config_module.get_language_for_platform(platform, lang)
     if not lang_result.valid then
       lang = nil

From f52244c534c7bef31e39b04f5e0e48760239ca89 Mon Sep 17 00:00:00 2001
From: Barrett Ruth <barrett.ruth@ramp.com>
Date: Fri, 24 Oct 2025 01:32:48 -0400
Subject: [PATCH 4/5] better errors

---
 lua/cp/config.lua | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/lua/cp/config.lua b/lua/cp/config.lua
index 4ae06d1..46b13b4 100644
--- a/lua/cp/config.lua
+++ b/lua/cp/config.lua
@@ -362,17 +362,22 @@ function M.get_language_for_platform(platform_id, language_id)
     return { valid = false, error = string.format("Unknown platform '%s'", platform_id) }
   end
 
+  local platform = cfg.platforms[platform_id]
+
   if not cfg.languages[language_id] then
-    return { valid = false, error = string.format("Unknown language '%s'", language_id) }
+    local available = table.concat(platform.enabled_languages, ', ')
+    return {
+      valid = false,
+      error = string.format("Unknown language '%s'. Available: [%s]", language_id, available),
+    }
   end
 
-  local platform = cfg.platforms[platform_id]
   if not vim.tbl_contains(platform.enabled_languages, language_id) then
     local available = table.concat(platform.enabled_languages, ', ')
     return {
       valid = false,
       error = string.format(
-        "Language '%s' not enabled for %s. Available: %s",
+        "Language '%s' not enabled for %s. Available: [%s]",
         language_id,
         platform_id,
         available

From 9d848eba225acc509289555be6894a3583d1631e Mon Sep 17 00:00:00 2001
From: Barrett Ruth <barrett.ruth@ramp.com>
Date: Fri, 24 Oct 2025 01:44:06 -0400
Subject: [PATCH 5/5] feat: improve autocomplete

---
 lua/cp/cache.lua         | 20 +++++++++
 lua/cp/commands/init.lua | 33 +++++++++++---
 plugin/cp.lua            | 93 ++++++++++++++++++++++++++++++----------
 3 files changed, 119 insertions(+), 27 deletions(-)

diff --git a/lua/cp/cache.lua b/lua/cp/cache.lua
index 5c56ef8..86d806f 100644
--- a/lua/cp/cache.lua
+++ b/lua/cp/cache.lua
@@ -92,6 +92,26 @@ function M.get_contest_data(platform, contest_id)
   return cache_data[platform][contest_id]
 end
 
+---Get all cached contest IDs for a platform
+---@param platform string
+---@return string[]
+function M.get_cached_contest_ids(platform)
+  vim.validate({
+    platform = { platform, 'string' },
+  })
+
+  if not cache_data[platform] then
+    return {}
+  end
+
+  local contest_ids = {}
+  for contest_id, _ in pairs(cache_data[platform]) do
+    table.insert(contest_ids, contest_id)
+  end
+  table.sort(contest_ids)
+  return contest_ids
+end
+
 ---@param platform string
 ---@param contest_id string
 ---@param problems Problem[]
diff --git a/lua/cp/commands/init.lua b/lua/cp/commands/init.lua
index fa4f658..f48727d 100644
--- a/lua/cp/commands/init.lua
+++ b/lua/cp/commands/init.lua
@@ -59,16 +59,15 @@ local function parse_command(args)
       local debug = false
       local test_index = nil
 
-      for i = 2, #args do
-        local arg = args[i]
-        if arg == '--debug' then
+      if #args == 2 then
+        if args[2] == '--debug' then
           debug = true
         else
-          local idx = tonumber(arg)
+          local idx = tonumber(args[2])
           if not idx then
             return {
               type = 'error',
-              message = ("Invalid argument '%s': expected test number or --debug"):format(arg),
+              message = ("Invalid argument '%s': expected test number or --debug"):format(args[2]),
             }
           end
           if idx < 1 or idx ~= math.floor(idx) then
@@ -76,6 +75,30 @@ local function parse_command(args)
           end
           test_index = idx
         end
+      elseif #args == 3 then
+        local idx = tonumber(args[2])
+        if not idx then
+          return {
+            type = 'error',
+            message = ("Invalid argument '%s': expected test number"):format(args[2]),
+          }
+        end
+        if idx < 1 or idx ~= math.floor(idx) then
+          return { type = 'error', message = ("'%s' is not a valid test index"):format(idx) }
+        end
+        if args[3] ~= '--debug' then
+          return {
+            type = 'error',
+            message = ("Invalid argument '%s': expected --debug"):format(args[3]),
+          }
+        end
+        test_index = idx
+        debug = true
+      elseif #args > 3 then
+        return {
+          type = 'error',
+          message = 'Too many arguments. Usage: :CP ' .. first .. ' [test_num] [--debug]',
+        }
       end
 
       return { type = 'action', action = first, test_index = test_index, debug = debug }
diff --git a/plugin/cp.lua b/plugin/cp.lua
index 5d4df32..b4954a4 100644
--- a/plugin/cp.lua
+++ b/plugin/cp.lua
@@ -22,12 +22,30 @@ end, {
       num_args = num_args + 1
     end
 
+    local function filter_candidates(candidates)
+      return vim.tbl_filter(function(cmd)
+        return cmd:find(ArgLead, 1, true) == 1
+      end, candidates)
+    end
+
+    local function get_enabled_languages(platform)
+      local config = require('cp.config').get_config()
+      if platform and config.platforms[platform] then
+        return config.platforms[platform].enabled_languages
+      end
+      return vim.tbl_keys(config.languages)
+    end
+
     if num_args == 2 then
       local candidates = {}
       local state = require('cp.state')
       local platform = state.get_platform()
       local contest_id = state.get_contest_id()
 
+      vim.list_extend(candidates, platforms)
+      table.insert(candidates, 'cache')
+      table.insert(candidates, 'pick')
+
       if platform and contest_id then
         vim.list_extend(candidates, actions)
         local cache = require('cp.cache')
@@ -39,44 +57,75 @@ end, {
           table.sort(ids)
           vim.list_extend(candidates, ids)
         end
-      else
-        vim.list_extend(candidates, platforms)
-        table.insert(candidates, 'cache')
-        table.insert(candidates, 'pick')
       end
 
-      return vim.tbl_filter(function(cmd)
-        return cmd:find(ArgLead, 1, true) == 1
-      end, candidates)
+      return filter_candidates(candidates)
     elseif num_args == 3 then
-      if args[2] == 'cache' then
-        return vim.tbl_filter(function(cmd)
-          return cmd:find(ArgLead, 1, true) == 1
-        end, { 'clear', 'read' })
+      if vim.tbl_contains(platforms, args[2]) then
+        local cache = require('cp.cache')
+        cache.load()
+        local contests = cache.get_cached_contest_ids(args[2])
+        return filter_candidates(contests)
+      elseif args[2] == 'cache' then
+        return filter_candidates({ 'clear', 'read' })
       elseif args[2] == 'interact' then
-        local cands = utils.cwd_executables()
-        return vim.tbl_filter(function(cmd)
-          return cmd:find(ArgLead, 1, true) == 1
-        end, cands)
+        return filter_candidates(utils.cwd_executables())
+      elseif args[2] == 'run' or args[2] == 'panel' then
+        local state = require('cp.state')
+        local platform = state.get_platform()
+        local contest_id = state.get_contest_id()
+        local problem_id = state.get_problem_id()
+        local candidates = { '--debug' }
+        if platform and contest_id and problem_id then
+          local cache = require('cp.cache')
+          cache.load()
+          local test_cases = cache.get_test_cases(platform, contest_id, problem_id)
+          if test_cases then
+            for i = 1, #test_cases do
+              table.insert(candidates, tostring(i))
+            end
+          end
+        end
+        return filter_candidates(candidates)
+      elseif args[2] == 'next' or args[2] == 'prev' or args[2] == 'pick' then
+        return filter_candidates({ '--lang' })
+      else
+        local state = require('cp.state')
+        if state.get_platform() and state.get_contest_id() then
+          return filter_candidates({ '--lang' })
+        end
       end
     elseif num_args == 4 then
       if args[2] == 'cache' and args[3] == 'clear' then
-        return vim.tbl_filter(function(cmd)
-          return cmd:find(ArgLead, 1, true) == 1
-        end, platforms)
+        return filter_candidates(platforms)
+      elseif args[3] == '--lang' then
+        local platform = require('cp.state').get_platform()
+        return filter_candidates(get_enabled_languages(platform))
+      elseif (args[2] == 'run' or args[2] == 'panel') and tonumber(args[3]) then
+        return filter_candidates({ '--debug' })
       elseif vim.tbl_contains(platforms, args[2]) then
         local cache = require('cp.cache')
         cache.load()
         local contest_data = cache.get_contest_data(args[2], args[3])
+        local candidates = { '--lang' }
         if contest_data and contest_data.problems then
-          local candidates = {}
           for _, problem in ipairs(contest_data.problems) do
             table.insert(candidates, problem.id)
           end
-          return vim.tbl_filter(function(cmd)
-            return cmd:find(ArgLead, 1, true) == 1
-          end, candidates)
         end
+        return filter_candidates(candidates)
+      end
+    elseif num_args == 5 then
+      if vim.tbl_contains(platforms, args[2]) then
+        if args[4] == '--lang' then
+          return filter_candidates(get_enabled_languages(args[2]))
+        else
+          return filter_candidates({ '--lang' })
+        end
+      end
+    elseif num_args == 6 then
+      if vim.tbl_contains(platforms, args[2]) and args[5] == '--lang' then
+        return filter_candidates(get_enabled_languages(args[2]))
       end
     end
     return {}