"""CLI experiment command and sub-commands."""

import json
import os
import shutil
import zipfile
from pathlib import Path
from typing import List, Optional

import requests
import typer
from rich import print  # pylint:disable=redefined-builtin
from rich.console import Console
from rich.table import Table
from typing_extensions import Annotated

from folding_studio.config import API_URL, REQUEST_TIMEOUT
from folding_studio.utils.headers import get_auth_headers

app = typer.Typer(
    no_args_is_help=True,
    help=(
        "Get experiment information and metadata, like its status, "
        "results or the generated features (msa, templates, etc.).\n"
        "Read more at https://int-bio-foldingstudio-gcp.nw.r.appspot.com/how-to-guides/af2_openfold/fetch_folding_job_status/."
    ),
)

experiment_ID_argument = typer.Argument(help="ID of the experiment.")


def _download_file_from_signed_url(
    exp_id: str,
    endpoint: str,
    output: Path,
    force: bool,
    unzip: bool = False,
    api_key: str | None = None,
) -> None:
    """Download a zip file from an experiment id.

    Args:
        exp_id (str): Experiment id.
        endpoint (str): API endpoint to call.
        output (Path): Output file path.
        force (bool): Force file writing if it already exists.
        unzip (bool): Unzip the zip file after downloading.

    Raises:
        typer.Exit: If output file path exists but force set to false.
        typer.Exit: If unzip set to true but the directory already exists and force set to false.
        typer.Exit: If an error occurred during the initial API call.
    """
    if output.exists() and not force:
        print(
            f"Warning: The file '{output}' already exists. Use the --force flag to overwrite it."
        )
        raise typer.Exit(code=1)

    if unzip:
        if not output.suffix == ".zip":
            print(
                "Error: The downloaded file is not a .zip file. Please ensure the correct file format."
            )
            raise typer.Exit(code=1)

        dir_path = output.with_suffix("")
        if dir_path.exists() and not force:
            print(
                f"Warning: The --unzip flag is raised but the directory '{dir_path}' "
                "already exists. Use the --force flag to overwrite it."
            )
            raise typer.Exit(code=1)

    headers = get_auth_headers(api_key)
    url = API_URL + endpoint

    response = requests.get(
        url,
        params={"experiment_id": exp_id},
        headers=headers,
        timeout=REQUEST_TIMEOUT,
    )
    if not response.ok:
        print(f"Failed to download the file: {response.content.decode()}.")
        raise typer.Exit(code=1)

    file_response = requests.get(
        response.json()["signed_url"],
        stream=True,
        timeout=REQUEST_TIMEOUT,
    )
    with output.open("wb") as f:
        file_response.raw.decode_content = True
        shutil.copyfileobj(file_response.raw, f)
    print(f"File downloaded successfully to {output}.")

    if unzip:
        dir_path.mkdir(parents=True, exist_ok=True)
        with zipfile.ZipFile(output, "r") as zip_ref:
            zip_ref.extractall(dir_path)
            print(f"Extracted all files to {dir_path}.")


@app.command()
def status(
    exp_id: Annotated[str, experiment_ID_argument],
    api_key: Annotated[str, typer.Option("--api-key", "-k")],
):
    """Get an experiment status."""
    headers = get_auth_headers(api_key)
    url = API_URL + "getExperimentStatus"
    response = requests.get(
        url,
        params={"experiment_id": exp_id},
        headers=headers,
        timeout=REQUEST_TIMEOUT,
    )

    if not response.ok:
        print(f"An error occurred : {response.content.decode()}")
        raise typer.Exit(code=1)

    message = response.json()
    print(message["status"])


