use chrono::Local;
use glob::glob;
use regex::Regex;
use rustyline::DefaultEditor;
use rustyline::error::ReadlineError;
use serde::{Deserialize, Serialize};
use std::env;
use std::fs;
use std::io::{self, Write};
use std::path::PathBuf;
use std::process::Command;
use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
};
use std::thread;
use std::time::Duration;
use walkdir::WalkDir;

use menta::{GenerateTextRequest, ModelMessage, Tool, ToolChoice, ToolExecute, generate_text};

const RESET: &str = "\x1b[0m";
const DIM: &str = "\x1b[2m";
const BLUE: &str = "\x1b[34m";
const CYAN: &str = "\x1b[36m";
const GREEN: &str = "\x1b[32m";
const YELLOW: &str = "\x1b[33m";
const RED: &str = "\x1b[31m";

fn paint(color: &str, value: impl AsRef<str>) -> String {
    format!("{color}{}{RESET}", value.as_ref())
}

struct SpinnerState {
    active: AtomicBool,
    paused: AtomicBool,
}

struct LoadingIndicator {
    state: Arc<SpinnerState>,
    handle: Option<thread::JoinHandle<()>>,
}

impl LoadingIndicator {
    fn start(message: &'static str) -> Self {
        let state = Arc::new(SpinnerState {
            active: AtomicBool::new(true),
            paused: AtomicBool::new(false),
        });
        let state_for_thread = Arc::clone(&state);

        let handle = thread::spawn(move || {
            let frames = ['|', '/', '-', '\\'];
            let mut index = 0;

            while state_for_thread.active.load(Ordering::Relaxed) {
                if state_for_thread.paused.load(Ordering::Relaxed) {
                    thread::sleep(Duration::from_millis(25));
                    continue;
                }

                print!(
                    "\r{} {}",
                    paint(CYAN, frames[index % frames.len()].to_string()),
                    paint(DIM, message),
                );
                let _ = io::stdout().flush();
                index += 1;
                thread::sleep(Duration::from_millis(100));
            }
        });

        Self {
            state,
            handle: Some(handle),
        }
    }

    fn state(&self) -> Arc<SpinnerState> {
        Arc::clone(&self.state)
    }

    fn stop(mut self) {
        self.state.active.store(false, Ordering::Relaxed);
        if let Some(handle) = self.handle.take() {
            let _ = handle.join();
        }
        print!("\r{}\r", " ".repeat(40));
        let _ = io::stdout().flush();
    }
}

struct SpinnerPause {
    state: Arc<SpinnerState>,
}

impl SpinnerPause {
    fn new(state: Arc<SpinnerState>) -> Self {
        state.paused.store(true, Ordering::Relaxed);
        print!("\r{}\r", " ".repeat(60));
        let _ = io::stdout().flush();
        Self { state }
    }
}

impl Drop for SpinnerPause {
    fn drop(&mut self) {
        self.state.paused.store(false, Ordering::Relaxed);
    }
}

fn logged_tool<T>(verbose: bool, spinner_state: Arc<SpinnerState>) -> Tool
where
    T: ToolExecute,
{
    let definition = T::definition();
    let name = definition.name.clone();

    Tool::new_async(
        definition.name,
        definition.description,
        definition.input_schema,
        definition.output_schema,
        move |input| {
            let name = name.clone();
            let spinner_state = Arc::clone(&spinner_state);
            async move {
                let _pause = SpinnerPause::new(spinner_state);
                println!("{} {} {}", paint(YELLOW, "tool>"), paint(YELLOW, &name), paint(DIM, &input));

                let parsed = serde_json::from_str::<T>(&input)
                    .map_err(|error| format!("invalid tool input for {name}: {error}"))?;
                let output = parsed.execute().await?;
                let output = serde_json::to_string(&output)
                    .map_err(|error| format!("invalid tool output for {name}: {error}"))?;

                if verbose {
                    println!("{} {} {}", paint(YELLOW, "tool<"), paint(YELLOW, &name), paint(DIM, &output));
                }
                Ok(output)
            }
        },
    )
}

#[derive(Deserialize, Tool)]
#[tool(description = "Get the weather in a location")]
struct WeatherTool {
    #[description = "The location to get the weather for"]
    location: String,
}

#[derive(Serialize)]
struct WeatherOutput {
    temperature: i32,
    conditions: String,
}

