#!/usr/bin/env python3

import asyncio
import datetime
import enum
import random
import typing

import i3ipc


class ComposableText:
    def __init__(self, parent: typing.Optional["ComposableText"] = None) -> None:
        self.parent = parent

        prevParent = self
        while parent:
            prevParent = parent
            parent = parent.parent
        assert isinstance(prevParent, Bar)
        self.bar: Bar = prevParent

        self.text: str
        self.needsComposition = True

    def composeText(self) -> str:
        raise NotImplementedError(f"{self} cannot compose text")

    def getText(self) -> str:
        if self.needsComposition:
            self.text = self.composeText()
            print(f"{self} composed {self.text}")
            self.needsComposition = False
        return self.text

    def setText(self, text: str) -> None:
        self.text = text
        self.needsComposition = False
        parent = self.parent
        while parent and not parent.needsComposition:
            parent.needsComposition = True
            parent = parent.parent
        self.bar.refresh.set()


def randomColor(seed: int | bytes | None = None) -> str:
    if seed is not None:
        random.seed(seed)
    return "#" + "".join(f"{random.randint(0, 0xff):02x}" for _ in range(3))


class Section(ComposableText):
    """
    Colorable block separated by chevrons
    """

    def __init__(self, parent: "Module") -> None:
        super().__init__(parent=parent)
        self.color = randomColor()


class Module(ComposableText):
    """
    Sections handled by a same updater
    """

    def __init__(self, parent: "Side") -> None:
        super().__init__(parent=parent)
        self.sections: list[Section] = []

    def appendSection(self, section: Section) -> None:
        self.sections.append(section)


class Alignment(enum.Enum):
    LEFT = "l"
    RIGHT = "r"
    CENTER = "c"


class Side(ComposableText):
    def __init__(self, parent: "Screen", alignment: Alignment) -> None:
        super().__init__(parent=parent)
        self.alignment = alignment
        self.modules: list[Module] = []

    def composeText(self) -> str:
        if not self.modules:
            return ""
        text = "%{" + self.alignment.value + "}"
        lastSection: Section | None = None
        for module in self.modules:
            for section in module.sections:
                if lastSection is None:
                    if self.alignment == Alignment.LEFT:
                        text += "%{B" + section.color + "}%{F-}"
                    else:
                        text += "%{B-}%{F" + section.color + "}%{R}%{F-}"
                else:
                    if self.alignment == Alignment.RIGHT:
                        if lastSection.color == section.color:
                            text += ""
                        else:
                            text += "%{F" + section.color + "}%{R}"
                    else:
                        if lastSection.color == section.color:
                            text += ""
                        else:
                            text += "%{R}%{B" + section.color + "}"
                    text += "%{F-}"
                text += " " + section.getText() + " "
                # FIXME Should be handled by animation
                # FIXME Support hidden sections
                lastSection = section
        if self.alignment != Alignment.RIGHT:
            text += "%{R}%{B-}"
        return text


class Screen(ComposableText):
    def __init__(self, parent: "Bar", output: str) -> None:
        super().__init__(parent=parent)
        self.output = output

        self.sides = dict()
        for alignment in Alignment:
            self.sides[alignment] = Side(parent=self, alignment=alignment)

    def composeText(self) -> str:
        return ("%{Sn" + self.output + "}") + "".join(
            side.getText() for side in self.sides.values()
        )


ScreenSelection = enum.Enum("ScreenSelection", ["ALL", "PRIMARY", "SECONDARY"])


class Bar(ComposableText):
    """
    Top-level
    """

    def __init__(self) -> None:
        super().__init__()
        self.providers: list["Provider"] = []
        self.refresh = asyncio.Event()

        self.screens = []
        i3 = i3ipc.Connection()
        for output in i3.get_outputs():
            if not output.active:
                continue
            screen = Screen(parent=self, output=output.name)
            self.screens.append(screen)

    async def run(self) -> None:
        proc = await asyncio.create_subprocess_exec(
            "lemonbar",
            "-b",
            "-a",
            "64",
            "-f",
            "DejaVuSansM Nerd Font:size=10",
            stdin=asyncio.subprocess.PIPE,
        )

        async def refresher() -> None:
            assert proc.stdin
            while True:
                await self.refresh.wait()
                self.refresh.clear()
                proc.stdin.write(self.getText().encode())

        async with asyncio.TaskGroup() as tg:
            tg.create_task(refresher())
            for updater in self.providers:
                tg.create_task(updater.run())

    def composeText(self) -> str:
        return "".join(section.getText() for section in self.screens) + "\n"

    def addProvider(
        self,
        provider: "Provider",
        alignment: Alignment = Alignment.LEFT,
        screenSelection: ScreenSelection = ScreenSelection.ALL,
    ) -> None:
        # FIXME Actually have a screenNum and screenCount args
        modules = list()
        for s, screen in enumerate(self.screens):
            if (
                screenSelection == ScreenSelection.ALL
                or (screenSelection == ScreenSelection.PRIMARY and s == 0)
                or (screenSelection == ScreenSelection.SECONDARY and s == 1)
            ):
                side = screen.sides[alignment]
                module = Module(parent=side)
                side.modules.append(module)
                modules.append(module)
        provider.modules = modules
        if modules:
            self.providers.append(provider)


class Provider:
    def __init__(self) -> None:
        self.modules: list[Module] = list()

    async def run(self) -> None:
        raise NotImplementedError()


class TimeProvider(Provider):
    async def run(self) -> None:
        sections = list()
        for module in self.modules:
            section = Section(parent=module)
            module.sections.append(section)
            sections.append(section)
        # FIXME Allow for mirror(ed) modules so no need for updaters to handle all

        while True:
            for s, section in enumerate(sections):
                now = datetime.datetime.now()
                section.setText(now.strftime("%a %y-%m-%d %H:%M:%S"))
            await asyncio.sleep(1)


async def main() -> None:
    bar = Bar()
    bar.addProvider(TimeProvider())
    # bar.addProvider(TimeProvider(), screenSelection=ScreenSelection.PRIMARY)
    # bar.addProvider(TimeProvider(), alignment=Alignment.CENTER)
    # bar.addProvider(TimeProvider(), alignment=Alignment.CENTER)
    # bar.addProvider(TimeProvider(), alignment=Alignment.RIGHT)
    # bar.addProvider(TimeProvider(), alignment=Alignment.RIGHT)
    await bar.run()
    # FIXME Why time not updating?


asyncio.run(main())