//! Utility functions for writing to [`stdout`](std::io::stdout) and [`stderr`](std::io::stderr).
//!
//! Originally from [cargo](https://github.com/rust-lang/cargo/blob/35814255a1dbaeca9219fae81d37a8190050092c/src/cargo/core/shell.rs).

use super::style::*;
use anstream::AutoStream;
use anstyle::Style;
use clap::ValueEnum;
use eyre::Result;
use serde::{Deserialize, Serialize};
use std::{
    fmt,
    io::{IsTerminal, prelude::*},
    ops::DerefMut,
    sync::{
        Mutex, OnceLock, PoisonError,
        atomic::{AtomicBool, Ordering},
    },
};

/// Returns the current color choice.
pub fn color_choice() -> ColorChoice {
    Shell::get().color_choice()
}

/// Returns the currently set verbosity level.
pub fn verbosity() -> Verbosity {
    Shell::get().verbosity()
}

/// Set the verbosity level.
pub fn set_verbosity(verbosity: Verbosity) {
    Shell::get().set_verbosity(verbosity);
}

/// Returns whether the output mode is [`OutputMode::Quiet`].
pub fn is_quiet() -> bool {
    Shell::get().output_mode().is_quiet()
}

/// Returns whether stderr is a terminal (tty).
///
/// Used to gate progress/spinner output that only makes sense for interactive use.
pub fn is_err_tty() -> bool {
    Shell::get().is_err_tty()
}

/// Returns whether the output format is [`OutputFormat::Json`].
pub fn is_json() -> bool {
    Shell::get().is_json()
}

/// Returns whether the output format is [`OutputFormat::Markdown`].
pub fn is_markdown() -> bool {
    Shell::get().is_markdown()
}

/// The global shell instance.
static GLOBAL_SHELL: OnceLock<Mutex<Shell>> = OnceLock::new();

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
/// The requested output mode.
pub enum OutputMode {
    /// Default output
    #[default]
    Normal,
    /// No output
    Quiet,
}

impl OutputMode {
    /// Returns true if the output mode is `Normal`.
    pub fn is_normal(self) -> bool {
        self == Self::Normal
    }

    /// Returns true if the output mode is `Quiet`.
    pub fn is_quiet(self) -> bool {
        self == Self::Quiet
    }
}

/// The requested output format.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub enum OutputFormat {
    /// Plain text output.
    #[default]
    Text,
    /// JSON output.
    Json,
    /// Plain text with markdown tables.
    Markdown,
}

impl OutputFormat {
    /// Returns true if the output format is `Text`.
    pub fn is_text(self) -> bool {
        self == Self::Text
    }

    /// Returns true if the output format is `Json`.
    pub fn is_json(self) -> bool {
        self == Self::Json
    }

    /// Returns true if the output format is `Markdown`.
    pub fn is_markdown(self) -> bool {
        self == Self::Markdown
    }
}

/// The verbosity level.
pub type Verbosity = u8;

/// An abstraction around console output that remembers preferences for output
/// verbosity and color.
pub struct Shell {
    /// Wrapper around stdout/stderr. This helps with supporting sending
    /// output to a memory buffer which is useful for tests.
    output: ShellOut,

    /// The format to use for message output.
    output_format: OutputFormat,

    /// The verbosity mode to use for message output.
    output_mode: OutputMode,

    /// The verbosity level to use for message output.
    verbosity: Verbosity,

    /// Flag that indicates the current line needs to be cleared before
    /// printing. Used when a progress bar is currently displayed.
    needs_clear: AtomicBool,
}

impl fmt::Debug for Shell {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut s = f.debug_struct("Shell");
        s.field("output_format", &self.output_format);
        s.field("output_mode", &self.output_mode);
        s.field("verbosity", &self.verbosity);
        if let ShellOut::Stream { color_choice, .. } = self.output {
            s.field("color_choice", &color_choice);
        }
        s.finish()
    }
}

/// A `Write`able object, either with or without color support.
enum ShellOut {
    /// Color-enabled stdio, with information on whether color should be used.
    Stream {
        stdout: AutoStream<std::io::Stdout>,
        stderr: AutoStream<std::io::Stderr>,
        stderr_tty: bool,
        color_choice: ColorChoice,
    },
    /// A write object that ignores all output.
    Empty(std::io::Empty),
    /// Captures stdout and stderr into in-memory buffers. Intended for tests.
    Captured { stdout: Vec<u8>, stderr: Vec<u8> },
}