impl ToolExecute for WeatherTool {
    type Output = WeatherOutput;

    async fn execute(&self) -> std::result::Result<Self::Output, String> {
        Ok(WeatherOutput {
            temperature: 72,
            conditions: format!("sunny in {}", self.location),
        })
    }
}

#[derive(Deserialize, Tool)]
#[tool(description = "Run a local shell command")]
struct BashTool {
    #[description = "The shell command to run"]
    command: String,
}

#[derive(Serialize)]
struct BashOutput {
    success: bool,
    status: i32,
    stdout: String,
    stderr: String,
}

impl ToolExecute for BashTool {
    type Output = BashOutput;

    async fn execute(&self) -> std::result::Result<Self::Output, String> {
        let output = Command::new("sh")
            .arg("-c")
            .arg(&self.command)
            .output()
            .map_err(|error| error.to_string())?;

        Ok(BashOutput {
            success: output.status.success(),
            status: output.status.code().unwrap_or(-1),
            stdout: String::from_utf8_lossy(&output.stdout).into_owned(),
            stderr: String::from_utf8_lossy(&output.stderr).into_owned(),
        })
    }
}

#[derive(Deserialize, Tool)]
#[tool(description = "Write content to a local file")]
struct WriteTool {
    #[description = "Path to the file to write"]
    path: String,
    #[description = "Full file content to write"]
    content: String,
}

#[derive(Serialize)]
struct WriteOutput {
    path: String,
    bytes_written: usize,
}

impl ToolExecute for WriteTool {
    type Output = WriteOutput;

    async fn execute(&self) -> std::result::Result<Self::Output, String> {
        if let Some(parent) = PathBuf::from(&self.path).parent() {
            if !parent.as_os_str().is_empty() {
                fs::create_dir_all(parent).map_err(|error| error.to_string())?;
            }
        }

        fs::write(&self.path, &self.content).map_err(|error| error.to_string())?;

        Ok(WriteOutput {
            path: self.path.clone(),
            bytes_written: self.content.len(),
        })
    }
}

#[derive(Deserialize, Tool)]
#[tool(description = "Read a local file")]
struct ReadTool {
    #[description = "Path to the file to read"]
    path: String,
}

#[derive(Serialize)]
struct ReadOutput {
    path: String,
    content: String,
}

impl ToolExecute for ReadTool {
    type Output = ReadOutput;

    async fn execute(&self) -> std::result::Result<Self::Output, String> {
        let content = fs::read_to_string(&self.path).map_err(|error| error.to_string())?;

        Ok(ReadOutput {
            path: self.path.clone(),
            content,
        })
    }
}

#[derive(Deserialize, Tool)]
#[tool(description = "Find files by glob pattern")]
struct GlobTool {
    #[description = "Glob pattern like **/*.rs or src/**/*.rs"]
    pattern: String,
}

#[derive(Serialize)]
struct GlobOutput {
    matches: Vec<String>,
}

impl ToolExecute for GlobTool {
    type Output = GlobOutput;

    async fn execute(&self) -> std::result::Result<Self::Output, String> {
        let matches = glob(&self.pattern)
            .map_err(|error| error.to_string())?
            .filter_map(|entry| entry.ok())
            .map(|path| path.display().to_string())
            .collect::<Vec<_>>();

        Ok(GlobOutput { matches })
    }
}

#[derive(Deserialize, Tool)]
#[tool(description = "Search file contents with a regex pattern")]
struct GrepTool {
    #[description = "Regex pattern to search for"]
    pattern: String,
    #[description = "Root path to search from, defaults to current directory"]
    path: Option<String>,
}

#[derive(Serialize)]
struct GrepOutput {
    matches: Vec<String>,
}

impl ToolExecute for GrepTool {
    type Output = GrepOutput;

    async fn execute(&self) -> std::result::Result<Self::Output, String> {
        let regex = Regex::new(&self.pattern).map_err(|error| error.to_string())?;
        let root = self.path.as_deref().unwrap_or(".");
        let mut matches = Vec::new();

        for entry in WalkDir::new(root) {
            let entry = entry.map_err(|error| error.to_string())?;
            if !entry.file_type().is_file() {
                continue;
            }

            let path = entry.path();
            let Ok(content) = fs::read_to_string(path) else {
                continue;
            };

            for (index, line) in content.lines().enumerate() {
                if regex.is_match(line) {
                    matches.push(format!("{}:{}:{}", path.display(), index + 1, line));
                }
            }
        }

        Ok(GrepOutput { matches })
    }
}

