use proc_macro2::{TokenStream, TokenTree, Delimiter};

// ── Public surface ────────────────────────────────────────────

pub struct TemplateAst {
    pub nodes: Vec<TemplateNode>,
}

/// Parse a proc_macro2::TokenStream containing HTML-ish template markup.
///
/// Recognised constructs:
/// - `{{ expr }}`          → escaped interpolation  (→ push_escaped)
/// - `{{{ expr }}}`        → raw interpolation      (→ push_str of .to_string())
/// - `<tag @if="cond">…</tag>` → conditional block
/// - `<tag @for="x in y">…</tag>` → loop block
/// - `<tag @else>…</tag>`  → else branch
/// - `<Component …/>`      → component rendering call (PascalCase tag)
/// - `<tag attr="val">…</tag>` → HTML element
/// - everything else       → emitted as literal text
pub fn parse_template(tokens: TokenStream) -> TemplateAst {
    let flat = flatten(tokens);
    let (nodes, _) = parse_nodes(&flat, 0, None);
    TemplateAst { nodes }
}

// ── AST types ─────────────────────────────────────────────────

#[derive(Debug, Clone)]
pub enum TemplateNode {
    Text(String),
    Escape { expr: String },
    Raw    { expr: String },
    Element(Box<Element>),
}

#[derive(Debug, Clone)]
pub struct Element {
    pub tag: String,
    pub is_component: bool,
    pub attrs: Vec<(String, String)>,
    pub directives: Vec<Directive>,
    pub children: Vec<TemplateNode>,
    pub self_closing: bool,
}

#[derive(Debug, Clone)]
pub enum Directive {
    If(String),
    For { var: String, iter: String },
    Else,
}

// ── Flat token representation ─────────────────────────────────

#[derive(Debug, Clone)]
enum Tok {
    Lt,
    Gt,
    Slash,
    At,
    Eq,
    Comma,
    Ident(String),
    Lit(String),    // string literal — includes the outer quotes
    BraceGroup(Vec<Tok>),  // { ... }
    Other(String),
}

fn flatten(ts: TokenStream) -> Vec<Tok> {
    ts.into_iter().flat_map(flatten_tt).collect()
}

fn flatten_tt(tt: TokenTree) -> Vec<Tok> {
    match tt {
        TokenTree::Ident(i) => vec![Tok::Ident(i.to_string())],
        TokenTree::Literal(l) => vec![Tok::Lit(l.to_string())],
        TokenTree::Punct(p) => vec![match p.as_char() {
            '<' => Tok::Lt,
            '>' => Tok::Gt,
            '/' => Tok::Slash,
            '@' => Tok::At,
            '=' => Tok::Eq,
            ',' => Tok::Comma,
            c   => Tok::Other(c.to_string()),
        }],
        TokenTree::Group(g) if g.delimiter() == Delimiter::Brace => {
            vec![Tok::BraceGroup(flatten(g.stream()))]
        }
        TokenTree::Group(g) => {
            // Other delimiters (parens, brackets) — flatten inline with pseudo-tokens
            let (open, close) = match g.delimiter() {
                Delimiter::Parenthesis => ("(", ")"),
                Delimiter::Bracket     => ("[", "]"),
                _                      => ("{", "}"),
            };
            let mut v = vec![Tok::Other(open.to_string())];
            v.extend(flatten(g.stream()));
            v.push(Tok::Other(close.to_string()));
            v
        }
    }
}

// ── Parser ────────────────────────────────────────────────────

