#!/usr/bin/python3
from enum import Enum
import re
import typing as T
from dataclasses import dataclass
from pathlib import Path

GRID_MAX_SQUARES = 50
GRID_SQUARE_SIZE = 12
GRID_SQUARE_MARGIN = 2
LEGEND_SQUARE_SIZE = 12
LEGEND_SQUARE_MARGIN = 2
LEGEND_ROW_PADDING = 3
LEGEND_MARGIN = 15
TEXT_SIZE = 15
TEXT_MARGIN = 5
DOCS_DIR = Path(__file__).parent
PROGRESS_TXT_FILE = DOCS_DIR / "progress.txt"
PROGRESS_SVG_FILE = DOCS_DIR / "progress.svg"


class FunctionStatus(Enum):
    todo = "todo"
    named = "named"
    decompiled = "decompiled"


COLOR_MAP = {
    FunctionStatus.decompiled: "forestgreen",
    FunctionStatus.named: "lightpink",
    FunctionStatus.todo: "mistyrose",
}


@dataclass
class Box:
    x1: float
    y1: float
    x2: float
    y2: float

    @property
    def dx(self) -> float:
        return self.x2 - self.x1

    @property
    def dy(self) -> float:
        return self.y2 - self.y1


@dataclass
class SquarifyResult(Box):
    item: T.Any


class Squarify:
    def __init__(
        self,
        items: T.Iterable[T.Any],
        key: T.Callable[[T.Any], float],
        box: Box,
        normalize: bool = True,
    ) -> None:
        if normalize:
            total_size = sum(map(key, items))
            total_area = box.dx * box.dy

            def normalized_key(item: T.Any) -> float:
                return key(item) * total_area / total_size

            self.key = normalized_key
        else:
            self.key = key

        self.items = list(sorted(items, key=self.key, reverse=True))
        self.box = box

    def layoutrow(
        self, items: T.Iterable[T.Any]
    ) -> T.Iterable[SquarifyResult]:
        covered_area = sum(self.key(item) for item in items)
        dx = covered_area / self.box.dy
        y = self.box.y1
        for item in items:
            yield SquarifyResult(
                item=item,
                x1=self.box.x1,
                y1=y,
                x2=self.box.x1 + dx,
                y2=y + self.key(item) / dx,
            )
            y += self.key(item) / dx

    def layoutcol(
        self, items: T.Iterable[T.Any]
    ) -> T.Iterable[SquarifyResult]:
        covered_area = sum(self.key(item) for item in items)
        dy = covered_area / self.box.dx
        x = self.box.x1
        for item in items:
            yield SquarifyResult(
                item=item,
                x1=x,
                y1=self.box.y1,
                x2=x + self.key(item) / dy,
                y2=self.box.y1 + dy,
            )
            x += self.key(item) / dy

    def layout(self, items: T.Iterable[T.Any]) -> T.Iterable[SquarifyResult]:
        yield from (
            self.layoutrow(items)
            if self.box.dx >= self.box.dy
            else self.layoutcol(items)
        )

    def leftoverrow(self, items: T.Iterable[T.Any]) -> Box:
        covered_area = sum(self.key(item) for item in items)
        dx = covered_area / self.box.dy
        return Box(
            x1=self.box.x1 + dx,
            y1=self.box.y1,
            x2=self.box.x1 + self.box.dx,
            y2=self.box.y1 + self.box.dy,
        )

    def leftovercol(self, items: T.Iterable[T.Any]) -> Box:
        covered_area = sum(self.key(item) for item in items)
        dy = covered_area / self.box.dx
        return Box(
            x1=self.box.x1,
            y1=self.box.y1 + dy,
            x2=self.box.x1 + self.box.dx,
            y2=self.box.y1 + self.box.dy,
        )

    def leftover(self, items: T.Iterable[T.Any]) -> Box:
        return (
            self.leftoverrow(items)
            if self.box.dx >= self.box.dy
            else self.leftovercol(items)
        )

    def worst_ratio(self, items: T.Iterable[T.Any]) -> float:
        return max(
            max(result.dx / result.dy, result.dy / result.dx)
            for result in self.layout(items)
        )

    def run(
        self, items: T.Optional[T.List[T.Any]] = None
    ) -> T.Iterable[SquarifyResult]:
        if not items:
            items = self.items

        if len(items) == 0:
            return

        if len(items) == 1:
            yield from self.layout(items)
            return

        i = 1
        while i < len(items) and self.worst_ratio(
            items[:i]
        ) >= self.worst_ratio(items[: i + 1]):
            i += 1
        current = items[:i]
        remaining = items[i:]

        leftover_box = self.leftover(current)
        yield from self.layout(current)
        yield from Squarify(
            remaining, key=self.key, box=leftover_box, normalize=False
        ).run()


