import dataclasses
from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import date, datetime, timedelta
from typing import Any, Callable, Iterable, Literal


@dataclass(frozen=True)
class HTML():
    html: str
    def _repr_html_(self):
        return self.html

@dataclass(frozen=True)
class Values(ABC):
    @abstractmethod
    def summary(self) -> str:
        pass
    @abstractmethod
    def __len__(self) -> int:
        pass

    @abstractmethod
    def __contains__(self, value: Any) -> bool:
        pass

    @abstractmethod
    def from_strings(self, values: list[str]) -> list['Values']:
        pass

@dataclass(frozen=True)
class Enum(Values):
    """
    The simplest kind of key value is just a list of strings.
    summary -> string1/string2/string....
    """
    values: list[Any]

    def __len__(self) -> int:
        return len(self.values)
    def summary(self) -> str:
        return '/'.join(sorted(self.values))
    def __contains__(self, value: Any) -> bool:
        return value in self.values
    def from_strings(self, values: list[str]) -> list['Values']:
        return [Enum(values)]

@dataclass(frozen=True)
class Range(Values, ABC):
    dtype: str = dataclasses.field(kw_only=True)

@dataclass(frozen=True)
class DateRange(Range):
    start: date
    end: date
    step: timedelta
    dtype: Literal["date"] = dataclasses.field(kw_only=True, default="date")

    @classmethod
    def from_strings(self, values: list[str]) -> list['DateRange']:
        dates = sorted([datetime.strptime(v, "%Y%m%d") for v in values])
        if len(dates) < 2:
            return [DateRange(
                start=dates[0],
                end=dates[0],
                step=timedelta(days=0)
            )]
        
        ranges = []
        current_range, dates = [dates[0],], dates[1:]
        while len(dates) > 1:
            if dates[0] - current_range[-1] == timedelta(days=1):
                current_range.append(dates.pop(0))
            
            elif len(current_range) == 1:
                ranges.append(DateRange(
                start=current_range[0],
                end=current_range[0],
                step=timedelta(days=0)
                ))
                current_range = [dates.pop(0),]

            else:
                ranges.append(DateRange(
                start=current_range[0],
                end=current_range[-1],
                step=timedelta(days=1)
                ))
                current_range = [dates.pop(0),]
        return ranges
    
    def __contains__(self, value: Any) -> bool:
        v = datetime.strptime(value, "%Y%m%d").date()
        return self.start <= v <= self.end and (v - self.start) % self.step == 0


    def __len__(self) -> int:
        return (self.end - self.start) // self.step
    
    def summary(self) -> str:
        def fmt(d): return d.strftime("%Y%m%d")
        if self.step == timedelta(days=0):
            return f"{fmt(self.start)}"
        if self.step == timedelta(days=1):
            return f"{fmt(self.start)}/to/{fmt(self.end)}"
        
        return f"{fmt(self.start)}/to/{fmt(self.end)}/by/{self.step // timedelta(days=1)}"

@dataclass(frozen=True)
class TimeRange(Range):
    start: int
    end: int
    step: int
    dtype: Literal["time"] = dataclasses.field(kw_only=True, default="time")

    @classmethod
    def from_strings(self, values: list[str]) -> list['TimeRange']:
        if len(values) == 0: return []

        times = sorted([int(v) for v in values])
        if len(times) < 2:
            return [TimeRange(
                start=times[0],
                end=times[0],
                step=100
            )]
        
        ranges = []
        current_range, times = [times[0],], times[1:]
        while len(times) > 1:
            if times[0] - current_range[-1] == 1:
                current_range.append(times.pop(0))
            
            elif len(current_range) == 1:
                ranges.append(TimeRange(
                start=current_range[0],
                end=current_range[0],
                step=0
                ))
                current_range = [times.pop(0),]

            else:
                ranges.append(TimeRange(
                start=current_range[0],
                end=current_range[-1],
                step=1
                ))
                current_range = [times.pop(0),]
        return ranges

    def __len__(self) -> int:
        return (self.end - self.start) // self.step
    
    def summary(self) -> str:
        def fmt(d): return f"{d:04d}"
        if self.step == 0:
            return f"{fmt(self.start)}"
        return f"{fmt(self.start)}/to/{fmt(self.end)}/by/{self.step}"
    
    def __contains__(self, value: Any) -> bool:
        v = int(value)
        return self.start <= v <= self.end and (v - self.start) % self.step == 0