/// Returns (nodes, next_pos).
/// `stop_at` — stop when we encounter `</stop_at>`.
fn parse_nodes(toks: &[Tok], mut pos: usize, stop_at: Option<&str>) -> (Vec<TemplateNode>, usize) {
    let mut nodes: Vec<TemplateNode> = Vec::new();

    while pos < toks.len() {
        // ── {{ expr }} / {{{ expr }}} interpolation ──────────
        if let Tok::BraceGroup(inner) = &toks[pos] {
            // Check for nested brace group → {{ }}
            if inner.len() == 1 {
                if let Tok::BraceGroup(inner2) = &inner[0] {
                    // {{{ }}} → raw
                    if inner2.len() == 1 {
                        if let Tok::BraceGroup(expr_toks) = &inner2[0] {
                            let expr = toks_to_str(expr_toks);
                            nodes.push(TemplateNode::Raw { expr });
                            pos += 1;
                            continue;
                        }
                    }
                    // {{ }} → escaped
                    let expr = toks_to_str(inner2);
                    nodes.push(TemplateNode::Escape { expr });
                    pos += 1;
                    continue;
                }
            }
            // Other brace group → treat as text `{ ... }`
            let text = format!("{{ {} }}", toks_to_str(inner));
            push_text(&mut nodes, &text);
            pos += 1;
            continue;
        }

        // ── HTML tag ─────────────────────────────────────────
        if let Tok::Lt = &toks[pos] {
            // closing tag?
            if pos + 1 < toks.len() {
                if let Tok::Slash = &toks[pos + 1] {
                    // </tag> — check if it matches our stop_at
                    if let Some(stop) = stop_at {
                        if let Some(Tok::Ident(tag)) = toks.get(pos + 2) {
                            if tag == stop {
                                // consume </tag>
                                let end = find_gt(toks, pos + 3).unwrap_or(pos + 3);
                                pos = end + 1;
                                return (nodes, pos);
                            }
                        }
                    }
                    // Unknown closing tag → emit as text
                    let end = find_gt(toks, pos + 1).unwrap_or(pos + 1);
                    let text = reconstruct_tag(toks, pos, end);
                    push_text(&mut nodes, &text);
                    pos = end + 1;
                    continue;
                }
            }

            // Opening / self-closing tag
            if let Some((elem, next_pos)) = parse_element(toks, pos) {
                let tag_name = elem.tag.clone();
                let is_self = elem.self_closing;
                let mut elem = elem;

                if !is_self {
                    // parse children until </tag>
                    let (children, after) = parse_nodes(toks, next_pos, Some(&tag_name));
                    elem.children = children;
                    pos = after;
                } else {
                    pos = next_pos;
                }

                nodes.push(TemplateNode::Element(Box::new(elem)));
                continue;
            }

            // Couldn't parse as element — treat `<` as text
            push_text(&mut nodes, "<");
            pos += 1;
            continue;
        }

        // ── Anything else → text ─────────────────────────────
        let text = tok_to_str(&toks[pos]);
        push_text(&mut nodes, &text);
        pos += 1;
    }

    (nodes, pos)
}

/// Parse an element starting at `<`.  Returns (Element, pos_after_closing_`>`).
fn parse_element(toks: &[Tok], start: usize) -> Option<(Element, usize)> {
    debug_assert!(matches!(toks.get(start), Some(Tok::Lt)));
    let mut pos = start + 1;

    // tag name
    let tag = match toks.get(pos) {
        Some(Tok::Ident(t)) => t.clone(),
        _ => return None,
    };
    pos += 1;

    let is_component = tag.chars().next().map(|c| c.is_uppercase()).unwrap_or(false);

    // attributes + directives
    let mut attrs: Vec<(String, String)> = Vec::new();
    let mut directives: Vec<Directive> = Vec::new();
    let mut self_closing = false;

    loop {
        match toks.get(pos) {
            None | Some(Tok::Gt) => {
                if matches!(toks.get(pos), Some(Tok::Gt)) { pos += 1; }
                break;
            }
            Some(Tok::Slash) => {
                // /> self-closing
                self_closing = true;
                if matches!(toks.get(pos + 1), Some(Tok::Gt)) { pos += 2; } else { pos += 1; }
                break;
            }
            Some(Tok::At) => {
                // directive: @name or @name="value"
                pos += 1;
                let name = match toks.get(pos) {
                    Some(Tok::Ident(n)) => n.clone(),
                    _ => continue,
                };
                pos += 1;
                if matches!(toks.get(pos), Some(Tok::Eq)) {
                    pos += 1;
                    let val = match toks.get(pos) {
                        Some(Tok::Lit(s)) => unquote(s),
                        Some(other) => tok_to_str(other),
                        None => String::new(),
                    };
                    pos += 1;
                    match name.as_str() {
                        "if"   => directives.push(Directive::If(val)),
                        "for"  => {
                            let (var, iter) = split_for_expr(&val);
                            directives.push(Directive::For { var, iter });
                        }
                        _other => attrs.push((format!("@{name}"), val)),
                    }
                } else {
                    match name.as_str() {
                        "else" => directives.push(Directive::Else),
                        _ => {}
                    }
                }
            }
            Some(Tok::Ident(name)) => {
                // regular attribute: name or name="value"
                let name = name.clone();
                pos += 1;
                if matches!(toks.get(pos), Some(Tok::Eq)) {
                    pos += 1;
                    let val = match toks.get(pos) {
                        Some(Tok::Lit(s)) => unquote(s),
                        Some(other) => tok_to_str(other),
                        None => String::new(),
                    };
                    pos += 1;
                    attrs.push((name, val));
                } else {
                    attrs.push((name, String::new()));
                }
            }
            _ => { pos += 1; }
        }
    }

    Some((Element { tag, is_component, attrs, directives, children: vec![], self_closing }, pos))
}

