from __future__ import annotations

import re
from typing import Any, Dict, List, Tuple


def extract_nodes_from_markdown(markdown_content: str) -> Tuple[List[Dict[str, Any]], List[str]]:
    header_pattern = r"^(#{1,6})\s+(.+)$"
    code_block_pattern = r"^```"
    node_list: List[Dict[str, Any]] = []
    lines = markdown_content.split("\n")
    in_code_block = False

    for line_num, line in enumerate(lines, 1):
        stripped = line.strip()
        if re.match(code_block_pattern, stripped):
            in_code_block = not in_code_block
            continue
        if not stripped:
            continue
        if in_code_block:
            continue
        match = re.match(header_pattern, stripped)
        if match:
            node_list.append({"node_title": match.group(2).strip(), "line_num": line_num, "level": len(match.group(1))})
    return node_list, lines


def extract_node_text_content(node_list: List[Dict[str, Any]], markdown_lines: List[str]) -> List[Dict[str, Any]]:
    out: List[Dict[str, Any]] = []
    for node in node_list:
        out.append(
            {
                "title": node["node_title"],
                "line_num": node["line_num"],
                "level": int(node["level"]),
            }
        )

    for i, node in enumerate(out):
        start_line = node["line_num"] - 1
        if i + 1 < len(out):
            end_line = out[i + 1]["line_num"] - 1
        else:
            end_line = len(markdown_lines)
        node["text"] = "\n".join(markdown_lines[start_line:end_line]).strip()
    return out


def build_tree_from_nodes(node_list: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    if not node_list:
        return []

    stack: List[Tuple[Dict[str, Any], int]] = []
    roots: List[Dict[str, Any]] = []
    counter = 0

    for node in node_list:
        counter += 1
        tree_node = {
            "title": node["title"],
            "node_id": str(counter).zfill(4),
            "text": node.get("text", ""),
            "line_num": node["line_num"],
            "level": node["level"],
            "nodes": [],
        }

        while stack and stack[-1][1] >= node["level"]:
            stack.pop()

        if not stack:
            roots.append(tree_node)
        else:
            stack[-1][0]["nodes"].append(tree_node)

        stack.append((tree_node, node["level"]))

    return roots


async def markdown_to_tree_from_content(markdown_content: str) -> Dict[str, Any]:
    node_list, lines = extract_nodes_from_markdown(markdown_content)
    nodes = extract_node_text_content(node_list, lines)
    tree = build_tree_from_nodes(nodes)
    return {"doc_name": "docling_markdown", "structure": tree}
