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

from __future__ import annotations

import json
import re
import unicodedata
from datetime import datetime
from typing import Annotated, Any, Dict, Generator, List

from pendingai.commands import append_client_context
from pendingai.commands.retrosynthesis.controller import Retrosynthesis
from pendingai.commands.retrosynthesis.models import (
    Engine,
    Job,
    JobPage,
    Library,
    NewJob,
)
from pendingai.commands.retrosynthesis.types import Options
from rich import print, print_json
from rich.console import Console
from rich.prompt import Confirm
from rich.status import Status
from rich.table import Table, box
from typer import Argument, BadParameter, Context, Exit, Option, Typer

app: Typer = Typer(context_settings={"max_content_width": 90})


@app.callback()
def _callback(ctx: Context) -> None:
    """
    High-throughput deep-learning retrosynthesis service.

    Retrosynthesis 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>.
    """


@app.command("engines")
def _engines(
    ctx: Context,
    render_json: Options.RenderJson = False,
) -> None:
    """
    List available retrosynthesis engines. Their unique IDs are required
    for submitting jobs. If no engines are available, contact support
    with <pendingai support>.
    """
    append_client_context(ctx)
    ctx.obj.controller = Retrosynthesis(ctx)

    engines: List[Engine] = ctx.obj.controller.get_engines()
    engines.sort(key=lambda x: x.last_alive, reverse=True)
    engines.sort(key=lambda x: x.default, reverse=True)

    if len(engines) == 0:
        print("[yellow]! No retrosynthesis engine available.")
        raise Exit(1)

    if render_json:
        print_json(data=[engine.model_dump(mode="json") for engine in engines])

    else:
        table: Table = Table(box=box.SQUARE)
        table.add_column("ID", style="bold")
        table.add_column("Name")
        table.add_column("Last Alive", style="dim", justify="right")
        for engine in engines:
            ts: float = (datetime.utcnow() - engine.last_alive).seconds
            table.add_row(
                engine.id + (" (default)" if engine.default else ""),
                engine.name,
                (f"{ts:3d}" if ts < 60 else ">60") + " seconds ago",
            )
        print(table)


@app.command("libraries")
def _libraries(
    ctx: Context,
    render_json: Options.RenderJson = False,
) -> None:
    """
    List available building block libraries. Their unique IDs are required
    for submitting jobs. If no libraries are available, contact support
    with <pendingai support>.
    """
    append_client_context(ctx)
    ctx.obj.controller = Retrosynthesis(ctx)

    libraries: List[Library] = ctx.obj.controller.get_libraries()
    libraries.sort(key=lambda x: x.name)

    if len(libraries) == 0:
        print("[yellow]! No building block libraries available.")
        raise Exit(1)

    if render_json:
        print_json(data=[library.model_dump(mode="json") for library in libraries])

    else:
        table: Table = Table(box=box.SQUARE)
        table.add_column("ID", style="bold")
        table.add_column("Name")
        table.add_column("Version", style="dim")
        for library in libraries:
            table.add_row(library.id, library.name, library.version)
        print(table)


def _sanitize_filename(pth: str) -> str:
    """
    TODO: Finish documentation.
    """
    if "." in pth:
        pth = ".".join(pth.split(".")[:-1])
    pth = unicodedata.normalize("NFKD", pth).encode("ascii", "ignore").decode("ascii")
    pth = re.sub(r"[^\.\w\s-]", "", pth.lower())
    return re.sub(r"[-\.\s]+", "_", pth).strip("-_")