// ── Codegen ───────────────────────────────────────────────────

impl TemplateAst {
    /// Returns the loop variable names used in `@for` directives.
    pub fn for_bindings(&self) -> Vec<String> {
        let mut out = Vec::new();
        collect_for_bindings(&self.nodes, &mut out);
        out
    }

    /// Produce the body of a `render()` function as a TokenStream:
    /// ```rust
    /// let mut __buf = misses_rt::Buffer::new();
    /// // ... push_str / push_escaped / if / for ...
    /// misses_rt::Markup::from_buffer(__buf)
    /// ```
    pub fn codegen(&self) -> TokenStream {
        let mut code = String::new();
        code.push_str("let mut __buf = misses_rt :: Buffer :: new () ;\n");
        for node in &self.nodes {
            codegen_node(&mut code, node, 0);
        }
        code.push_str("misses_rt :: Markup :: from_buffer ( __buf )\n");
        code.parse()
            .unwrap_or_else(|_| TokenStream::new())
    }
}

fn collect_for_bindings(nodes: &[TemplateNode], out: &mut Vec<String>) {
    for node in nodes {
        if let TemplateNode::Element(e) = node {
            for d in &e.directives {
                if let Directive::For { var, .. } = d {
                    out.push(var.clone());
                }
            }
            collect_for_bindings(&e.children, out);
        }
    }
}

fn codegen_node(code: &mut String, node: &TemplateNode, depth: usize) {
    let indent = "    ".repeat(depth);
    match node {
        TemplateNode::Text(s) => {
            if !s.trim().is_empty() {
                let escaped = s.replace('\\', "\\\\").replace('"', "\\\"");
                code.push_str(&format!("{indent}__buf . push_str ( \"{}\" ) ;\n", escaped));
            }
        }
        TemplateNode::Escape { expr } => {
            code.push_str(&format!(
                "{indent}__buf . push_escaped ( & ({expr}) . to_string () ) ;\n"
            ));
        }
        TemplateNode::Raw { expr } => {
            code.push_str(&format!(
                "{indent}__buf . push_str ( & ({expr}) . to_string () ) ;\n"
            ));
        }
        TemplateNode::Element(elem) => codegen_element(code, elem, depth),
    }
}

