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

import sys
import typing
from datetime import datetime, timezone

import pytz
import rich
import rich.prompt
import rich.table
import typer

from pendingai import config
from pendingai.commands.retro.batch import command as batch_command
from pendingai.commands.retro.controller import RetroController
from pendingai.commands.retro.job import command as job_command
from pendingai.commands.retro.models import BuildingBlockLibrary, RetrosynthesisEngine
from pendingai.context import Context

console = rich.console.Console(theme=config.RICH_CONSOLE_THEME, soft_wrap=True)

app = typer.Typer(
    name="retro",
    help=(
        "High-throughput deep-learning retrosynthesis service."
        "\n\nRetrosynthesis predictions are generated by submitting "
        "jobs to an engine, monitoring their status, and retrieving "
        "the completed results. For more information refer to the "
        "documentation with <pendingai docs>."
    ),
    short_help="High-throughput deep-learning retrosynthesis service.",
    no_args_is_help=True,
    add_completion=False,
    pretty_exceptions_show_locals=False,
    rich_markup_mode=None,
)
app.add_typer(job_command.app)
app.add_typer(batch_command.app)

# region callback: auth ------------------------------------------------


@app.callback()
def enforce_authorized_access(context: Context) -> None:
    """
    Require access to all commands to precheck access permissions with
    an aliveness check and authentication check to verify that service
    access is allowed for the user.

    Args:
        context (Context): App runtime context.

    Raises:
        typer.Exit: User is not authorized for the subcommand service or
            there is no active session found in the app runtime context.
    """
    controller = RetroController(context)
    if "--help" not in sys.argv:
        # check the app runtime context cache to see if the access token
        # exists and also that it has not yet expired.
        if context.obj.cache.access_token is None:
            console.print(
                "[warn]! You are not logged in, use [code]pendingai auth "
                "login[/] to create a new session."
            )
            raise typer.Exit(1)
        elif context.obj.cache.access_token.is_expired():
            console.print(
                "[warn]! You are not logged in, use [code]pendingai auth "
                "login[/] to create a new session."
            )
            raise typer.Exit(1)

        # make api requests to check if a user is not subscribed to the
        # service i.e., a non 200 response is given, also check that the
        # service is alive and reachable.
        if not controller.api.check_authorization():
            console.print(
                "[warn]! Unauthorized access, please contact support "
                "if you would like to access this service."
            )
            raise typer.Exit(1)
        elif not controller.api.check_aliveness():
            console.print(
                "[warn]! Service is temporarily unavaiable, "
                "please contact support if the service remains unavaiable."
            )
            raise typer.Exit(1)


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


@app.command(
    "engines",
    help=(
        "List available retrosynthesis engines. Their unique IDs are required "
        "for submitting jobs. If no engines are available please contact support."
    ),
    short_help="List available retrosynthesis engines.",
)
def retrieve_engines(
    context: Context,
    render_json: typing.Annotated[
        bool,
        typer.Option(
            "--json",
            help="Render output as JSON.",
            is_flag=True,
        ),
    ] = False,
) -> None:
    """
    Retrieve a list of available retrosynthesis engines stored in the
    database and output with relevant information for the user.

    Args:
        context (Context): App runtime context.
        render_json (bool, optional): Output results as json.

    Raises:
        typer.Exit: No retrosynthesis engines were returned.
    """
    # request retrosynthesis engines from the api client; sort the items
    # by their aliveness and then by putting the default engine first if
    # one is returned.
    controller = RetroController(context)
    items: list[RetrosynthesisEngine] = controller.retrieve_retrosynthesis_engines()
    items.sort(key=lambda x: x.last_alive, reverse=True)
    items.sort(key=lambda x: x.default, reverse=True)

    # capture the edge case where no engines exist from the response by
    # the api client; give a non-zero exit status.
    if len(items) == 0:
        console.print("[warn]! No retrosynthesis engines available.")
        raise typer.Exit(1)

    # output the retrosynthesis engine data either as json data if the
    # --json flag is given, else as a table output.
    if render_json:
        console.print_json(data=[x.model_dump(mode="json") for x in items])
    else:
        table = rich.table.Table(
            rich.table.Column("ID", style="b"),
            rich.table.Column("Name"),
            rich.table.Column("Last Alive", style="dim", justify="right"),
            box=rich.table.box.SQUARE,
        )
        for item in items:
            # preprocess the last alive timestamp to get the remaining
            # seconds from the current datetime, requires localizing the
            # database datetime into utc.
            item.last_alive = pytz.utc.localize(item.last_alive, is_dst=True)
            timestamp: float = (
                datetime.now(timezone.utc) - item.last_alive
            ).total_seconds()
            table.add_row(
                item.id + (" (default)" if item.default else ""),
                item.name,
                (f"{timestamp:3.0f}" if timestamp < 60 else "Over 60")
                + " second(s) ago",
            )
        console.print(table)


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