def squarify(
    items: T.List[T.Any], key: T.Callable[[T.Any], float], box: Box
) -> T.Iterable[SquarifyResult]:
    yield from Squarify(items, key, box).run()


@dataclass
class Function:
    name: str
    offset: int
    size: int
    flags: str

    @property
    def is_decompiled(self) -> bool:
        return "+" in self.flags

    @property
    def is_called(self) -> bool:
        return "*" in self.flags

    @property
    def is_named(self) -> bool:
        return not self.name.startswith("sub_")

    @property
    def status(self) -> FunctionStatus:
        if self.is_decompiled:
            return FunctionStatus.decompiled
        elif self.is_named:
            return FunctionStatus.named
        return FunctionStatus.todo


def collect_functions() -> T.Iterable[Function]:
    for line in PROGRESS_TXT_FILE.open():
        line = line.strip()
        if line.startswith("#") or not line:
            continue
        func_name, offset, size, flags = re.split(r"\s+", line)
        if not offset.replace("-", ""):
            continue
        yield Function(
            name=func_name,
            offset=int(offset, 16),
            size=int(size, 16),
            flags=flags,
        )


class Shape:
    @property
    def box(self) -> Box:
        raise NotImplementedError("not implemented")

    def render(self) -> str:
        raise NotImplementedError("not implemented")


def get_common_bbox(shapes: T.List[Shape]) -> Box:
    return Box(
        x1=min(shape.box.x1 for shape in shapes),
        y1=min(shape.box.y1 for shape in shapes),
        x2=max(shape.box.x2 for shape in shapes),
        y2=max(shape.box.y2 for shape in shapes),
    )


@dataclass
class Rectangle(Shape):
    x: float
    y: float
    class_name: str
    dx: float
    dy: float
    title: T.Optional[str] = None

    @property
    def box(self) -> Box:
        return Box(
            x1=self.x,
            y1=self.y,
            x2=self.x + self.dx,
            y2=self.y + self.dy,
        )

    def render(self) -> str:
        return (
            f"<rect "
            f'width="{self.dx:.02f}" '
            f'height="{self.dy:.02f}" '
            f'x="{self.x:.02f}" '
            f'y="{self.y:.02f}" '
            f'class="{self.class_name}"'
            + (f"><title>{self.title}</title></rect>" if self.title else "/>")
        )


class Square(Rectangle):
    def __init__(
        self,
        x: float,
        y: float,
        class_name: str,
        size: float = GRID_SQUARE_SIZE,
        title: T.Optional[str] = None,
    ) -> None:
        super().__init__(
            x=x, y=y, class_name=class_name, dx=size, dy=size, title=title
        )


@dataclass
class Text(Shape):
    x: float
    y: float
    text: str

    @property
    def box(self) -> Box:
        return Box(x1=self.x, y1=self.y, x2=self.x, y2=self.y + TEXT_SIZE)

    def render(self) -> str:
        return (
            f"<text "
            f'x="{self.x:.02f}" y="{self.y + TEXT_SIZE/2:.02f}">'
            f"{self.text}"
            f"</text>"
        )


@dataclass
class LegendText(Shape):
    x: float
    y: float
    class_name: str
    text: str

    @property
    def _square(self) -> Square:
        return Square(
            x=self.x,
            y=self.y + (TEXT_SIZE - LEGEND_SQUARE_SIZE) / 2,
            class_name=self.class_name,
            size=LEGEND_SQUARE_SIZE,
        )

    @property
    def _text(self) -> Text:
        return Text(
            x=LEGEND_SQUARE_SIZE + TEXT_MARGIN,
            y=self.y,
            text=self.text,
        )

    @property
    def box(self) -> Box:
        return get_common_bbox([self._square, self._text])

    def render(self) -> str:
        return self._square.render() + self._text.render()


