# Copyright 2026 Marimo. All rights reserved.
from __future__ import annotations

import msgspec


class DependencyTag(msgspec.Struct, rename="camel"):
    kind: str
    value: str


class DependencyTreeNode(msgspec.Struct, rename="camel"):
    name: str
    version: str | None
    # List of {"kind": "extra"|"group", "value": str}
    tags: list[DependencyTag]
    dependencies: list[DependencyTreeNode]


def parse_name_version(content: str) -> tuple[str, str | None]:
    """Parse package name and version from uv tree output."""
    if " v" in content:
        name, version = content.split(" v", 1)
        return name.strip(), version.split()[0]  # Take only version part
    return content.strip(), None


def parse_uv_tree(text: str) -> DependencyTreeNode:
    """Parse the text output of `uv tree` into a nested data structure."""
    lines = text.strip().split("\n")

    # Create a virtual root to hold all top-level dependencies
    tree = DependencyTreeNode(
        name="<root>", version=None, tags=[], dependencies=[]
    )
    stack = [(tree, -1)]  # (node, level)

    for line in lines:
        line = line.rstrip()
        if (
            not line
            or "Package tree already displayed" in line
            or "Package tree is a cycle" in line
        ):
            continue

        # Calculate indentation level by counting characters before tree symbols
        if not any(symbol in line for symbol in ["├──", "└──"]):
            level = 0  # Top-level package
        else:
            # Find the tree symbol position and divide by 4 (standard tree indentation)
            for symbol in ["├──", "└──"]:
                pos = line.find(symbol)
                if pos != -1:
                    level = (pos // 4) + 1
                    break

        # content after tree symbols
        content = line.lstrip("│ ├└─").strip()

        # Check for cycle indicator
        is_cycle = content.endswith("(*)")
        if is_cycle:
            content = content[:-3].strip()

        # tags (extras/groups)
        tags: list[DependencyTag] = []

        while "(extra:" in content or "(group:" in content:
            start = (
                content.rfind("(extra:")
                if "(extra:" in content
                else content.rfind("(group:")
            )
            if start == -1:
                break
            end = content.find(")", start)
            if end == -1:
                break
            tag_text = content[start + 1 : end]
            kind, value = tag_text.split(":", 1)
            assert kind == "extra" or kind == "group"
            tags.append(DependencyTag(kind=kind, value=value.strip()))
            content = content[:start].strip()

        name, version = parse_name_version(content)

        # Add cycle indicator as a special tag
        if is_cycle:
            tags.append(DependencyTag(kind="cycle", value="true"))

        node = DependencyTreeNode(
            name=name,
            version=version,
            tags=tags,
            dependencies=[],
        )

        # Adjust stack to correct level
        while len(stack) > 1 and stack[-1][1] >= level:
            stack.pop()

        # Add to parent and push to stack
        stack[-1][0].dependencies.append(node)
        stack.append((node, level))

    return tree
