mod config;
mod llm;
mod tools;

use std::io::{IsTerminal, Read, Write};
use std::path::PathBuf;
use std::time::{Duration, Instant};

use llm::{Chat, Image, Item};
use tools::Policy;

/// No args on a terminal starts an interactive session. Anything else is one
/// shot:
///
///   si                                  # interactive chat
///   si "what is a 3DS?"                 # one shot
///   si "look around" --yes              # one shot, tools run without asking
///   si "describe this" photo.jpg        # one shot, with an image
///   echo "what is a 3DS?" | si          # one shot from stdin
///
/// In the REPL, image paths typed in the prompt are attached to the message.
#[tokio::main]
async fn main() {
    if let Err(err) = run().await {
        eprintln!("error: {err:#}");
        std::process::exit(1);
    }
}

async fn run() -> anyhow::Result<()> {
    let raw_args: Vec<String> = std::env::args().skip(1).collect();

    let auto_approve = raw_args.iter().any(|a| a == "--yes" || a == "-y");
    let args: Vec<String> = raw_args
        .into_iter()
        .filter(|a| a != "--yes" && a != "-y")
        .collect();

    // Nothing but a terminal to talk to: interactive session.
    if args.is_empty() && std::io::stdin().is_terminal() {
        return repl(Policy { auto_approve }, Vec::new()).await;
    }

    let (prompt, images) = split_args(args)?;
    let policy = Policy { auto_approve };

    // Images but no question, on a real terminal: rather than blocking on
    // stdin, open a session with the images loaded so the user can just ask.
    if prompt.is_empty() && !images.is_empty() && std::io::stdin().is_terminal() {
        return repl(policy, images).await;
    }

    let mut chat = Chat::new(policy)?;
    let mut sink = OutputSink::new();

    let result = chat.send(&prompt, &images, |item| sink.on_item(item)).await;

    match result {
        Ok(_) => {
            sink.finish();
            Ok(())
        }
        Err(err) => {
            sink.abort();
            Err(err)
        }
    }
}

/// Multi-turn session: read a line, stream a reply, repeat until Ctrl-D.
///
/// `pending_images` are attached to the first message sent, which is how
/// `si photo.jpg` gets its image into the conversation.
async fn repl(policy: Policy, pending_images: Vec<Image>) -> anyhow::Result<()> {
    let mut chat = Chat::new(policy)?;
    let mut sink = OutputSink::new();
    let mut pending_images = pending_images;

    println!("si {} — /help for commands", config::MODEL);
    if !pending_images.is_empty() {
        println!(
            "\x1b[2m{} image(s) attached — ask about them\x1b[0m",
            pending_images.len()
        );
    }

    let stdin = std::io::stdin();
    loop {
        print!("\n> ");
        let _ = std::io::stdout().flush();

        let mut line = String::new();
        if stdin.read_line(&mut line)? == 0 {
            println!();
            break; // Ctrl-D
        }

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

        match input {
            "/quit" | "/exit" => break,
            "/clear" => {
                chat.reset();
                pending_images.clear();
                sink.note("history cleared");
                continue;
            }
            "/help" => {
                sink.note(
                    "/quit exit · /clear forget history · /turns count · image paths are attached",
                );
                continue;
            }
            "/turns" => {
                sink.note(&format!("{} turn(s) in history", chat.turn_count()));
                continue;
            }
            _ => {}
        }

        // Any word naming an image file becomes an attachment.
        let (text, images) = match split_repl_input(input) {
            Ok(split) => split,
            Err(err) => {
                sink.note(&format!("error: {err:#}"));
                continue;
            }
        };

        // Images handed over at startup ride along with the first question.
        let mut images = images;
        if !pending_images.is_empty() {
            images.append(&mut pending_images);
        }

        match chat.send(&text, &images, |item| sink.on_item(item)).await {
            Ok(()) => sink.finish(),
            // Keep the session alive; the turn was rolled back.
            Err(err) => {
                sink.abort();
                eprintln!("error: {err:#}");
            }
        }
    }

    Ok(())
}

/// Pull image paths out of REPL input. Words with an image extension are
/// attachments; everything else is prompt text.
fn split_repl_input(input: &str) -> anyhow::Result<(String, Vec<Image>)> {
    let mut words: Vec<&str> = Vec::new();
    let mut images: Vec<Image> = Vec::new();

    for word in input.split_whitespace() {
        let path = PathBuf::from(word);
        if Image::mime_for(&path).is_some() {
            if !path.is_file() {
                anyhow::bail!("image not found: {}", path.display());
            }
            images.push(Image::load(&path)?);
        } else {
            words.push(word);
        }
    }

    Ok((words.join(" "), images))
}

/// Split argv into text prompt and image attachments. Anything with an image
/// extension is an image, not prompt text.
fn split_args(args: Vec<String>) -> anyhow::Result<(String, Vec<Image>)> {
    let mut words: Vec<String> = Vec::new();
    let mut images: Vec<Image> = Vec::new();

    for arg in args {
        let path = PathBuf::from(&arg);
        if Image::mime_for(&path).is_some() {
            if !path.is_file() {
                anyhow::bail!("image not found: {}", path.display());
            }
            images.push(Image::load(&path)?);
        } else {
            words.push(arg);
        }
    }

    // No prompt on the command line: fall back to stdin, but only when it's a
    // pipe. Blocking on a terminal's stdin here would silently hang.
    if words.is_empty() && !std::io::stdin().is_terminal() {
        let mut buf = String::new();
        std::io::stdin()
            .read_to_string(&mut buf)
            .map_err(|e| anyhow::anyhow!("reading stdin: {e}"))?;
        words.push(buf);
    }

    let prompt = words.join(" ").trim().to_string();
    if prompt.is_empty() && images.is_empty() {
        anyhow::bail!("empty prompt (pass text and/or an image path)");
    }
    Ok((prompt, images))
}