@dataclass(frozen=True)
class IntRange(Range):
    dtype: Literal["int"]
    start: int
    end: int
    step: int
    dtype: Literal["int"] = dataclasses.field(kw_only=True, default="int")

    def __len__(self) -> int:
        return (self.end - self.start) // self.step
    
    def summary(self) -> str:
        def fmt(d): return d.strftime("%Y%m%d")
        return f"{fmt(self.start)}/to/{fmt(self.end)}/by/{self.step}"
    
    def __contains__(self, value: Any) -> bool:
        v = int(value)
        return self.start <= v <= self.end and (v - self.start) % self.step == 0
    

def values_from_json(obj) -> Values:
    if isinstance(obj, list): 
        return Enum(obj)

    match obj["dtype"]:
        case "date": return DateRange(**obj)
        case "time": return TimeRange(**obj)
        case "int": return IntRange(**obj)
        case _: raise ValueError(f"Unknown dtype {obj['dtype']}")

@dataclass(frozen=True)
class Node:
    key: str
    values: Values # Must support len()
    metadata: dict[str, str] # Applies to all children
    payload: list[Any] # List of size product(len(n.values) for n in  ancestors(self))
    children: list['Node']

def summarize_node(node: Node) -> tuple[str, Node]:
    """
    Extracts a summarized representation of the node while collapsing single-child paths.
    Returns the summary string and the last node in the chain that has multiple children.
    """
    summary = []
    
    while True:
        values_summary = node.values.summary()
        if len(values_summary) > 50:
            values_summary = values_summary[:50] + "..."
        summary.append(f"{node.key}={values_summary}")

        # Move down if there's exactly one child, otherwise stop
        if len(node.children) != 1:
            break
        node = node.children[0]

    return ", ".join(summary), node

def node_tree_to_string(node : Node, prefix : str = "", depth = None) -> Iterable[str]:
    summary, node = summarize_node(node)
    
    if depth is not None and depth <= 0:
        yield summary + " - ...\n"
        return
    # Special case for nodes with only a single child, this makes the printed representation more compact
    elif len(node.children) == 1:
        yield summary + ", "
        yield from node_tree_to_string(node.children[0], prefix, depth = depth)
        return
    else:
        yield summary + "\n"
    
    for index, child in enumerate(node.children):
        connector = "└── " if index == len(node.children) - 1 else "├── "
        yield prefix + connector
        extension = "    " if index == len(node.children) - 1 else "│   "
        yield from node_tree_to_string(child, prefix + extension, depth = depth - 1 if depth is not None else None)

def node_tree_to_html(node : Node, prefix : str = "", depth = 1, connector = "") -> Iterable[str]:
    summary, node = summarize_node(node)
    
    if len(node.children) == 0:
        yield f'<span class="leaf">{connector}{summary}</span>'
        return
    else:
        open = "open" if depth > 0 else ""
        yield f"<details {open}><summary>{connector}{summary}</summary>"

    for index, child in enumerate(node.children):
        connector = "└── " if index == len(node.children) - 1 else "├── "
        extension = "    " if index == len(node.children) - 1 else "│   "
        yield from node_tree_to_html(child, prefix + extension, depth = depth - 1, connector = prefix+connector)
    yield "</details>"

