use anyhow::anyhow;
use content_inspector::{ContentType, inspect};
use std::fs;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::path::Path;
use tiktoken_rs::CoreBPE;

const MAX_TOKENIZER_LINE_BYTES: usize = 200_000;

fn has_overly_long_line(bytes: &[u8]) -> bool {
    let mut current = 0usize;
    for byte in bytes {
        match byte {
            b'\n' | b'\r' => current = 0,
            _ => {
                current += 1;
                if current > MAX_TOKENIZER_LINE_BYTES {
                    return true;
                }
            }
        }
    }
    false
}

pub(crate) fn init_tokenizer(model: &str) -> anyhow::Result<CoreBPE> {
    tiktoken_rs::get_bpe_from_model(model)
        .map_err(|err| anyhow!("Failed to initialize tokenizer model '{model}': {err}"))
}

pub(crate) fn count_tokens_in_text_file(path: &Path, tokenizer: &CoreBPE) -> Option<u64> {
    let bytes = fs::read(path).ok()?;

    if inspect(&bytes) == ContentType::BINARY {
        return None;
    }

    if has_overly_long_line(&bytes) {
        eprintln!(
            "Skipping '{}': file contains a line longer than {} bytes which can trigger tokenizer regex overflows",
            path.display(),
            MAX_TOKENIZER_LINE_BYTES
        );
        return None;
    }

    let text = String::from_utf8_lossy(&bytes);
    match catch_unwind(AssertUnwindSafe(|| {
        tokenizer.encode_with_special_tokens(&text).len() as u64
    })) {
        Ok(token_count) => Some(token_count),
        Err(_) => {
            eprintln!(
                "Skipping '{}': tokenizer panicked while counting tokens for this file",
                path.display()
            );
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;
    use std::time::{SystemTime, UNIX_EPOCH};

    fn temp_dir(prefix: &str) -> PathBuf {
        let ts = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("valid clock")
            .as_nanos();
        let pid = std::process::id();
        let path = std::env::temp_dir().join(format!("dua-{prefix}-{pid}-{ts}"));
        fs::create_dir_all(&path).expect("create temp dir");
        path
    }

    #[test]
    fn text_file_counts_tokens() {
        let dir = temp_dir("token-text");
        let path = dir.join("text.txt");
        let text = "hello world";
        fs::write(&path, text).expect("write text file");

        let tokenizer = init_tokenizer("gpt-5").expect("init tokenizer");
        let expected = tokenizer.encode_with_special_tokens(text).len() as u64;

        assert_eq!(count_tokens_in_text_file(&path, &tokenizer), Some(expected));

        fs::remove_dir_all(dir).ok();
    }

    #[test]
    fn binary_file_is_skipped() {
        let dir = temp_dir("token-binary");
        let path = dir.join("binary.bin");
        fs::write(&path, [0_u8, 1, 2, 3, 0, 4]).expect("write binary file");

        let tokenizer = init_tokenizer("gpt-5").expect("init tokenizer");
        assert_eq!(count_tokens_in_text_file(&path, &tokenizer), None);

        fs::remove_dir_all(dir).ok();
    }

    #[test]
    fn invalid_model_returns_actionable_error() {
        let err = match init_tokenizer("definitely-not-a-real-model") {
            Ok(_) => panic!("invalid model must fail"),
            Err(err) => err.to_string(),
        };

        assert!(
            err.contains("definitely-not-a-real-model"),
            "error should mention model name: {err}"
        );
    }

    #[test]
    fn long_single_line_file_is_skipped() {
        let dir = temp_dir("token-long-line");
        let path = dir.join("long.txt");
        let long_line = "a".repeat(MAX_TOKENIZER_LINE_BYTES + 1);
        fs::write(&path, long_line).expect("write long line file");

        let tokenizer = init_tokenizer("gpt-5").expect("init tokenizer");
        assert_eq!(count_tokens_in_text_file(&path, &tokenizer), None);

        fs::remove_dir_all(dir).ok();
    }
}