@app.command("submit", no_args_is_help=True)
def _submit(
    ctx: Context,
    output_file: Options.OutputIdFile,
    targets_list: Options.InputTargetsList = [],
    targets_file: Options.InputTargetsFile = None,
    eng_id: Options.Parameters.EngineId = None,
    lib_ids: Options.Parameters.LibraryIds = None,
    number_of_routes: Options.Parameters.NumRoutes = 20,
    processing_time: Options.Parameters.ProcessingTime = 300,
    reaction_limit: Options.Parameters.ReactionLimit = 8,
    building_block_limit: Options.Parameters.BlockLimit = 8,
    job_tags: Options.Parameters.JobTags = [],
) -> None:
    """
    Submit jobs to retrosynthesis engine. Writes IDs for submitted jobs
    to output file. Jobs are automatically tagged with submission
    time and output filename.

    WARNING: Jobs including their results will be automatically deleted
    30 days after completion.

    \b
    Tags and building block libraries are repeatable:
      <pendingai retro submit --output-file job1.ids --tag tag1 --tag tag2>
      <pendingai retro submit --output-file job1.ids --lib-id lib1 --lib-id lib2>
    """
    append_client_context(ctx)
    ctx.obj.controller = Retrosynthesis(ctx)

    console: Console = Console(quiet=False, highlighter=None)
    targets: List[str] = targets_list + (targets_file if targets_file else [])

    # Validate input values; targets must be provided and not exceed an
    # upper-limit to not overload the system.
    if len(targets) == 0:
        raise BadParameter("At least 1 input structure is required for retrosynthesis.")
    if len(targets) > 10_000:
        raise BadParameter("Exceeded request limit of 10,000.")
    console.print(f"[green]\u2713[/] Validated {len(targets)} retrosynthesis target(s).")

    job_tags.append(datetime.now().strftime("%Y%m%d-%H%M%S"))
    console.print(f'[green]\u2713[/] Added additional batch tag "{job_tags[-1]}"')
    job_tags.append(_sanitize_filename(output_file.name)[:16])
    console.print(f'[green]\u2713[/] Added additional batch tag "{job_tags[-1]}"')

    # Results are streamed back to the user in order to handle batched
    # requests and partial results. That way any incomplete requests are
    # still returned and end on a failed batch submission.
    results: Generator[List[NewJob], Any, None] = ctx.obj.controller.submit_molecules(
        targets=targets,
        engine=eng_id,
        libraries=lib_ids,
        number_of_routes=number_of_routes,
        processing_time=processing_time,
        reaction_limit=reaction_limit,
        building_block_limit=building_block_limit,
        tags=job_tags,
        console=console,
    )

    with open(output_file, "w") as fp:
        for batch in results:
            fp.writelines([result.id + "\n" for result in batch])


@app.command("delete", no_args_is_help=True)
def _delete(
    ctx: Context,
    ids_list: Options.InputIdList = [],
    ids_file: Options.InputIdFile = None,
    output_file: Options.OutputFile = None,
    render_json: Options.RenderJson = False,
) -> None:
    """
    Delete one or more retrosynthesis jobs by ID. Deletion is irreversible,
    completed in order, and will skip any missing or failed IDs.

    WARNING: Jobs including their results will be automatically deleted 30 days
    after completion.
    """
    append_client_context(ctx)
    ctx.obj.controller = Retrosynthesis(ctx)

    if output_file is not None:
        output_file = open(output_file, "w")  # type: ignore

    console: Console = Console(quiet=render_json, highlighter=None)
    ids: List[str] = ids_list + (ids_file if ids_file else [])

    # Validate input values; targets must be provided and not exceed an
    # upper-limit to not overload the system.
    if len(ids) == 0:
        raise BadParameter("At least 1 job ID is required for deletion.")
    if len(ids) > 1000:
        raise BadParameter("Exceeded request limit of 1000.")

    # If the number of targets exceeds a limit of 25, there are too many
    # results to display in the terminal, this is then saved to a local
    # temporary filepath.
    if len(ids) > 25 and output_file is None:
        console.quiet = False
        path: str = f"delete_{datetime.now().strftime('%Y-%m-%dT%H:%M:%S')}.json"
        print(f"[yellow not b]! Over 25 IDs requested, saving results to file: {path}")
        output_file = open(path, "w")  # type: ignore

    results: Dict = ctx.obj.controller.delete_jobs(ids=ids, console=console)

    if output_file:
        json.dump(results, output_file, indent=2)  # type: ignore

    elif render_json:
        print_json(data=results)

    else:
        table: Table = Table("ID", "Status", box=box.SQUARE)
        [table.add_row(id, "[green]Success") for id in results["success"]]
        [table.add_row(id, "[red]Failed") for id in results["failure"]]
        [table.add_row(id, "[dim]Missing") for id in results["missing"]]
        print(table)