/// Whether messages should use color output.
#[derive(Debug, Default, PartialEq, Eq, Clone, Copy, Serialize, Deserialize, ValueEnum)]
pub enum ColorChoice {
    /// Intelligently guess whether to use color output (default).
    #[default]
    Auto,
    /// Force color output.
    Always,
    /// Force disable color output.
    Never,
}

impl Default for Shell {
    fn default() -> Self {
        Self::new()
    }
}

impl Shell {
    /// Creates a new shell (color choice and verbosity), defaulting to 'auto' color and verbose
    /// output.
    pub fn new() -> Self {
        Self::new_with(
            OutputFormat::Text,
            OutputMode::Normal,
            ColorChoice::Auto,
            Verbosity::default(),
        )
    }

    /// Creates a new shell with the given color choice and verbosity.
    pub fn new_with(
        format: OutputFormat,
        mode: OutputMode,
        color: ColorChoice,
        verbosity: Verbosity,
    ) -> Self {
        Self {
            output: ShellOut::Stream {
                stdout: AutoStream::new(std::io::stdout(), color.to_anstream_color_choice()),
                stderr: AutoStream::new(std::io::stderr(), color.to_anstream_color_choice()),
                color_choice: color,
                stderr_tty: std::io::stderr().is_terminal(),
            },
            output_format: format,
            output_mode: mode,
            verbosity,
            needs_clear: AtomicBool::new(false),
        }
    }

    /// Creates a shell that ignores all output.
    pub const fn empty() -> Self {
        Self {
            output: ShellOut::Empty(std::io::empty()),
            output_format: OutputFormat::Text,
            output_mode: OutputMode::Quiet,
            verbosity: 0,
            needs_clear: AtomicBool::new(false),
        }
    }

    /// Creates a shell that captures stdout and stderr into in-memory buffers.
    ///
    /// Intended for tests that want to assert how a piece of code routes output
    /// between stdout and stderr. Use [`Shell::captured_stdout`] and
    /// [`Shell::captured_stderr`] to read the buffers back.
    pub const fn captured() -> Self {
        Self {
            output: ShellOut::Captured { stdout: Vec::new(), stderr: Vec::new() },
            output_format: OutputFormat::Text,
            output_mode: OutputMode::Normal,
            verbosity: 0,
            needs_clear: AtomicBool::new(false),
        }
    }

    /// Returns the captured stdout buffer, if this shell was created via [`Shell::captured`].
    pub fn captured_stdout(&self) -> Option<&[u8]> {
        match &self.output {
            ShellOut::Captured { stdout, .. } => Some(stdout),
            _ => None,
        }
    }

    /// Returns the captured stderr buffer, if this shell was created via [`Shell::captured`].
    pub fn captured_stderr(&self) -> Option<&[u8]> {
        match &self.output {
            ShellOut::Captured { stderr, .. } => Some(stderr),
            _ => None,
        }
    }

