use std::fmt::Display;

use crate::{lexer::SourceLocation, writeln_indented};

use super::{
    intern::{self, InternPoolWrapper as InternPool},
    visitor::AstVisitorTrait,
    Ast, Children, ComptimeCache, Index, Tag, TypeCache,
};

struct AstRenderVisitor<'a, W: core::fmt::Write> {
    syms: &'a crate::symbol_table::syms2::Symbols,
    ip: &'a InternPool,
    cache: TypeCache,
    comptime_cache: ComptimeCache,
    scopes: Vec<Index>,
    w: &'a mut W,
}

impl<'a, W: core::fmt::Write> AstVisitorTrait<&'a Ast> for AstRenderVisitor<'_, W> {
    type Error = core::fmt::Error;

    type Value = ();

    const UNIMPL: Self::Error = core::fmt::Error;

    fn visit_any(&mut self, ast: &'a Ast, idx: Index) -> Result<Self::Value, Self::Error> {
        let display = NodeDisplay::new(ast, self.ip, idx).with_indent(self.scopes.len() as u32 * 2);

        match display.tag {
            Tag::File
            | Tag::FunctionDecl
            | Tag::GlobalDecl
            | Tag::Block
            | Tag::BlockTrailingExpr => {
                self.scopes.push(idx);
            }
            _ => {}
        }

        write!(self.w, "{display}")?;

        for child in display.children.0 {
            self.visit_any(ast, child)?;
        }

        match display.tag {
            Tag::File
            | Tag::FunctionDecl
            | Tag::GlobalDecl
            | Tag::Block
            | Tag::BlockTrailingExpr => {
                self.scopes.pop();
            }
            _ => {}
        }

        Ok(())
    }
}

pub struct AstRenderer<'a> {
    ast: &'a Ast,
    #[allow(dead_code)]
    syms: &'a crate::symbol_table::syms2::Symbols,
    ip: &'a InternPool,
    cache: TypeCache,
    comptime_cache: ComptimeCache,
}

pub struct NodeDisplay {
    node: Index,
    tag: Tag,
    loc: SourceLocation,
    children: Children,
    is_comptime: bool,
    ty: intern::Index,
    indent: u32,
}

impl NodeDisplay {
    pub fn new(ast: &Ast, ip: &InternPool, node: Index) -> NodeDisplay {
        let tag = ast.tags[node];
        let loc = ast.source_locs[node];

        let children = Children(ast.get_node_children(node));
        let ty = ast.get_type_of_node(ip, &mut TypeCache::new(), node);

        let is_comptime = ast.is_node_comptime_evaluable(&mut ComptimeCache::default(), node);

        Self {
            node,
            tag,
            loc,
            children,
            is_comptime,
            ty,
            indent: 0,
        }
    }
    fn with_indent(self, indent: u32) -> Self {
        Self { indent, ..self }
    }
}

impl Display for NodeDisplay {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let Self {
            node,
            tag,
            loc,
            children,
            is_comptime,
            ty,
            indent,
        } = self;
        writeln_indented!(
            *indent,
            f,
            "{node} {}({ty}) = ({loc}) {tag:?} {children}",
            if *is_comptime { "CONST " } else { "" }
        )?;

        Ok(())
    }
}

impl<'a> AstRenderer<'a> {
    pub fn new(
        ast: &'a Ast,
        ip: &'a InternPool,
        syms: &'a crate::symbol_table::syms2::Symbols,
    ) -> Self {
        Self {
            ast,
            syms,
            ip,
            cache: TypeCache::new(),
            comptime_cache: ComptimeCache::default(),
        }
    }

    pub(crate) fn render<W: core::fmt::Write>(&mut self, w: &mut W) -> core::fmt::Result {
        let mut visitor = AstRenderVisitor {
            syms: self.syms,
            ip: self.ip,
            cache: TypeCache::new(),
            comptime_cache: ComptimeCache::default(),
            scopes: Vec::new(),
            w,
        };

        for idx in self.ast.get_root_file_indices() {
            visitor.visit_any(self.ast, idx)?;
        }

        Ok(())
    }
}