@app.command("status", no_args_is_help=True)
def _status(
    ctx: Context,
    ids_list: Options.InputIdList = [],
    ids_file: Options.InputIdFile = None,
    output_file: Options.OutputFile = None,
    render_json: Options.RenderJson = False,
) -> None:
    """
    Get the status of one or more retrosynthesis jobs by ID.
    """
    append_client_context(ctx)
    ctx.obj.controller = Retrosynthesis(ctx)

    if output_file is not None:
        output_file = open(output_file, "w")  # type: ignore

    console: Console = Console(quiet=render_json, highlighter=None)
    ids: List[str] = ids_list + (ids_file if ids_file else [])

    # Validate input values; targets must be provided and not exceed an
    # upper-limit to not overload the system.
    if len(ids) == 0:
        raise BadParameter("At least 1 job ID is required for checking a status.")
    if len(ids) > 1000:
        raise BadParameter("Exceeded request limit of 1000.")

    # If the number of targets exceeds a limit of 25, there are too many
    # results to display in the terminal, this is then saved to a local
    # temporary filepath.
    if len(ids) > 25 and output_file is None:
        console.quiet = False
        path: str = f"status_{datetime.now().strftime('%Y-%m-%dT%H:%M:%S')}.json"
        print(f"[yellow not b]! Over 25 IDs requested, saving results to file: {path}")
        output_file = open(path, "w")  # type: ignore

    results: List[Dict] = ctx.obj.controller.get_job_status(ids=ids, console=console)

    if output_file:
        json.dump(results, output_file, indent=2)  # type: ignore

    elif render_json:
        print_json(data=results)

    else:
        table: Table = Table("ID", "Status", box=box.SQUARE)
        for result in results:
            if result["status"] == "completed":
                result["status"] = "[green]Completed"
            elif result["status"] == "submitted":
                result["status"] = "Submitted"
            elif result["status"] == "processing":
                result["status"] = "[yellow]Processing"
            elif result["status"] == "missing":
                result["status"] = "[dim]Missing"
            else:
                result["status"] = "[red]Failed"
            table.add_row(result["id"], result["status"])
        print(table)


@app.command("view", no_args_is_help=True)
def _view(
    ctx: Context,
    ids_list: Options.InputIdList = [],
    ids_file: Options.InputIdFile = None,
    output_file: Options.OutputFile = None,  # type: ignore
    render_json: Options.RenderJson = False,
) -> None:
    """
    Retrieve retrosynthesis job results by ID. IDs are processed
    sequentially. Retrosynthetic routes are only returned in JSON output.
    """
    append_client_context(ctx)
    ctx.obj.controller = Retrosynthesis(ctx)

    if output_file is not None:
        output_file = open(output_file, "w")  # type: ignore

    console: Console = Console(quiet=render_json, highlighter=None)
    ids: List[str] = ids_list + (ids_file if ids_file else [])

    # Validate input values; targets must be provided and not exceed an
    # upper-limit to not overload the system.
    if len(ids) == 0:
        raise BadParameter("At least 1 job ID is required for retrieving results.")
    if len(ids) > 10_000:
        raise BadParameter("Exceeded request limit of 10,000.")

    # If the number of targets exceeds a limit of 25, there are too many
    # results to display in the terminal, this is then saved to a local
    # temporary filepath.
    if len(ids) > 25 and output_file is None:
        console.quiet = False
        path: str = f"results_{datetime.now().strftime('%Y-%m-%dT%H:%M:%S')}.json"
        print(f"[yellow not b]! Over 10 IDs requested, saving results to file: {path}")
        output_file = open(path, "w")  # type: ignore

    results: Generator[Job | Dict, Any, None] = ctx.obj.controller.get_jobs_by_id(
        ids=ids,
        console=console,
    )

    # Special file writing is needed to comply with JSON formatting and
    # can handle different response formats depending on if the job ID
    # has existing results or does not.
    if output_file:
        output_file.write("[\n\t")  # type: ignore
        for i, result in enumerate(results):
            if i > 0:
                output_file.write(",\n\t")  # type: ignore
            try:
                output_file.write(result.model_dump_json())  # type: ignore
            except Exception:
                output_file.write(json.dumps(result))  # type: ignore
        output_file.write("\n]")  # type: ignore

    elif render_json:
        print_json(
            data=[
                x.model_dump(mode="json") if not isinstance(x, dict) else x
                for x in results
            ]
        )

    else:
        table: Table = Table(box=box.SQUARE)
        table.add_column("ID", width=24)
        table.add_column("Structure", width=50, style="dim")
        table.add_column("Status")
        table.add_column("Routes", justify="right")
        results_: list = [i for i in results]
        has_tabular_view: bool = False
        for result in results_:
            if isinstance(result, dict):
                print(f"[yellow not b]! Job is not completed or missing: {result['id']}")
            else:
                has_tabular_view = True
                table.add_row(
                    result.id,
                    result.query,
                    result.status.title(),
                    str(len(result.routes)) if len(result.routes) else "0",
                )
        print("[yellow]! Use --json to retrieve detailed retrosynthetic routes.")
        if has_tabular_view:
            print(table)