/// Minimum gap between live re-renders. Without this, a fast token stream
/// repaints the terminal hundreds of times a second.
const MIN_REPAINT_INTERVAL: Duration = Duration::from_millis(40);

/// Streams the reply to stdout.
///
/// On a terminal, the whole reply is kept as raw Markdown and re-rendered
/// through termimad periodically, so it always ends up fully formatted. When
/// piped, the raw Markdown is emitted straight through and never touched.
struct OutputSink {
    styled: bool,
    source: String,
    /// Terminal rows the last repaint occupied, so we know how far to move the
    /// cursor back up before overwriting.
    painted_rows: usize,
    last_repaint: Instant,
    /// When the in-flight tool call started, for timing its completion line.
    tool_started: Option<Instant>,
    skin: termimad::MadSkin,
}

impl OutputSink {
    fn new() -> Self {
        Self {
            styled: std::io::stdout().is_terminal(),
            source: String::new(),
            painted_rows: 0,
            last_repaint: Instant::now(),
            tool_started: None,
            skin: termimad::MadSkin::default(),
        }
    }

    /// A short non-model message (status, help). Ends the current paint so the
    /// next reply repaints from a clean slate.
    fn note(&mut self, text: &str) {
        if self.styled {
            self.clear();
        }
        println!("\x1b[2m{text}\x1b[0m");
        let _ = std::io::stdout().flush();
    }

    /// One ordered stream of model text and tool activity.
    ///
    /// Tool lines are dimmed and column-aligned so they read as machinery
    /// running alongside the answer rather than as part of it.
    fn on_item(&mut self, item: Item) {
        match item {
            Item::Text(delta) => self.on_text(&delta),

            Item::ToolStart { name, summary } => {
                // The reply pauses while a tool runs, so drop whatever was
                // painted; the next text repaint starts from a clean slate.
                if self.styled {
                    self.clear();
                }
                self.tool_started = Some(Instant::now());

                if !self.styled {
                    return; // keep piped output as pure model text
                }
                let args = summary
                    .strip_prefix(&format!("{name}("))
                    .and_then(|s| s.strip_suffix(')'))
                    .unwrap_or(&summary);
                println!("  \x1b[36m{name:<11}\x1b[0m\x1b[2m{args}\x1b[0m");
                let _ = std::io::stdout().flush();
            }

            Item::ToolDone { outcome } => {
                let Some(started) = self.tool_started.take() else {
                    return;
                };
                if !self.styled {
                    return;
                }
                let secs = started.elapsed().as_secs_f32();
                let timing = if secs >= 1.0 {
                    format!("{secs:.1}s")
                } else {
                    format!("{:.0}ms", secs * 1000.0)
                };
                // Errors in red, everything else green; both dimmed so the
                // answer stays the brightest thing on screen.
                let color = if outcome.starts_with("error:") {
                    "31"
                } else if outcome == "skipped" {
                    "33"
                } else {
                    "32"
                };
                println!(
                    "  \x1b[2m╰─\x1b[0m \x1b[2;{color}m{outcome}\x1b[0m\x1b[2m · {timing}\x1b[0m"
                );
                let _ = std::io::stdout().flush();
            }
        }
    }

    fn on_text(&mut self, delta: &str) {
        self.source.push_str(delta);

        if !self.styled {
            self.flush_raw(delta);
            return;
        }

        // Newlines change the document's shape, so repaint immediately rather
        // than waiting out the throttle.
        let structural = delta.contains('\n');
        if structural || self.last_repaint.elapsed() >= MIN_REPAINT_INTERVAL {
            self.repaint();
        }
    }

    fn finish(&mut self) {
        if self.styled {
            self.repaint();
        }
        println!();
        // The trailing newline moved us past this reply; the next repaint
        // must not clear it.
        self.painted_rows = 0;
        self.source.clear();
        let _ = std::io::stdout().flush();
    }

    /// On error: drop the half-painted render and leave a clean error line for
    /// the caller to print.
    fn abort(&mut self) {
        if self.styled {
            self.clear();
        }
        self.source.clear();
        self.tool_started = None;
        let _ = std::io::stdout().flush();
    }

    fn flush_raw(&mut self, delta: &str) {
        print!("{delta}");
        let _ = std::io::stdout().flush();
    }

    /// Move the cursor back above the previous render and wipe it.
    fn clear(&mut self) {
        if self.painted_rows > 0 {
            // ESC[<n>A: cursor up n rows. ESC[J: erase from cursor to screen end.
            print!("\x1b[{}A\x1b[J", self.painted_rows);
        }
        self.painted_rows = 0;
    }

    fn repaint(&mut self) {
        let width = termimad::terminal_size().0 as usize;
        // FmtText renders itself to a string; counting the lines in that gives
        // exactly how many terminal rows the repaint will occupy.
        let rendered = self.skin.text(&self.source, Some(width)).to_string();
        let rows = rendered.matches('\n').count();

        self.clear();
        print!("{rendered}");
        let _ = std::io::stdout().flush();

        self.painted_rows = rows;
        self.last_repaint = Instant::now();
    }
}