from __future__ import annotations

from contextlib import contextmanager
import traceback

import rich
import rich.console
import rich.panel
import rich.table
import rich.columns

from .context import Context, ContextData


class RichLogger:
    def __init__(self, record: bool = False, indent_width=4) -> None:
        self._console = rich.console.Console(record=record)
        self._indent = 0
        self._indent_width = max(2, indent_width)

    def export_text(self) -> str:
        return self._console.export_text()

    def _indented(self, r, lc="│ "):
        sep = "│"
        i = f"{{:>{self._indent_width}}}".format(sep)

        ilc = f"??{{:>{self._indent_width}}}".format(lc)
        ilc = ((self._indent_width - 1) * " ") + lc
        col = rich.columns.Columns(
            [(self._indent // 2) * i + ilc, r],
            padding=(0, 0),
            # column_first=True,
            align="right",
        )
        return col

    def message(self, message: rich.console.RenderableType):
        r = self._indented(message, lc="├─ ")
        self._console.print(r)

    def log_dict(
        self,
        d: dict[rich.console.RenderableType, rich.console.RenderableType],
        title: str | None = None,
    ):
        table = rich.table.Table(title=title)
        for k, v in d.items():
            table.add_row(str(k), repr(v))
        self._console.print(self._indented(table))

    def _section_header(self, title: str):
        p = rich.panel.Panel(title, expand=False)
        return p

    def _section_footer(self):
        return ""

    @contextmanager
    def group(self, title: str):
        self.message(title)
        self._indent += 2
        try:
            yield
        finally:
            # self._indent -= 1
            self._console.print(self._indented("", lc=f"╰🠆 ✓"))
            self._indent -= 2

    @contextmanager
    def section(self, title: str):
        w = len(title)
        i = (self._indent_width - 2) * "─"
        self._console.print(self._indented("", lc="│╭─" + w * "─" + "─╮"))
        self._console.print(self._indented("", lc="├┤ " + title + " │"))
        self._console.print(
            self._indented("", lc=f"│╰{i}┬" + (w - self._indent_width + 2) * "─" + "─╯")
        )
        # self._console.print(rich.padding.Padding.indent(p, (indent * 4) - 2))
        self._indent += 2
        try:
            yield
        except Exception as err:

            if footer := self._section_footer():
                self._console.print(self._indented("", lc=f"├🠆 {footer}"))

            self._console.print(self._indented("[red]✘ Error:[/red]", lc=f"├🠆 "))
            tb = traceback.format_exc()
            self._console.print(rich.panel.Panel(tb, style="red"))
            if str(err):
                self._console.print(self._indented(f"[red]{err}[/red]", lc=f"├🠆 "))

        else:
            if footer := self._section_footer():
                self._console.print(self._indented("", lc=f"├🠆 {footer}"))
            msg = f"[green]🗹[/green] Done {title}"
            self._console.print(self._indented("", lc=f"╰🠆 {msg}"))
        finally:
            self._indent -= 2
            # self._console.print(self._indented("", lc="├─╯"))


@Context._decorate_context_type
class LoggableContext(ContextData):
    @classmethod
    def create_logger(cls, record: bool = False) -> RichLogger:
        return RichLogger(record=record)

    logger: RichLogger | None

    def _log_section_summary(self) -> str:
        return ""

    @contextmanager
    def _on_entering(self):
        if self.logger is not None:
            with self.logger.section(self._log_section_summary()):
                yield
        return

    def log(self, message: str):
        if self.logger is None:
            return
        self.logger.message(message)

    @contextmanager
    def log_group(self, title: str):
        if self.logger is None:
            return
        with self.logger.group(title):
            yield

    @contextmanager
    def log_section(self, title: str):
        if self.logger is None:
            return
        with self.logger.section(title):
            yield
