// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

use {
  crate::{
    builder::{
      self,
      FieldDebug,
    },
    prefix::Prefix,
    styles,
    writer::{
      self,
      Writer,
    },
    Bluegum,
    Builder,
    Styles,
  },
  owo_colors::OwoColorize,
  std::{
    cell::RefCell,
    fmt::Write,
    rc::Rc,
    sync::Arc,
  },
  unicode_width::UnicodeWidthStr,
};

const SEPARATOR_PAD_LEN: usize = 6;

/// The main printer for rendering [`Bluegum`] trees to formatted output.
///
/// The `Printer` handles the conversion from the internal tree representation
/// to formatted text with colors, line numbers, and tree structure
/// visualization. It supports customizable styling and can output both colored
/// and plain text.
///
/// # Examples
///
/// ```rust
/// use bluegum::{
///   Bluegum,
///   Printer,
///   Styles,
/// };
///
/// struct Example;
/// impl Bluegum for Example {
///   fn node(&self, b: &mut bluegum::Builder) {
///     b.name("Example");
///   }
/// }
///
/// let example = Example;
/// let mut printer = Printer::default();
/// printer.render(&example);
/// println!("{}", printer.strip_color());
/// ```
pub struct Printer {
  buf:    String,
  styles: styles::Styles,
  title:  Option<String>,
}

impl Printer {
  /// Create a new printer with the specified styles.
  ///
  /// # Examples
  ///
  /// ```rust
  /// use bluegum::{
  ///   Printer,
  ///   Styles,
  /// };
  ///
  /// let styles = Styles::default();
  /// let printer = Printer::new(styles);
  /// ```
  pub fn new(styles: styles::Styles) -> Self {
    Self {
      styles,
      buf: String::new(),
      title: None,
    }
  }

  /// Set an optional title for the output.
  ///
  /// The title will be displayed at the top of the rendered output.
  ///
  /// # Examples
  ///
  /// ```rust
  /// use bluegum::Printer;
  ///
  /// let printer = Printer::default().set_title("My Tree Structure");
  /// ```
  pub fn set_title(mut self, title: &str) -> Self {
    self.title = Some(title.to_string());
    self
  }
}

impl Default for Printer {
  fn default() -> Self {
    Self::new(styles::Styles::default())
  }
}

impl Printer {
  /// Get the rendered output as a string with ANSI color codes removed.
  ///
  /// This is useful when you want plain text output without colors,
  /// such as for logging or writing to files.
  ///
  /// # Examples
  ///
  /// ```rust
  /// use bluegum::{
  ///   Bluegum,
  ///   Printer,
  /// };
  ///
  /// struct Example;
  /// impl Bluegum for Example {
  ///   fn node(&self, b: &mut bluegum::Builder) {
  ///     b.name("Example");
  ///   }
  /// }
  ///
  /// let example = Example;
  /// let mut printer = Printer::default();
  /// printer.render(&example);
  /// let plain_text = printer.strip_color();
  /// ```
  pub fn strip_color(&self) -> String {
    anstream::adapter::strip_str(&self.buf).to_string()
  }

  /// Get the rendered output as a string with ANSI color codes preserved.
  ///
  /// This returns the full colored output suitable for display in
  /// terminals that support ANSI colors.
  ///
  /// # Examples
  ///
  /// ```rust
  /// use bluegum::{
  ///   Bluegum,
  ///   Printer,
  /// };
  ///
  /// struct Example;
  /// impl Bluegum for Example {
  ///   fn node(&self, b: &mut bluegum::Builder) {
  ///     b.name("Example");
  ///   }
  /// }
  ///
  /// let example = Example;
  /// let mut printer = Printer::default();
  /// printer.render(&example);
  /// let colored_text = printer.with_color();
  /// ```
  pub fn with_color(&self) -> &str {
    &self.buf
  }

  pub(crate) fn buffer(&self) -> &str {
    &self.buf
  }

  /// Print the rendered output to stdout with colors.
  ///
  /// This is a convenience method that prints the colored output
  /// directly to standard output.
  ///
  /// # Examples
  ///
  /// ```rust
  /// use bluegum::{
  ///   Bluegum,
  ///   Printer,
  /// };
  ///
  /// struct Example;
  /// impl Bluegum for Example {
  ///   fn node(&self, b: &mut bluegum::Builder) {
  ///     b.name("Example");
  ///   }
  /// }
  ///
  /// let example = Example;
  /// let mut printer = Printer::default();
  /// printer.render(&example);
  /// printer.to_stdout();
  /// ```
  pub fn to_stdout(&self) {
    println!("{}", &self.buf);
  }
}

