leetrs 1.0.0

A command-line tool to interact with LeetCode.
Documentation
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 {
    /// Authenticate with LeetCode
    Auth,
    /// Launch the TUI (Placeholder for now)
    Tui,
    /// Check auth status
    Status,
    /// Pick a problem
    Pick {
        identifier: String,
        language: Option<Language>,
        #[arg(short, long)]
        preview: bool,
    },
    /// Submit a problem to leetcode
    Submit {
        /// The path to your solution file (e.g., 'two_sum.rs')
        file: String,
    },
    /// Check leetrs version
    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:");
                    // LEETCODE_SESSION is massive. Printing it nicely so it wraps well.
                    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 {
                // 4. launch neovim with a vertical split
                println!("๐Ÿš€ launching neovim...");
                if !*preview {
                    let status = Command::new("nvim")
                        .arg(&desc)
                        .arg("-c")
                        .arg(format!("vsplit {}", code)) // Force a vertical split with the code file
                        .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;
                }
            };

            // 1. Read the file content
            let code = match std::fs::read_to_string(&file) {
                Ok(c) => c,
                Err(e) => {
                    eprintln!("โŒ Failed to read file '{}': {}", file, e);
                    return;
                }
            };

            // 2. Extract the slug from the filename (e.g., "two_sum.rs" -> "two-sum")
            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(),
            );
            // 3. Fetch the question to get its internal ID
            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;
                }
            };

            // 4. Submit the code
            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;
                }
            };

            // 5. Poll for results
            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;
                }
            };

            // 6. Display the formatted results
            println!("\n==================================================");

            let status = result.status_msg.unwrap_or_else(|| "Unknown".to_string());

            if status == "Accepted" {
                // Print Accepted in Green
                // println!("  โœ… \x1b[32m{}\x1b[0m", status);
                println!("  โœ… {}", status);
            } else {
                // Print Errors/Wrong Answers in Red
                // println!("  โŒ \x1b[31m{}\x1b[0m", status);
                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");
        }
    }
}

/// Handles prompting the user to paste their tokens manually
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,
    })
}

/// Automatically extracts LeetCode cookies from the specified browser
fn auto_extract_flow(browser: &str) -> Result<LeetCodeCredentials, String> {
    println!("\n๐Ÿ” Attempting to extract cookies from {}...", browser);

    // We only want to query cookies belonging to LeetCode to speed up the process
    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;

    // Search the returned cookies for the two we care about
    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(),
        ),
    }
}