"""SoloSeq folding submission command."""

from datetime import datetime
from pathlib import Path
from typing import Annotated

import typer
from rich.panel import Panel

from folding_studio.client import Client
from folding_studio.commands.utils import (
    success_fail_catch_print,
    success_fail_catch_spinner,
)
from folding_studio.config import FOLDING_API_KEY
from folding_studio.console import console
from folding_studio.query.soloseq import SoloSeqQuery


def soloseq(
    source: Annotated[
        Path,
        typer.Argument(
            help=(
                "Path to the data source. Either a fasta file or a directory of fasta files."
            ),
            exists=True,
        ),
    ],
    project_code: Annotated[
        str,
        typer.Option(
            help="Project code. If unknown, contact your PM or the Folding Studio team.",
            envvar="FOLDING_PROJECT_CODE",
            exists=True,
        ),
    ],
    seed: Annotated[int, typer.Option(help="Random seed.")] = 0,
    skip_relaxation: Annotated[
        bool,
        typer.Option(help="Skip the relaxation process."),
    ] = False,
    subtract_plddt: Annotated[
        bool,
        typer.Option(help="Output (100 - pLDDT) instead of the pLDDT itself."),
    ] = False,
    output: Annotated[
        Path,
        typer.Option(
            help="Local path to download the result zip and query parameters to. "
            "Default to 'soloseq_results'."
        ),
    ] = "soloseq_results",
    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="Unzip the file after its download.")
    ] = False,
    spinner: Annotated[
        bool, typer.Option(help="Use live spinner in log output.")
    ] = True,
):
    """Synchronous SoloSeq folding submission"""
    success_fail_catch = (
        success_fail_catch_spinner if spinner else success_fail_catch_print
    )
    console.print(
        Panel("[bold cyan]:dna: SoloSeq Folding submission [/bold cyan]", expand=False)
    )
    output_dir = output / f"submission_{datetime.now().strftime('%Y%m%d%H%M%S')}"

    # Create a client using API key or JWT
    with success_fail_catch(":key: Authenticating client"):
        client = Client.authenticate()

    # Define a query
    with success_fail_catch(":package: Generating query"):
        query_builder = (
            SoloSeqQuery.from_file if source.is_file() else SoloSeqQuery.from_directory
        )
        query: SoloSeqQuery = query_builder(
            source,
            seed=seed,
            skip_relaxation=skip_relaxation,
            subtract_plddt=subtract_plddt,
        )
        query.save_parameters(output_dir)

    console.print("[blue]Generated query: [/blue]")
    console.print_json(data=query.payload)

    # Send a request
    with success_fail_catch(":brain: Processing folding job"):
        response = client.send_request(query, project_code)

    # Access confidence data
    console.print("[blue]Confidence Data:[/blue]", end=" ")
    console.print_json(data=response.confidence_data)

    with success_fail_catch(
        f":floppy_disk: Downloading results to `[green]{output_dir}[/green]`"
    ):
        response.download_results(output_dir=output_dir, force=force, unzip=unzip)