impl Printer {
  /// Render a [`Bluegum`] tree into the printer's internal buffer.
  ///
  /// This method takes a root node implementing [`Bluegum`] and renders
  /// it into the printer's internal buffer. The rendered output can then
  /// be retrieved using [`strip_color`], [`with_color`], or printed
  /// using [`to_stdout`].
  ///
  /// # Examples
  ///
  /// ```rust
  /// use bluegum::{
  ///   Bluegum,
  ///   Printer,
  /// };
  ///
  /// struct Example;
  /// impl Bluegum for Example {
  ///   fn node(&self, b: &mut bluegum::Builder) {
  ///     b.name("Example");
  ///   }
  /// }
  ///
  /// let example = Example;
  /// let mut printer = Printer::default();
  /// printer.render(&example);
  /// ```
  ///
  /// [`strip_color`]: Printer::strip_color
  /// [`with_color`]: Printer::with_color
  /// [`to_stdout`]: Printer::to_stdout
  pub fn render(&mut self, root: &impl Bluegum) -> &mut Self {
    let mut p = Prefix::new(&self.styles);
    let mut w = Writer::new();
    let mut inner = PrinterInner {
      prefix:              &mut p,
      writer:              &mut w,
      verbose:             self.styles.verbose,
      styles:              &self.styles,
      title:               self.title.clone(),
      last_line_was_blank: false,
    };

    let mut b = Builder::new();

    Bluegum::node(root, &mut b);

    inner.render(&b, &mut self.buf);

    self
  }

  pub fn render_builder(&mut self, builder: Builder) -> &mut Self {
    let mut p = Prefix::new(&self.styles);
    let mut w = Writer::new();
    let mut inner = PrinterInner {
      prefix:              &mut p,
      writer:              &mut w,
      verbose:             self.styles.verbose,
      styles:              &self.styles,
      title:               self.title.clone(),
      last_line_was_blank: false,
    };

    inner.render(&builder, &mut self.buf);

    self
  }

  pub fn render_builder_with<F>(&mut self, callback: F) -> &mut Self
  where
    F: FnOnce(&mut Builder),
  {
    let mut p = Prefix::new(&self.styles);
    let mut w = Writer::new();
    let mut inner = PrinterInner {
      prefix:              &mut p,
      writer:              &mut w,
      verbose:             self.styles.verbose,
      styles:              &self.styles,
      title:               self.title.clone(),
      last_line_was_blank: false,
    };

    let mut b = Builder::new();

    callback(&mut b);

    inner.render(&b, &mut self.buf);

    self
  }
}

pub(crate) struct PrinterInner<'a> {
  pub(crate) prefix:   &'a mut Prefix<'a>,
  writer:              &'a mut Writer,
  verbose:             bool,
  styles:              &'a styles::Styles,
  title:               Option<String>,
  last_line_was_blank: bool,
}

impl<'a, 'b: 'a> PrinterInner<'a> {
  fn render(&'a mut self, b: &'a Builder, buf: &'b mut String) {
    // render the tree starting from the root node
    self.render_node(b, 0, false, 1);

    self.writer.trim_line_endings();

    // commit the written lines to the buffer
    self.write_to_buf(buf);
  }

