use std::{
fs::{self},
process::Command,
};
use clap::{Parser, Subcommand};
use dialoguer::{Password, Select, theme::ColorfulTheme};
use leetrs::{auth::LeetCodeCredentials, client::LeetCodeClient, models::Language, picker::Picker};
#[derive(Parser, Debug)]
#[command(name = "leetrs")]
#[command(about = "A Neovim-integrated LeetCode TUI", long_about = None)]
struct Cli {
#[command(subcommand)]
command: Commands,
}
#[derive(Subcommand, Debug)]
enum Commands {
Auth,
Tui,
Status,
Pick {
identifier: String,
language: Option<Language>,
#[arg(short, long)]
preview: bool,
},
Submit {
file: String,
},
Version,
}
#[tokio::main]
async fn main() {
let cli = Cli::parse();
match &cli.command {
Commands::Auth => {
println!("๐ LeetCode Authentication\n");
let options = &[
"Paste tokens manually",
"Extract from Firefox",
"Extract from Chrome",
];
let selection = Select::with_theme(&ColorfulTheme::default())
.with_prompt("How would you like to authenticate?")
.default(0)
.items(&options[..])
.interact()
.unwrap();
let credentials_result = match selection {
0 => manual_auth_flow(),
1 => auto_extract_flow("firefox"),
2 => auto_extract_flow("chrome"),
_ => unreachable!(),
};
match credentials_result {
Ok(creds) => match creds.save() {
Ok(_) => println!("\nโ
Authentication successful!"),
Err(e) => eprintln!("\nโ Failed to save credentials: {}", e),
},
Err(e) => {
eprintln!("\nโ Authentication failed: {}", e);
if selection != 0 {
eprintln!(
"Tip: Make sure you are logged into leetcode.com on that browser, or try the manual option."
);
}
}
}
}
Commands::Tui => {
println!("TUI interface coming soon!");
}
Commands::Status => {
match LeetCodeCredentials::load() {
Some(creds) => {
println!("โ
Currently authenticated!");
println!("๐ csrftoken:");
println!("{}\n", creds.csrf_token);
println!("๐ LEETCODE_SESSION:");
println!("{}", creds.session_cookie);
}
None => {
eprintln!("โ Not authenticated. No valid credentials found.");
eprintln!("Run `leetrs auth` to set up your account.");
}
}
}
Commands::Pick {
identifier,
language,
preview,
} => {
let creds = match LeetCodeCredentials::load() {
Some(c) => c,
None => {
eprintln!("โ Not authenticated. Please run `leetrs auth` first.");
return;
}
};
let client = match LeetCodeClient::new(creds) {
Ok(c) => c,
Err(e) => {
eprintln!("โ Failed to initialize client: {}", e);
return;
}
};
let picker = Picker::new(client);
if let Ok((code, desc)) = picker.pick(identifier, language).await {
println!("๐ launching neovim...");
if !*preview {
let status = Command::new("nvim")
.arg(&desc)
.arg("-c")
.arg(format!("vsplit {}", code)) .status();
match status {
Ok(exit_status) if exit_status.success() => {
println!("\n๐ neovim closed.");
}
Ok(exit_status) => {
eprintln!("โ ๏ธ neovim exited with an error code: {}", exit_status);
}
Err(e) => {
eprintln!(
"โ failed to launch neovim. is it installed and in your path? error: {}",
e
);
}
}
} else {
let content = fs::read_to_string(desc);
if let Ok(content) = content {
print!("{}", content);
}
}
}
}
Commands::Submit { file } => {
let creds = match LeetCodeCredentials::load() {
Some(c) => c,
None => {
eprintln!("โ Not authenticated. Please run `leetrs auth` first.");
return;
}
};
let client = match LeetCodeClient::new(creds) {
Ok(c) => c,
Err(e) => {
eprintln!("โ Failed to initialize client: {}", e);
return;
}
};
let code = match std::fs::read_to_string(&file) {
Ok(c) => c,
Err(e) => {
eprintln!("โ Failed to read file '{}': {}", file, e);
return;
}
};
let path = std::path::Path::new(&file);
let file_stem = path
.file_stem()
.unwrap_or_default()
.to_str()
.unwrap_or_default();
let slug = file_stem.replace("_", "-");
println!("๐ Resolving ID for '{}'...", slug);
let language = Language::from_extension(
path.extension()
.and_then(|s| s.to_str())
.unwrap_or_default(),
);
let question = match client.get_question_by_slug(&slug, &language).await {
Ok(q) => q,
Err(e) => {
eprintln!(
"โ Failed to fetch question ID. Does the filename match the problem slug? Error: {}",
e
);
return;
}
};
println!("๐ Submitting {}...", file);
let submission_id = match client
.submit_code(&slug, &question.question_id, language.to_lang_slug(), &code)
.await
{
Ok(id) => id,
Err(e) => {
eprintln!("โ Submission failed: {}", e);
return;
}
};
println!("โณ Code queued. Waiting for execution results...");
let result = match client.check_submission(submission_id).await {
Ok(r) => r,
Err(e) => {
eprintln!("โ Failed to check submission status: {}", e);
return;
}
};
println!("\n==================================================");
let status = result.status_msg.unwrap_or_else(|| "Unknown".to_string());
if status == "Accepted" {
println!(" โ
{}", status);
} else {
println!(" โ {}", status);
}
println!("==================================================\n");
if let (Some(correct), Some(total)) = (result.total_correct, result.total_testcases) {
println!("๐งช Testcases: {} / {} passed", correct, total);
}
if status == "Accepted" {
if let Some(runtime) = result.status_runtime {
println!("โฑ๏ธ Runtime: {}", runtime);
}
if let Some(memory) = result.status_memory {
println!("๐พ Memory: {}", memory);
}
} else if status == "Compile Error" {
if let Some(err_msg) = result.compile_error {
println!("๐ฅ Compiler Output:\n{}", err_msg);
}
}
}
Commands::Version => {
println!("leetrs 1.0");
}
}
}
fn manual_auth_flow() -> Result<LeetCodeCredentials, String> {
println!("\nPlease extract your cookies from your browser session.");
println!("(Developer Tools -> Application -> Cookies -> leetcode.com)\n");
let session_cookie = Password::with_theme(&ColorfulTheme::default())
.with_prompt("Enter LEETCODE_SESSION cookie")
.interact()
.map_err(|e| e.to_string())?;
let csrf_token = Password::with_theme(&ColorfulTheme::default())
.with_prompt("Enter csrftoken cookie")
.interact()
.map_err(|e| e.to_string())?;
Ok(LeetCodeCredentials {
session_cookie,
csrf_token,
})
}
fn auto_extract_flow(browser: &str) -> Result<LeetCodeCredentials, String> {
println!("\n๐ Attempting to extract cookies from {}...", browser);
let domains = Some(vec!["leetcode.com".to_string()]);
let cookies = match browser {
"chrome" => {
rookie::chrome(domains).map_err(|e| format!("Chrome extraction failed: {}", e))?
}
"firefox" => {
rookie::firefox(domains).map_err(|e| format!("Firefox extraction failed: {}", e))?
}
_ => return Err("Unsupported browser".into()),
};
let mut session_cookie = None;
let mut csrf_token = None;
for cookie in cookies {
if cookie.name == "LEETCODE_SESSION" {
session_cookie = Some(cookie.value);
} else if cookie.name == "csrftoken" {
csrf_token = Some(cookie.value);
}
}
match (session_cookie, csrf_token) {
(Some(session), Some(csrf)) => Ok(LeetCodeCredentials {
session_cookie: session,
csrf_token: csrf,
}),
_ => Err(
"Could not find both LEETCODE_SESSION and csrftoken in the browser's database.".into(),
),
}
}