@app.command(
    "libraries",
    help=(
        "List available building block libraries. Their unique IDs are required "
        "for submitting jobs. If no libraries are available please contact support."
    ),
    short_help="List available building block libraries.",
)
def retrieve_libraries(
    context: Context,
    render_json: typing.Annotated[
        bool,
        typer.Option(
            "--json",
            help="Render output as JSON.",
            is_flag=True,
        ),
    ] = False,
) -> None:
    """
    Retrieve a list of available building block libraries stored in the
    database and output with relevant information for the user.

    Args:
        context (Context): App runtime context.
        render_json (bool, optional): Output results as json.

    Raises:
        typer.Exit: No building block libraries were returned.
    """
    # request building block libraries from the api client; sort items
    # by their name alphabetically.
    controller = RetroController(context)
    items: list[BuildingBlockLibrary] = controller.retrieve_building_block_libraries()
    items.sort(key=lambda x: x.name, reverse=True)

    # capture the edge case where no libraries exist from the response
    # by the api client; give a non-zero exit status.
    if len(items) == 0:
        console.print("[warn]! No building block libraries available.")
        raise typer.Exit(1)

    # output the building block library data either as json data if the
    # --json flag is given, else as a table output.
    if render_json:
        console.print_json(data=[x.model_dump(mode="json") for x in items])
    else:
        table = rich.table.Table(
            rich.table.Column("ID", style="b"),
            rich.table.Column("Name"),
            rich.table.Column("Version", style="dim"),
            box=rich.table.box.SQUARE,
        )
        for item in items:
            table.add_row(item.id, item.name, item.version)
        console.print(table)


# region command: tag --------------------------------------------------


def list_tags_callback(context: Context, list_tags: bool = False) -> None:
    """
    List tags belonging the the user as a list. Catch the edge case
    where the user has no retrosynthesis job tags.

    Args:
        context (Context): App runtime context.
        list_tags (bool, optional): Eager `--list` option flag.

    Raises:
        typer.Exit: Exit application once eager callback completes.
    """
    controller = RetroController(context)
    if list_tags:
        # retrieve tags and respond with non-zero status if no tags are
        # returned, otherwise print feedback for the user.
        tags: list[str] = controller.retrieve_tags()
        if len(tags) == 0:
            console.print("[warn]! No retrosynthesis job tags found.")
            raise typer.Exit(1)
        else:
            console.print(f"[success]✓ Found {len(tags)} job tag(s):")
            for tag in tags:
                console.print(f"[success]- [not b]{tag}")
        raise typer.Exit(0)


@app.command(
    "tag",
    help=(
        "Tag-based bulk operations. Perform collective operations on a "
        "job tag such as archiving or counting the number of jobs."
    ),
    short_help="Tag-based bulk operations.",
)
def tag_operations(
    context: Context,
    tag: typing.Annotated[
        str,
        typer.Argument(
            metavar="TAG",
            help="The operation is carried out on this single tag.",
        ),
    ],
    list_tags: typing.Annotated[
        bool,
        typer.Option(
            "--list",
            help="List all job tags then exit.",
            callback=list_tags_callback,
            is_flag=True,
            is_eager=True,
        ),
    ] = False,
    delete_tagged_jobs: typing.Annotated[
        bool,
        typer.Option(
            "--delete",
            help="Delete jobs with matching tag.",
            is_flag=True,
        ),
    ] = False,
    count_tagged_jobs: typing.Annotated[
        bool,
        typer.Option(
            "--count",
            help="Count jobs with matching tag.",
            is_flag=True,
        ),
    ] = False,
) -> None:
    """
    Peform a bulk operation on a retrosynthesis job tag. Operations will
    also be completed where a job contains more than one tag.

    Args:
        context (Context): App runtime context.
        tag (str): Tag the operation is performed on.
        list_tags (bool, optional): List all tags and exit.
        delete_tagged_jobs (bool, optional): Delete all jobs containing
            the specified tag.
        count_tagged_jobs (bool, optional): Count all jobs containing
            the specified tag.
    """
    controller = RetroController(context)
    with console.status(f"Finding jobs with tag: {tag}"):
        num_jobs: int = controller.retrieve_tag_count(tag)

    if delete_tagged_jobs:
        # confirm with the user that they want to delete the same number
        # of jobs and then perform the delete operation; exit with zero
        # status for success.
        console.print(f"[warn]! Warning: Deleting {num_jobs} retrosynthesis job(s).")
        console.print("[warn]- There may also be jobs matching different tags.")
        console.print("[warn]- This action cannot be undone.")
        if rich.prompt.Confirm.ask("[warn]! Are you sure?", console=console):
            deleted: int = controller.delete_tag(tag)
            console.print(f"[success]✓ Deleted {deleted} retrosynthesis job(s): '{tag}'")
        raise typer.Exit(0)

    elif count_tagged_jobs:
        # output the number of retrosynthesis jobs with matching tag and
        # exit with zero status for success.
        console.print(f"[success]✓ Found {num_jobs} job(s) with tag: '{tag}'")
        raise typer.Exit(0)