  fn render_node(
    &mut self,
    b: &Builder,
    idx: usize,
    hide_index: bool,
    sibling_len: usize,
  ) -> std::fmt::Result {
    // If this is a Vec, render the Vec without a blank node in between
    if b.name.name.is_empty()
      && b.nodes.is_empty()
      && b.fields.is_empty()
      && !b.unlabled_nodes.is_empty()
    {
      // render the unlabled nodes
      let nodes_len = b.unlabled_nodes.len();
      for (i, node) in b.unlabled_nodes.iter().enumerate() {
        self.render_node(node, i, false, nodes_len)?;
      }
      return Ok(());
    }

    // write the name
    self.on_node_enter(
      b.name.clone(),
      Some(&b.fields),
      idx,
      hide_index,
      sibling_len,
      !b.nodes.is_empty() || !b.unlabled_nodes.is_empty(),
      !b.fields.is_empty(),
      idx == sibling_len - 1,
    );

    // unlabled nodes are printed after fields, but inline as if they are fields
    if !b.unlabled_nodes.is_empty() {
      let nodes_len = b.unlabled_nodes.len();
      for (i, node) in b.unlabled_nodes.iter().enumerate() {
        self.render_node(node, i, true, nodes_len);
      }
    }

    if !b.nodes.is_empty() {
      let nodes_len = b.nodes.len();
      for (i, (name, nodes)) in b.nodes.iter().enumerate() {
        'render_nodes: {
          self.on_nodes_enter(
            name,
            i,
            nodes.len(),
            nodes_len,
            b.total_nodes_len,
          );

          'render_node: {
            for (i, node) in nodes.iter().enumerate() {
              self.render_node(node, i, false, nodes.len());
            }
            self.on_nodes_exit();
          }
        }
      }
    }

    self.on_node_exit()
  }

  fn print_title(
    &mut self,
    buf: &mut String,
    line_column_width: usize,
    longest_line: usize,
  ) {
    if let Some(title) = &self.title {
      let title = title.trim();
      writeln!(
        buf,
        "{title_l_pad}{title_l}{title}{title_r}{title_r_pad}",
        title_l_pad = "=".repeat(line_column_width).purple().dimmed(),
        title_l = "[ ".purple().dimmed(),
        title = title.white(),
        title_r = " ]".purple().dimmed(),
        title_r_pad = "="
          .repeat(100.max({
            // longest_line
            // +4 accounts for the "[ " and " ]"
            let title_len = title.len() + line_column_width + 4;
            // longest_line
            if longest_line > title_len {
              longest_line - title_len
            } else {
              longest_line
            }
          }))
          .purple()
          .dimmed()
      );
    }
  }

  fn write_to_buf(&mut self, buf: &'a mut String) -> std::fmt::Result {
    let line_count = self.writer.lines.len();
    let line_count_width = line_count.to_string().width();

    let hide_lines = self.styles.hide_lines;
    let hide_line_numbers = self.styles.hide_line_numbers;

    // first measure the column width of the line segments
    let widths: Vec<(usize, usize, usize, usize)> = self
      .writer
      .lines
      .iter()
      .enumerate()
      .map(
        |(
          i,
          writer::Line {
            prefix,
            line,
            alt,
            debug,
          },
        )| {
          let prefix_width =
            anstream::adapter::strip_str(prefix).to_string().width();
          let width = anstream::adapter::strip_str(line).to_string().width();

          let alt_width = anstream::adapter::strip_str(alt).to_string().width();
          let debug_width =
            anstream::adapter::strip_str(debug).to_string().width();

          (prefix_width, width, alt_width, debug_width)
        },
      )
      .collect();

    // then calculate the longest line
    // we use this to align the debug output
    // TODO: only do this if printing debug (gold-kym)
    let longest_line = widths
      .iter()
      .map(|(p, width, ..)| *p + *width)
      .max()
      .unwrap_or(0);

    let longest_prefix =
      widths.iter().map(|(prefix, ..)| *prefix).max().unwrap_or(0);

    let longest_alt =
      widths.iter().map(|(_, _, alt, _)| *alt).max().unwrap_or(0);

    let longest_line_alt = (longest_line).min(self.styles.max_debug_pad);

    let longest_line_total = widths
      .iter()
      .map(|(prefix_width, width, alt_width, debug_width)| {
        *prefix_width + *width + *alt_width + *debug_width
      })
      .max()
      .unwrap_or(0);

    self.print_title(buf, line_count_width, longest_line_total);

    // write the lines in the writer to the buffer
    for (
      i,
      writer::Line {
        prefix,
        line,
        alt,
        debug,
      },
    ) in self.writer.lines.iter().enumerate()
    {
      let mut line_buf = String::new();

      let Some((prefix_width, width, alt_width, debug_width)) = widths.get(i)
      else {
        panic!("widths.get({}) failed", i);
      };

      let alt_pad_count: usize = {
        let leading_width = (*prefix_width + *width);
        longest_line_alt.saturating_sub(leading_width)
      };
      let max_debug_pad_count: usize = if *width > longest_line_alt {
        0
      } else {
        longest_line_alt.min(self.styles.max_debug_pad) - *width
      };

      // write the line number
      if !hide_line_numbers {
        write!(
          line_buf,
          "{}",
          format!("{i:0>line_count_width$}│").style(self.styles.line_numbers)
        )?;

        // reset any coloring
        write!(line_buf, "{}", "".on_default_color().default_color());
      }

      write!(line_buf, "{prefix}{line}")?;
      // reset any coloring
      write!(line_buf, "{}", "".on_default_color().default_color());

      let debug_pad_count = {
        if *alt_width > 0 {
          write!(
            line_buf,
            "{alt_pad}{alt}",
            alt_pad = " ".repeat(alt_pad_count + SEPARATOR_PAD_LEN),
            alt = alt,
          )?;

          // reset any coloring
          write!(line_buf, "{}", "".on_default_color().default_color());
          longest_alt - alt_width
        } else {
          alt_pad_count + longest_alt + SEPARATOR_PAD_LEN
        }
      };

      if self.styles.show_debug && !debug.is_empty() {
        write!(
          line_buf,
          "{debug_pad}{debug}",
          debug_pad = " ".repeat(debug_pad_count + SEPARATOR_PAD_LEN),
          debug = debug,
        )?;

        // reset any coloring
        write!(line_buf, "{}", "".on_default_color().default_color());
      }

      writeln!(buf, "{}", line_buf.trim())?;
    }

    // write a single new line to indicate the end
    writeln!(buf)?;

    Ok(())
  }
}

