#!/usr/bin/env python3
# -*- coding:utf-8 -*-

from __future__ import annotations

from datetime import datetime, timezone, tzinfo

from rich.table import Table, box
from typer import Exit, Typer

from pendingai.cli.console import Console
from pendingai.cli.context import PendingAiContext
from pendingai.cli.retro import batch, job
from pendingai.cli.shared import JsonOption
from pendingai.cli.utils import catch_exception
from pendingai.services.retrosynthesis.engines import Engine
from pendingai.services.retrosynthesis.libraries import Library

cout = Console()
tz: tzinfo | None = datetime.now().astimezone().tzinfo


app = Typer(
    name="retro",
    help=(
        "High-throughput deep-learning retrosynthesis service."
        "\n\nRetrosynthesis predictions are generated by submitting "
        "a batch of jobs to an engine, monitoring its status, and retrieving "
        "completed results. For more information refer to the "
        "documentation with <pendingai docs>."
    ),
    short_help="High-throughput deep-learning retrosynthesis service.",
)
app.add_typer(job.app)
app.add_typer(batch.app)

# region command: engines ----------------------------------------------


class EngineTable(Table):
    def __init__(self, engines: list[Engine] = []):
        super().__init__("ID", "Name", "Last Alive", box=box.SQUARE)
        for engine in engines:
            engine.last_alive = engine.last_alive.replace(tzinfo=timezone.utc)
            seconds: float = (datetime.now(tz) - engine.last_alive).total_seconds()
            seconds_t: str = f"{abs(seconds):3.0f}" if abs(seconds) < 60 else "Over 60"
            self.add_row(engine.id, engine.name, seconds_t + " second(s) ago")


@app.command("engines", help="List available retrosynthesis engines.")
@catch_exception()
def _engines(ctx: PendingAiContext, json: JsonOption = False):
    engines: list[Engine] = ctx.obj["client"].retrosynthesis.engines.list()
    engines.sort(key=lambda x: x.last_alive, reverse=True)
    if len(engines) == 0:
        cout.print("[warn]! No retrosynthesis engines available.")
        raise Exit(0)
    if json:
        cout.print_json(data=engines)
    else:
        cout.print(EngineTable(engines))


# region command: libraries --------------------------------------------


class LibrariesTable(Table):
    def __init__(self, libraries: list[Library] = []):
        super().__init__("ID", "Name", "Version", box=box.SQUARE)
        for library in libraries:
            self.add_row(library.id, library.name, library.version)


@app.command("libraries", help="List available building block libraries.")
@catch_exception()
def _libraries(ctx: PendingAiContext, json: JsonOption = False):
    libraries: list[Library] = ctx.obj["client"].retrosynthesis.libraries.list()
    if len(libraries) == 0:
        cout.print("[warn]! No building block libraries available.")
        raise Exit(0)
    if json:
        cout.print_json(data=libraries)
    else:
        cout.print(LibrariesTable(libraries))