@app.command("list")
def _list(
    ctx: Context,
    page: Options.Pagination.PageNumber = 1,
    page_size: Options.Pagination.PageSize = 10,
    status: Options.Pagination.Status = None,
    tags: Options.Parameters.Tags = [],
    render_json: Options.RenderJson = False,
) -> None:
    """
    Retrieve a page of retrosynthesis jobs.
    Displays select search parameters and a binary synthesizability flag.
    Use <pendingai retro view> to retrieve retrosynthetic routes of
    completed jobs.

    \b
    Tags are repeatable:
      <pendingai retro list --tag tag1 --tag tag2>
    """
    append_client_context(ctx)
    ctx.obj.controller = Retrosynthesis(ctx)

    console: Console = Console(quiet=render_json, highlighter=None)
    page_data: JobPage = ctx.obj.controller.get_paginated_jobs(
        page=page,
        page_size=page_size,
        status=status.value if status else None,
        tags=tags,
    )

    if page_data.limit > 0 and page > page_data.limit:
        console.print(f"[yellow]! Page limit of {page_data.limit} exceeded.")
        raise Exit(1)
    elif len(page_data.results) == 0:
        console.print("[yellow]! No results found with the given filter(s).")
        raise Exit(1)

    if render_json:
        print_json(data=[x.model_dump(mode="json") for x in page_data.results])

    else:
        table: Table = Table(box=box.SQUARE, caption=f"Page {page} of {page_data.limit}")
        table.add_column("ID", style="bold green")
        table.add_column("Structure", width=50)
        table.add_column("Submitted At", style="dim")
        table.add_column("Status", style="cyan")
        table.add_column("Engine", width=15, style="dim")
        table.add_column("Libraries", width=15, style="dim", no_wrap=True)
        table.add_column("Tags", width=20)
        table.add_column("Synthesizable", style="bold cyan")
        for result in page_data.results:
            table.add_row(
                result.id,
                result.query,
                result.created.strftime("%Y-%m-%d %H:%M:%S"),
                result.status.title(),
                result.parameters.retrosynthesis_engine,
                ", ".join(result.parameters.building_block_libraries),
                ", ".join(result.tags) if result.tags else "",
                str(len(result.routes) > 0) if result.status == "completed" else "",
            )
        print(table)


def _tag_list_callback(ctx: Context, value: bool) -> None:
    if value:
        append_client_context(ctx)
        ctx.obj.controller = Retrosynthesis(ctx)
        with Status("Retrieving all job tags"):
            tag_ids: list[str] = ctx.obj.controller.get_job_tags()

        if len(tag_ids) == 0:
            print("[yellow]! No job tags found.")
        else:
            print(f"Retrieved {len(tag_ids)} job tags:")
            for tag_id in tag_ids:
                print(f"- '{tag_id}'")
        raise Exit(0)


@app.command("tag")
def _tag(
    ctx: Context,
    tag_id: Annotated[
        str,
        Argument(
            metavar="TAG",
            help="The operations is carried out on this single tag.",
        ),
    ],
    list_tags: Annotated[
        bool,
        Option(
            "--list",
            help="List all job tags.",
            callback=_tag_list_callback,
            is_eager=True,
        ),
    ] = False,
    bulk_delete_tag_flag: Annotated[
        bool,
        Option(
            "--delete",
            help="Delete jobs with matching tag.",
            is_flag=True,
        ),
    ] = False,
    retrieve_tag_count_flag: Annotated[
        bool,
        Option(
            "--count",
            help="Count jobs with matching tag.",
            is_flag=True,
        ),
    ] = False,
) -> None:
    """
    Tag-based operations.
    """
    append_client_context(ctx)
    ctx.obj.controller = Retrosynthesis(ctx)

    console = Console(highlighter=None)

    # Display a wait status while performing the count operation
    with Status(f"Finding jobs with tag: '{tag_id}'"):
        count: int = ctx.obj.controller.get_job_tag_count(tag_id)

    if bulk_delete_tag_flag:
        # Confirm the user wants to delete all jobs with the tag
        print(f"[yellow not b]! Warning: Deleting {count} jobs tagged with '{tag_id}'.")
        print("[yellow]- This action cannot be undone.")
        if Confirm.ask("[yellow]! Are you sure?"):
            # Run the delete operation and output the number deleted
            with Status(f"Deleting jobs with tag: '{tag_id}'"):
                deleted: int = ctx.obj.controller.delete_job_tag(tag_id)
            console.print(f"[green]✓[/] Deleted {deleted} jobs with tag: '{tag_id}'")
        raise Exit(0)

    if retrieve_tag_count_flag:
        # Output the retrieved count
        console.print(f"[green]✓[/] Found {count} job(s) with tag: '{tag_id}'")
        raise Exit(0)