/// Methods for printing the tree
impl PrinterInner<'_> {
  #[allow(clippy::too_many_arguments)]
  pub(crate) fn on_node_enter(
    &mut self,
    name: builder::BuilderName,
    fields: Option<&builder::Fields>,
    index: usize,
    hide_index: bool,
    len: usize,
    has_node: bool,
    has_fields: bool,
    last_node: bool,
  ) -> std::fmt::Result {
    {
      // reset the last line was blank flag
      self.last_line_was_blank = false;

      if has_node {
        self.prefix.push_stack_node(has_fields, last_node);
      } else {
        self.prefix.push_stack_leaf(has_fields, last_node);
      }

      self.prefix.render_prefix_stack(self.writer, self.styles)?;

      if !hide_index && len > 1 {
        // Only print the index number, when theres enough nodes that it becomes
        // helpful. Ohterwise there's a whole screen of 0's for things like
        // a binary tree.
        self.writer.write_line(
          format!(
            "{icon} {idx} {name}",
            icon = "▣".style(self.styles.icon_node),
            name = name.name.style(self.styles.type_name),
            idx = format!("{index}:").style(self.styles.type_name_trivia),
          )
          .to_string(),
        )?;
      } else {
        self.writer.write_line(
          format!(
            "{icon} {name} ",
            icon = "▣".style(self.styles.icon_node),
            name = name.name.style(self.styles.type_name),
          )
          .to_string(),
        )?;
      }

      if let Some(alt) = name.alt {
        self.print_field_alt(&alt)?;
      }

      if !name.debug.is_empty() {
        self.print_field_debug(&name.debug)?;
      }
    }

    if let Some(fields) = fields {
      self.print_fields(fields, has_node, last_node)
    } else {
      Ok(())
    }
  }

  pub(crate) fn on_node_exit(&mut self) -> std::fmt::Result {
    if self.styles.new_line_after_node && !self.last_line_was_blank {
      self.print_blank_line()?;
    }

    self.prefix.pop_stack();

    Ok(())
  }

  /// Call before printing the nodes of a given Node or Fragment.
  pub(crate) fn on_nodes_enter(
    &mut self,
    name: &str,
    index: usize,
    our_len: usize,
    len: usize,
    nodes_len: usize,
  ) -> std::fmt::Result {
    // reset the last line was blank flag
    self.last_line_was_blank = false;

    let first_node = index == 0;
    let last_node = index == len - 1;

    self
      .prefix
      .push_stack_nodes(first_node, last_node, nodes_len != 0);

    self.prefix.render_prefix_stack(self.writer, self.styles)?;

    match our_len {
      | 1 => {
        self.writer.write_line(
          format!(
            "{icon} {name}",
            icon = "○".style(self.styles.nodes_label_trivia),
            name = name.style(self.styles.nodes_label),
          )
          .to_string(),
        )?;
      },
      | _ => {
        self.writer.write_line(
          format!(
            "{icon} {name} {len}",
            icon = "○".style(self.styles.nodes_label_trivia),
            name = name.style(self.styles.nodes_label),
            len = format!("[len={our_len}]").style(self.styles.trivia_len_tag),
          )
          .to_string(),
        )?;
      },
    }

    'write_no_nodes: {
      if nodes_len == 0 {
        self.prefix.push_stack_leaf(false, true);

        self.prefix.render_prefix_stack(self.writer, self.styles)?;

        self.writer.write_line(" no nodes".dimmed().to_string())?;

        self.on_node_exit()?;
      }
    }

    Ok(())
  }

  /// Called when all nodes of a given Node or Fragment have been printed.
  pub(crate) fn on_nodes_exit(&mut self) {
    self.prefix.pop_stack();
  }
}