@app.command()
def list(
    limit: Annotated[
        int,
        typer.Option(
            help=("Max number of experiment to display in the terminal."),
        ),
    ] = 100,
    output: Annotated[
        Optional[Path],
        typer.Option(
            "--output",
            "-o",
            help=(
                "Path to the file where the job metadata returned by the server are written."
            ),
        ),
    ] = None,
):  # pylint:disable=redefined-builtin
    """Get all your done and pending experiment ids. The IDs are provided in the order of submission, starting with the most recent."""
    headers = get_auth_headers()
    url = API_URL + "getDoneAndPendingExperiments"
    response = requests.get(
        url,
        headers=headers,
        timeout=REQUEST_TIMEOUT,
    )

    if not response.ok:
        print(f"An error occurred : {response.content.decode()}")
        raise typer.Exit(code=1)

    response_json = response.json()
    if output:
        with open(output, "w") as f:
            json.dump(response_json, f, indent=4)

    print(f"Done and pending experiments list written to [bold]{output}[/bold]")

    table = Table(title="Done and pending experiments")

    table.add_column("Experiment ID", justify="right", style="cyan", no_wrap=True)
    table.add_column("Status", style="magenta")

    total_exp_nb = 0
    for status, exp_list in response_json.items():
        total_exp_nb += len(exp_list)
        for exp in exp_list:
            table.add_row(exp, status)
        if limit < total_exp_nb:
            print(
                f"The table below is truncated to the last [bold]{limit}[/bold] submitted experiments. Increase '--limit' to see more."
            )
            if not output:
                print("Use '--output' to get the full list in file format.")
            else:
                print(f"See the full list in file format at [bold]{output}[/bold]")
            break

    console = Console()
    console.print(table)


@app.command()
def features(
    exp_id: Annotated[str, experiment_ID_argument],
    output: Annotated[
        Optional[Path],
        typer.Option(
            help="Local path to download the zip to. Default to '<exp_id>_features.zip'."
        ),
    ] = None,
    force: Annotated[
        bool,
        typer.Option(
            help=(
                "Forces the download to overwrite any existing file "
                "with the same name in the specified location."
            )
        ),
    ] = False,
    unzip: Annotated[
        bool, typer.Option(help="Automatically unzip the file after its download.")
    ] = False,
):
    """Get an experiment features."""
    if output is None:
        output = Path(f"{exp_id}_features.zip")

    _download_file_from_signed_url(
        exp_id=exp_id,
        endpoint="getZippedExperimentFeatures",
        output=output,
        force=force,
        unzip=unzip,
    )


@app.command()
def results(
    exp_id: Annotated[str, experiment_ID_argument],
    api_key: Annotated[str, typer.Option("--api-key", "-k")],
    output: Annotated[
        Optional[Path],
        typer.Option(
            help="Local path to download the zip to. Default to '<exp_id>_results.zip'."
        ),
    ] = None,
    force: Annotated[
        bool,
        typer.Option(
            help=(
                "Forces the download to overwrite any existing file "
                "with the same name in the specified location."
            )
        ),
    ] = False,
    unzip: Annotated[
        bool, typer.Option(help="Automatically unzip the file after its download.")
    ] = False,
):
    """Get an experiment results."""

    if output is None:
        output = Path(f"{exp_id}_results.zip")

    _download_file_from_signed_url(
        exp_id=exp_id,
        endpoint="getZippedExperimentResults",
        output=output,
        force=force,
        unzip=unzip,
        api_key=api_key,
    )


@app.command()
def cancel(
    exp_id: Annotated[List[str], experiment_ID_argument],
):
    """Cancel experiments job executions.

    You can pass one or more experiment id
    """
    headers = get_auth_headers()

    url = API_URL + "cancelJob"
    response = requests.post(
        url,
        data={"experiment_ids": exp_id},
        headers=headers,
        timeout=REQUEST_TIMEOUT,
    )

    if not response.ok:
        print(f"An error occurred : {response.content.decode()}")
        raise typer.Exit(code=1)
    message = response.json()
    print(message)


@app.command()
def logs(
    exp_id: Annotated[str, experiment_ID_argument],
    output: Annotated[
        Optional[Path],
        typer.Option(
            help="Local path to download the logs to. Default to '<exp_id>_logs.txt'."
        ),
    ] = None,
    force: Annotated[
        bool,
        typer.Option(
            help=(
                "Forces the download to overwrite any existing file "
                "with the same name in the specified location."
            )
        ),
    ] = False,
):
    """Get an experiment logs."""
    if output is None:
        output = Path(f"{exp_id}_logs.txt")

    _download_file_from_signed_url(
        exp_id=exp_id,
        endpoint="getExperimentLogs",
        output=output,
        force=force,
    )