    /// Acquire a lock to the global shell.
    ///
    /// Initializes it with the default values if it has not been set yet.
    pub fn get() -> impl DerefMut<Target = Self> + 'static {
        GLOBAL_SHELL.get_or_init(Default::default).lock().unwrap_or_else(PoisonError::into_inner)
    }

    /// Set the global shell.
    ///
    /// # Panics
    ///
    /// Panics if the global shell has already been set.
    #[track_caller]
    pub fn set(self) {
        GLOBAL_SHELL
            .set(Mutex::new(self))
            .unwrap_or_else(|_| panic!("attempted to set global shell twice"))
    }

    /// Sets whether the next print should clear the current line and returns the previous value.
    pub fn set_needs_clear(&self, needs_clear: bool) -> bool {
        self.needs_clear.swap(needs_clear, Ordering::Relaxed)
    }

    /// Returns `true` if the output format is JSON.
    pub fn is_json(&self) -> bool {
        self.output_format.is_json()
    }

    /// Returns `true` if the output format is Markdown.
    pub fn is_markdown(&self) -> bool {
        self.output_format.is_markdown()
    }

    /// Returns `true` if the verbosity level is `Quiet`.
    pub fn is_quiet(&self) -> bool {
        self.output_mode.is_quiet()
    }

    /// Returns `true` if the `needs_clear` flag is set.
    pub fn needs_clear(&self) -> bool {
        self.needs_clear.load(Ordering::Relaxed)
    }

    /// Returns `true` if the `needs_clear` flag is unset.
    pub fn is_cleared(&self) -> bool {
        !self.needs_clear()
    }

    /// Gets the output format of the shell.
    pub const fn output_format(&self) -> OutputFormat {
        self.output_format
    }

    /// Gets the output mode of the shell.
    pub const fn output_mode(&self) -> OutputMode {
        self.output_mode
    }

    /// Gets the verbosity of the shell when [`OutputMode::Normal`] is set.
    pub const fn verbosity(&self) -> Verbosity {
        self.verbosity
    }

    /// Sets the verbosity level.
    pub const fn set_verbosity(&mut self, verbosity: Verbosity) {
        self.verbosity = verbosity;
    }

    /// Sets the output mode.
    pub const fn set_output_mode(&mut self, output_mode: OutputMode) {
        self.output_mode = output_mode;
    }

    /// Gets the current color choice.
    ///
    /// If we are not using a color stream, this will always return `Never`, even if the color
    /// choice has been set to something else.
    pub const fn color_choice(&self) -> ColorChoice {
        match self.output {
            ShellOut::Stream { color_choice, .. } => color_choice,
            ShellOut::Empty(_) | ShellOut::Captured { .. } => ColorChoice::Never,
        }
    }

    /// Returns `true` if stderr is a tty.
    pub const fn is_err_tty(&self) -> bool {
        match self.output {
            ShellOut::Stream { stderr_tty, .. } => stderr_tty,
            ShellOut::Empty(_) | ShellOut::Captured { .. } => false,
        }
    }

    /// Whether `stderr` supports color.
    pub fn err_supports_color(&self) -> bool {
        match &self.output {
            ShellOut::Stream { stderr, .. } => supports_color(stderr.current_choice()),
            ShellOut::Empty(_) | ShellOut::Captured { .. } => false,
        }
    }

    /// Whether `stdout` supports color.
    pub fn out_supports_color(&self) -> bool {
        match &self.output {
            ShellOut::Stream { stdout, .. } => supports_color(stdout.current_choice()),
            ShellOut::Empty(_) | ShellOut::Captured { .. } => false,
        }
    }

    /// Gets a reference to the underlying stdout writer.
    pub fn out(&mut self) -> &mut dyn Write {
        self.maybe_err_erase_line();
        self.output.stdout()
    }

    /// Gets a reference to the underlying stderr writer.
    pub fn err(&mut self) -> &mut dyn Write {
        self.maybe_err_erase_line();
        self.output.stderr()
    }

    /// Erase from cursor to end of line if needed.
    pub fn maybe_err_erase_line(&mut self) {
        if self.err_supports_color() && self.set_needs_clear(false) {
            // This is the "EL - Erase in Line" sequence. It clears from the cursor
            // to the end of line.
            // https://en.wikipedia.org/wiki/ANSI_escape_code#CSI_sequences
            let _ = self.output.stderr().write_all(b"\x1B[K");
        }
    }

    /// Prints a red 'error' message. Use the [`sh_err!`] macro instead.
    /// This will render a message in [ERROR] style with a bold `Error: ` prefix.
    ///
    /// **Note**: will log regardless of the verbosity level.
    pub fn error(&mut self, message: impl fmt::Display) -> Result<()> {
        self.maybe_err_erase_line();
        self.output.message_stderr(&"Error", &ERROR, Some(&message), false)
    }

    /// Prints an amber 'warning' message. Use the [`sh_warn!`] macro instead.
    /// This will render a message in [WARN] style with a bold `Warning: `prefix.
    ///
    /// **Note**: if `verbosity` is set to `Quiet`, this is a no-op.
    pub fn warn(&mut self, message: impl fmt::Display) -> Result<()> {
        match self.output_mode {
            OutputMode::Quiet => Ok(()),
            _ => self.print(&"Warning", &WARN, Some(&message), false),
        }
    }

    /// Write a styled fragment.
    ///
    /// Caller is responsible for deciding whether [`Shell::verbosity`] is affects output.
    pub fn write_stdout(&mut self, fragment: impl fmt::Display, color: &Style) -> Result<()> {
        self.output.write_stdout(fragment, color)
    }

    /// Write a styled fragment with the default color. Use the [`sh_print!`] macro instead.
    ///
    /// **Note**: if `verbosity` is set to `Quiet`, this is a no-op.
    //
    // TODO: stdout is the canonical machine-readable result of a command and should NOT be
    // suppressed by `--quiet` (see `docs/dev/output-channels.md`). Flip this once the major
    // prose `sh_println!` call sites in forge/script have been migrated to `sh_status!`.
    pub fn print_out(&mut self, fragment: impl fmt::Display) -> Result<()> {
        match self.output_mode {
            OutputMode::Quiet => Ok(()),
            _ => self.write_stdout(fragment, &Style::new()),
        }
    }

    /// Write a styled fragment
    ///
    /// Caller is responsible for deciding whether [`Shell::verbosity`] is affects output.
    pub fn write_stderr(&mut self, fragment: impl fmt::Display, color: &Style) -> Result<()> {
        self.output.write_stderr(fragment, color)
    }

    /// Write a styled fragment with the default color. Use the [`sh_eprint!`] macro instead.
    ///
    /// **Note**: if `verbosity` is set to `Quiet`, this is a no-op.
    pub fn print_err(&mut self, fragment: impl fmt::Display) -> Result<()> {
        match self.output_mode {
            OutputMode::Quiet => Ok(()),
            _ => self.write_stderr(fragment, &Style::new()),
        }
    }

    /// Prints a message, where the status will have `color` color, and can be justified. The
    /// messages follows without color.
    fn print(
        &mut self,
        status: &dyn fmt::Display,
        style: &Style,
        message: Option<&dyn fmt::Display>,
        justified: bool,
    ) -> Result<()> {
        match self.output_mode {
            OutputMode::Quiet => Ok(()),
            _ => {
                self.maybe_err_erase_line();
                self.output.message_stderr(status, style, message, justified)
            }
        }
    }
}