impl PrinterInner<'_> {
  fn print_blank_line(&mut self) -> std::fmt::Result {
    self.prefix.push_stack_blank();
    self.prefix.render_prefix_stack(self.writer, self.styles)?;
    self.prefix.pop_stack();

    self.last_line_was_blank = true;

    Ok(())
  }

  fn print_fields(
    &mut self,
    fields: &super::builder::Fields,
    has_node: bool,
    last_node: bool,
  ) -> std::fmt::Result {
    if self.verbose {
      let longest_key = fields
        .iter()
        .filter_map(|f| {
          match f {
            | builder::Field::KeyValue { name, .. } => Some(name.len()),
            | _ => None,
          }
        })
        .max()
        .unwrap_or(0);
      let len = fields.len();

      for (i, field) in fields.iter().enumerate() {
        match field {
          | builder::Field::KeyValue {
            name,
            value,
            alt,
            debug,
          } => {
            self.prefix.push_stack_field(
              has_node,
              i == 0,
              i == len - 1,
              last_node,
            );

            self.prefix.render_prefix_stack(self.writer, self.styles)?;

            self.writer.write_line(format!(
              "{name}{equal:<width$}{value}",
              name = name.style(self.styles.field_name),
              // +1 for zero-indexing, +1 for the colon
              width = longest_key - name.len() + 2,
              equal = ": ".style(self.styles.field_name),
              value = value.style(self.styles.field_value),
            ))?;

            if let Some(alt) = alt.as_ref() {
              self.print_field_alt(alt)?
            };

            self.print_field_debug(debug)?;
            // pop the field off the prefix stack
            self.prefix.pop_stack();
          },
          | builder::Field::Divider => {
            self.prefix.push_stack_divider();

            self.prefix.render_prefix_stack(self.writer, self.styles)?;

            self.writer.write_line(format!(
              "{}",
              "-".repeat(longest_key + 1).style(self.styles.field_name)
            ))?;

            self.prefix.pop_stack();
          },
        }
      }
    } else {
      for (i, field) in fields.iter().enumerate() {
        match field {
          | builder::Field::KeyValue {
            name,
            value,
            alt,
            debug,
          } => {
            self.writer.write_line(format!(
              "{name}{equal}{quote}{value}{quote}{comma}",
              name = name.style(self.styles.field_name),
              equal = "=".style(self.styles.trivia),
              quote = '`'.style(self.styles.trivia),
              value = value.style(self.styles.field_value),
              comma = {
                // if this is the last field, don't print a comma
                if i + 1 == fields.len() {
                  ""
                } else {
                  ", "
                }
              }
              .style(self.styles.trivia)
            ))?;

            if let Some(alt) = alt.as_ref() {
              self.print_field_alt(alt)?
            };

            self.print_field_debug(debug)?;
          },
          | builder::Field::Divider => {
            self.writer.write_line(format!(
              "{}",
              "-".repeat(12).style(self.styles.field_name)
            ))?;
          },
        }
      }
    }

    Ok(())
  }

  fn print_field_alt(&mut self, alt: &str) -> std::fmt::Result {
    self.writer.write_alt(format!(
      "{value}",
      value = alt.style(self.styles.alt_value),
    ))?;

    Ok(())
  }

  fn print_field_debug(
    &mut self,
    debug: &super::builder::FieldDebug,
  ) -> std::fmt::Result {
    for (name, value) in debug {
      self.writer.write_debug(format!(
        " {l}{name}{c}{value}{r}",
        l = "[".purple().dimmed(),
        name = name.style(self.styles.field_name),
        c = ": ".purple().dimmed(),
        value = value.style(self.styles.field_value),
        r = "]".purple().dimmed(),
      ))?;
    }

    Ok(())
  }
}