@dataclass(frozen=True)
class CompressedTree:
    root: Node

    @classmethod
    def from_json(cls, json: dict) -> 'CompressedTree':
        def from_json(json: dict) -> Node:
            return Node(
                key=json["key"],
                values=values_from_json(json["values"]),
                metadata=json["metadata"] if "metadata" in json else {},
                payload=json["payload"] if "payload" in json else [],
                children=[from_json(c) for c in json["children"]]
            )
        return CompressedTree(root=from_json(json))
    
    def __str__(self):
        return "".join(node_tree_to_string(node=self.root))
    
    def html(self, depth = 2) -> HTML:
        return HTML(self._repr_html_(depth = depth))
    
    def _repr_html_(self, depth = 2):
        css = """
        <style>
        .qubed-tree-view {
            font-family: monospace;
            white-space: pre;
        }
        .qubed-tree-view details {
            # display: inline;
            margin-left: 0;
        }
        .qubed-tree-view summary {
            list-style: none;
            cursor: pointer;
            text-overflow: ellipsis;
            overflow: hidden;
            text-wrap: nowrap;
            display: block;
        }

        .qubed-tree-view .leaf {
            text-overflow: ellipsis;
            overflow: hidden;
            text-wrap: nowrap;
            display: block;
        }

        .qubed-tree-view summary:hover,span.leaf:hover {
            background-color: #f0f0f0;
        }
        .qubed-tree-view details > summary::after {
            content: ' ';
        }
        .qubed-tree-view details:not([open]) > summary::after {
            content: " ▼";
        }
        </style>

        """
        nodes = "".join(cc for c in self.root.children for cc in node_tree_to_html(node=c, depth=depth))
        return f"{css}<pre class='qubed-tree-view'>{nodes}</pre>"
    
    def print(self, depth = None):
        print("".join(cc for c in self.root.children for cc in node_tree_to_string(node=c, depth = depth)))

    def transform(self, func: Callable[[Node], Node]) -> 'CompressedTree':
        "Call a function on every node of the tree, any changes to the children of a node will be ignored."
        def transform(node: Node) -> Node:
            new_node = func(node)
            return dataclasses.replace(new_node, children = [transform(c) for c in node.children])
        return CompressedTree(root=transform(self.root))

    def guess_datatypes(self) -> 'CompressedTree':
        def guess_datatypes(node: Node) -> list[Node]:
            # Try to convert enum values into more structured types
            children = [cc for c in node.children for cc in guess_datatypes(c)]

            if isinstance(node.values, Enum):
                match node.key:
                    case "time": range_class = TimeRange
                    case "date": range_class = DateRange
                    case _: range_class = None

                if range_class is not None:
                    return [
                        dataclasses.replace(node, values = range, children = children)
                        for range in range_class.from_strings(node.values.values)
                    ]
            return [dataclasses.replace(node, children = children)]

        children = [cc for c in self.root.children for cc in guess_datatypes(c)]
        return CompressedTree(root=dataclasses.replace(self.root, children = children))

    
    def select(self, selection : dict[str, str | list[str]], mode: Literal["strict", "relaxed"] = "relaxed") -> 'CompressedTree':
        # make all values lists
        selection = {k : v if isinstance(v, list) else [v] for k,v in selection.items()}

        def not_none(xs): return [x for x in xs if x is not None]

        def select(node: Node) -> Node | None: 
            # Check if the key is specified in the selection
            if node.key not in selection: 
                if mode == "strict":
                    return None
                return dataclasses.replace(node, children = not_none(select(c) for c in node.children))
            
            # If the key is specified, check if any of the values match
            values = Enum([ c for c in selection[node.key] if c in node.values])

            if not values: 
                return None 
            
            return dataclasses.replace(node, values = values, children = not_none(select(c) for c in node.children))
            
        return CompressedTree(root=dataclasses.replace(self.root, children = not_none(select(c) for c in self.root.children)))
    
    def to_list_of_cubes(self):
        def to_list_of_cubes(node: Node) -> list[list[Node]]:
            return [[node] + sub_cube for c in node.children for sub_cube in to_list_of_cubes(c)]

        return to_list_of_cubes(self.root)

    def info(self):
        cubes = self.to_list_of_cubes()
        print(f"Number of distinct paths: {len(cubes)}")

        

        
# What should the interace look like?

# tree = CompressedTree.from_json(...)
# tree = CompressedTree.from_protobuf(...)

# tree.print(depth = 5) # Prints a nice tree representation