fn codegen_element(code: &mut String, elem: &Element, depth: usize) {
    let indent = "    ".repeat(depth);

    // Determine wrapping control flow from directives
    for dir in &elem.directives {
        match dir {
            Directive::If(cond) => {
                code.push_str(&format!("{indent}if {cond} {{\n"));
            }
            Directive::For { var, iter } => {
                code.push_str(&format!("{indent}for {var} in {iter} {{\n"));
            }
            Directive::Else => {
                code.push_str(&format!("{indent}}} else {{\n"));
            }
        }
    }

    let inner_depth = if elem.directives.is_empty() { depth } else { depth + 1 };
    let inner_indent = "    ".repeat(inner_depth);

    if elem.is_component {
        // Component invocation
        let tag = &elem.tag;
        let mut prop_args = String::new();
        for (k, v) in &elem.attrs {
            if !k.starts_with('@') {
                prop_args.push_str(&format!("{k} : ({v}) . into () , "));
            }
        }
        code.push_str(&format!(
            "{inner_indent}__buf . push_str ( & {tag} :: new ( {prop_args}) . render () . into_string () ) ;\n"
        ));
    } else {
        // HTML element: open tag
        let open_tag = build_open_tag(elem);
        let escaped = open_tag.replace('\\', "\\\\").replace('"', "\\\"");
        code.push_str(&format!("{inner_indent}__buf . push_str ( \"{escaped}\" ) ;\n"));

        // Children
        for child in &elem.children {
            codegen_node(code, child, inner_depth);
        }

        // Close tag
        if !elem.self_closing {
            let close = format!("</{}>", elem.tag);
            code.push_str(&format!("{inner_indent}__buf . push_str ( \"{}\" ) ;\n", close));
        }
    }

    if !elem.directives.is_empty() {
        code.push_str(&format!("{indent}}}\n"));
    }
}

fn build_open_tag(elem: &Element) -> String {
    let mut s = format!("<{}", elem.tag);
    for (name, val) in &elem.attrs {
        if name.starts_with('@') { continue; }
        if val.is_empty() {
            s.push_str(&format!(" {name}"));
        } else {
            s.push_str(&format!(" {name}=\"{val}\""));
        }
    }
    if elem.self_closing {
        s.push_str(" />");
    } else {
        s.push('>');
    }
    s
}

// ── Helpers ───────────────────────────────────────────────────

fn toks_to_str(toks: &[Tok]) -> String {
    toks.iter().map(tok_to_str).collect::<Vec<_>>().join(" ")
}

fn tok_to_str(tok: &Tok) -> String {
    match tok {
        Tok::Lt          => "<".into(),
        Tok::Gt          => ">".into(),
        Tok::Slash       => "/".into(),
        Tok::At          => "@".into(),
        Tok::Eq          => "=".into(),
        Tok::Comma       => ",".into(),
        Tok::Ident(s)    => s.clone(),
        Tok::Lit(s)      => s.clone(),
        Tok::BraceGroup(inner) => format!("{{ {} }}", toks_to_str(inner)),
        Tok::Other(s)    => s.clone(),
    }
}

/// Remove surrounding `"..."` from a string literal token.
fn unquote(s: &str) -> String {
    if s.starts_with('"') && s.ends_with('"') && s.len() >= 2 {
        s[1..s.len()-1].to_string()
    } else {
        s.to_string()
    }
}

/// Split `"item in &self.items"` → ("item", "&self.items")
fn split_for_expr(s: &str) -> (String, String) {
    if let Some(idx) = s.find(" in ") {
        let var = s[..idx].trim().to_string();
        let iter = s[idx + 4..].trim().to_string();
        (var, iter)
    } else {
        (s.to_string(), String::new())
    }
}

fn find_gt(toks: &[Tok], from: usize) -> Option<usize> {
    for i in from..toks.len() {
        if matches!(toks[i], Tok::Gt) { return Some(i); }
    }
    None
}

fn reconstruct_tag(toks: &[Tok], start: usize, end: usize) -> String {
    toks[start..=end].iter().map(tok_to_str).collect::<Vec<_>>().join("")
}

fn push_text(nodes: &mut Vec<TemplateNode>, text: &str) {
    // Normalise whitespace runs to a single space when appending to existing Text
    let text = text.replace('\n', " ").replace('\r', "");
    if text.trim().is_empty() {
        return;
    }
    match nodes.last_mut() {
        Some(TemplateNode::Text(prev)) => prev.push_str(&text),
        _ => nodes.push(TemplateNode::Text(text)),
    }
}