def render_grid(
    all_functions: T.List[Function], by: float
) -> T.Iterable[Shape]:
    for i, function in enumerate(all_functions):
        x = (i % GRID_MAX_SQUARES) * (GRID_SQUARE_SIZE + GRID_SQUARE_MARGIN)
        y = (i // GRID_MAX_SQUARES) * (GRID_SQUARE_SIZE + GRID_SQUARE_MARGIN)
        yield Square(
            x=x,
            y=by + y,
            class_name=function.status.value,
            title=function.name,
        )


def render_tree_grid(
    all_functions: T.List[Function], box: Box
) -> T.Iterable[Shape]:
    for result in squarify(
        all_functions, key=lambda function: float(function.size), box=box
    ):
        result.x2 = max(0, result.x2 - GRID_SQUARE_MARGIN)
        result.y2 = max(0, result.y2 - GRID_SQUARE_MARGIN)

        yield Rectangle(
            x=result.x1,
            y=result.y1,
            dx=result.dx,
            dy=result.dy,
            class_name=result.item.status.value,
            title=result.item.name,
        )


def render_svg(all_functions: T.List[Function]) -> T.Iterable[Shape]:
    y = 0.0
    yield Text(
        0,
        y,
        "TombATI functions, arranged according to the physical file layout:",
    )
    y = TEXT_SIZE + GRID_SQUARE_MARGIN

    shapes = list(render_grid(all_functions, y))
    yield from shapes

    y = max(shape.box.y2 for shape in shapes)
    y += GRID_SQUARE_SIZE + LEGEND_MARGIN
    yield Text(
        0, y, "TombATI functions, arranged according to the function sizes:"
    )
    y += TEXT_SIZE + GRID_SQUARE_MARGIN
    dx = max(shape.box.x2 for shape in shapes)

    shapes = list(
        render_tree_grid(all_functions, Box(x1=0, y1=y, x2=dx, y2=y + y))
    )
    yield from shapes

    y = max(shape.box.y2 for shape in shapes)
    y += GRID_SQUARE_SIZE + LEGEND_MARGIN

    ready_functions = [func for func in all_functions if func.is_decompiled]
    named_functions = [
        func
        for func in all_functions
        if func.is_named and func not in ready_functions
    ]
    todo_functions = [
        func
        for func in all_functions
        if func not in ready_functions and func not in named_functions
    ]

    def sum_size(functions: T.Iterable[Function]) -> int:
        return sum(func.size for func in functions)

    for (status, text, functions) in [
        (FunctionStatus.decompiled, "Functions decompiled", ready_functions),
        (
            FunctionStatus.named,
            "Functions not decompiled, but with known names",
            named_functions,
        ),
        (
            FunctionStatus.todo,
            "Functions not decompiled, with unknown names",
            todo_functions,
        ),
    ]:
        yield LegendText(
            x=0,
            y=y,
            class_name=status.value,
            text=f"{text} (count): {len(functions)/len(all_functions):.02%}",
        )
        y += TEXT_SIZE + LEGEND_ROW_PADDING

        yield LegendText(
            x=0,
            y=y,
            class_name=status.value,
            text=f"{text} (bytesize): {sum_size(functions)/sum_size(all_functions):.02%}",
        )
        y += TEXT_SIZE + LEGEND_ROW_PADDING


def main() -> None:
    functions = list(collect_functions())

    with PROGRESS_SVG_FILE.open("w") as handle:
        shapes = list(render_svg(functions))
        bbox = get_common_bbox(shapes)

        svg = (
            f'<svg version="1.1" '
            f'width="{bbox.dx:.02f}" '
            f'height="{bbox.dy:.02f}" '
            f'xmlns="http://www.w3.org/2000/svg">'
            f"<style>"
        )
        for status in FunctionStatus:
            svg += f".{status.value}{{fill:{COLOR_MAP[status]}}}"
        svg += f"text{{alignment-baseline:central;font-family:sans-serif;font-size:{TEXT_SIZE}px}}"
        svg += "</style>"

        print(
            svg,
            file=handle,
        )

        for shape in shapes:
            print(shape.render(), file=handle)

        print("</svg>", file=handle)


if __name__ == "__main__":
    main()