impl ShellOut {
    /// Prints out a message with a status to stderr. The status comes first, and is bold plus the
    /// given color. The status can be justified, in which case the max width that will right
    /// align is 12 chars.
    fn message_stderr(
        &mut self,
        status: &dyn fmt::Display,
        style: &Style,
        message: Option<&dyn fmt::Display>,
        justified: bool,
    ) -> Result<()> {
        let buffer = Self::format_message(status, message, style, justified)?;
        self.stderr().write_all(&buffer)?;
        Ok(())
    }

    /// Write a styled fragment
    fn write_stdout(&mut self, fragment: impl fmt::Display, style: &Style) -> Result<()> {
        let mut buffer = Vec::new();
        write!(buffer, "{style}{fragment}{style:#}")?;
        self.stdout().write_all(&buffer)?;
        Ok(())
    }

    /// Write a styled fragment
    fn write_stderr(&mut self, fragment: impl fmt::Display, style: &Style) -> Result<()> {
        let mut buffer = Vec::new();
        write!(buffer, "{style}{fragment}{style:#}")?;
        self.stderr().write_all(&buffer)?;
        Ok(())
    }

    /// Gets stdout as a [`io::Write`](Write) trait object.
    fn stdout(&mut self) -> &mut dyn Write {
        match self {
            Self::Stream { stdout, .. } => stdout,
            Self::Empty(e) => e,
            Self::Captured { stdout, .. } => stdout,
        }
    }

    /// Gets stderr as a [`io::Write`](Write) trait object.
    fn stderr(&mut self) -> &mut dyn Write {
        match self {
            Self::Stream { stderr, .. } => stderr,
            Self::Empty(e) => e,
            Self::Captured { stderr, .. } => stderr,
        }
    }

    /// Formats a message with a status and optional message.
    fn format_message(
        status: &dyn fmt::Display,
        message: Option<&dyn fmt::Display>,
        style: &Style,
        justified: bool,
    ) -> Result<Vec<u8>> {
        let bold = anstyle::Style::new().bold();

        let mut buffer = Vec::new();
        if justified {
            write!(buffer, "{style}{status:>12}{style:#}")?;
        } else {
            write!(buffer, "{style}{status}{style:#}{bold}:{bold:#}")?;
        }
        match message {
            Some(message) => {
                writeln!(buffer, " {message}")?;
            }
            None => write!(buffer, " ")?,
        }

        Ok(buffer)
    }
}

impl ColorChoice {
    /// Converts our color choice to [`anstream`]'s version.
    const fn to_anstream_color_choice(self) -> anstream::ColorChoice {
        match self {
            Self::Always => anstream::ColorChoice::Always,
            Self::Never => anstream::ColorChoice::Never,
            Self::Auto => anstream::ColorChoice::Auto,
        }
    }
}

const fn supports_color(choice: anstream::ColorChoice) -> bool {
    match choice {
        anstream::ColorChoice::Always
        | anstream::ColorChoice::AlwaysAnsi
        | anstream::ColorChoice::Auto => true,
        anstream::ColorChoice::Never => false,
    }
}