#[tokio::main]
async fn main() {
    let verbose = env::args().skip(1).any(|arg| arg == "--verbose");

    println!("{}", paint(BLUE, "menta agent example"));
    println!("{}", paint(DIM, "commands: /new, /model [model-id], exit, quit, !<shell-command>"));
    if verbose {
        println!("{}", paint(DIM, "verbose tool logging enabled"));
    }

    let mut model = String::from("openai/gpt-4.1-mini");
    let mut history = vec![ModelMessage::system(system_prompt())];
    let mut editor = DefaultEditor::new().expect("failed to initialize line editor");

    loop {
        let prompt = match editor.readline(&format!("[{model}] > ")) {
            Ok(prompt) => prompt,
            Err(ReadlineError::Interrupted) => {
                println!();
                continue;
            }
            Err(ReadlineError::Eof) => break,
            Err(error) => {
                eprintln!("error: failed to read input: {error}");
                break;
            }
        };

        let prompt = prompt.trim();
        if prompt.is_empty() {
            continue;
        }

        let _ = editor.add_history_entry(prompt);

        if matches!(prompt, "exit" | "quit") {
            break;
        }

        if prompt == "/new" {
            history = vec![ModelMessage::system(system_prompt())];
            println!("started new session");
            continue;
        }

        if let Some(next_model) = prompt.strip_prefix("/model") {
            let next_model = next_model.trim();
            if next_model.is_empty() {
                println!("{} {}", paint(BLUE, "current model:"), paint(CYAN, &model));
            } else {
                model = next_model.to_string();
                println!("{} {}", paint(BLUE, "switched model to"), paint(CYAN, &model));
            }
            continue;
        }

        if let Some(command) = prompt.strip_prefix('!') {
            let command = command.trim();
            if command.is_empty() {
                eprintln!("{} missing shell command after !", paint(RED, "error:"));
                continue;
            }

            match Command::new("sh").arg("-c").arg(command).output() {
                Ok(output) => {
                    let stdout = String::from_utf8_lossy(&output.stdout);
                    let stderr = String::from_utf8_lossy(&output.stderr);

                    if !stdout.trim().is_empty() {
                        print!("{stdout}");
                    }

                    if !stderr.trim().is_empty() {
                        eprint!("{stderr}");
                    }

                    if !output.status.success() {
                        eprintln!("{} command exited with status {}", paint(RED, "error:"), output.status);
                    }
                }
                Err(error) => eprintln!("{} failed to run command: {error}", paint(RED, "error:")),
            }
            continue;
        }

        let mut messages = history.clone();
        messages.push(ModelMessage::user(prompt));

        let loading = LoadingIndicator::start("thinking...");
        let spinner_state = loading.state();

        let result = generate_text(
            GenerateTextRequest::new()
                .model(model.clone())
                .messages(messages)
                .tools(vec![
                    logged_tool::<WeatherTool>(verbose, Arc::clone(&spinner_state)),
                    logged_tool::<BashTool>(verbose, Arc::clone(&spinner_state)),
                    logged_tool::<WriteTool>(verbose, Arc::clone(&spinner_state)),
                    logged_tool::<ReadTool>(verbose, Arc::clone(&spinner_state)),
                    logged_tool::<GlobTool>(verbose, Arc::clone(&spinner_state)),
                    logged_tool::<GrepTool>(verbose, Arc::clone(&spinner_state)),
                ])
                .tool_choice(ToolChoice::Auto)
                .max_steps(4),
        )
        .await;

        loading.stop();

        match result {
            Ok(result) => {
                println!("{} {}", paint(GREEN, "assistant>"), result.output);
                history.push(ModelMessage::user(prompt));
                history.push(ModelMessage::assistant_text(result.text));
            }
            Err(error) => eprintln!("{} {error}", paint(RED, "error:")),
        }
    }
}

fn system_prompt() -> String {
    let date = Local::now().format("%Y-%m-%d").to_string();
    let pwd = env::current_dir()
        .map(|path| path.display().to_string())
        .unwrap_or_else(|_| String::from("unknown"));

    format!(
        "You are a local coding assistant. Current date: {date}. Current working directory: {pwd}. Use tools when needed: weather, bash, write, read, glob, grep. Prefer tools for filesystem and shell tasks."
    )
}
