import json
from pathlib import Path
from typing import Dict, List, Optional

import click
import yaml

from ..client import (
    HLJSONEncoder,
    TrainingRunArtefactType,
    TrainingRunArtefactTypeEnum,
    TrainingRunType,
    get_latest_training_run_artefact,
    get_training_run,
    get_training_run_artefact,
)
from ..client import create_training_run as create_training_run_client
from ..client import update_training_run as update_training_run_client
from ..core import DEPRECATED_CAPABILITY_IMPLEMENTATION_FILE, GQLBaseModel
from .common import _to_pathlib


def load_config(cfg_path: Optional[Path]) -> Dict:
    if cfg_path is None:
        return {}

    loaders = {
        ".yaml": yaml.safe_load,
        ".json": json.load,
    }

    with cfg_path.open("r") as f:
        cfg = loaders[cfg_path.suffix](f)

    if cfg.get("parameters", {}).get("cropper", {}).get("scale", 1.0) != 1.0:
        raise ValueError(
            "Got'cha!!!, Dont push an artefact with cropping scale other than 1.0 or remove from dict"
        )

    return cfg


@click.group("training-run")
@click.pass_context
def training_run_group(ctx):
    pass


@training_run_group.command("read")
@click.option("-i", "--id", "training_run_id", type=int, required=True, help="Training run ID")
@click.option(
    "-o",
    "--output",
    type=click.Path(dir_okay=False, writable=True),
    required=False,
    callback=_to_pathlib,
)
@click.pass_context
def read_training_run(ctx, training_run_id, output):
    """Download a training run and optionally save it to a local JSON/YAML file."""
    client = ctx.obj["client"]
    training_run = get_training_run(client, training_run_id)

    if output is None:
        click.echo(json.dumps(training_run.model_dump(), cls=HLJSONEncoder))
        return

    if output.suffix.lower() in [".yaml", ".yml"]:
        training_run.dump_yaml(output)
    elif output.suffix.lower() == ".json":
        output.write_text(json.dumps(training_run.model_dump(), cls=HLJSONEncoder, indent=2))
    else:
        raise click.ClickException("Output path must end with .yaml, .yml, or .json")


@training_run_group.group("artefact")
@click.pass_context
def artefact_group(ctx):
    pass


@artefact_group.command("read")
@click.option(
    "-i",
    "--id",
    type=str,
    required=False,
)
@click.option(
    "-s",
    "--save-path",
    type=click.Path(file_okay=True, writable=True),
    required=True,
    callback=_to_pathlib,
)
@click.option(
    "-a",
    "--artefact-type",
    type=click.Choice(
        list(TrainingRunArtefactTypeEnum.__members__.keys()) + [DEPRECATED_CAPABILITY_IMPLEMENTATION_FILE]
    ),
    required=False,
)
@click.option(
    "-u",
    "--artefact-uuid",
    type=str,
    required=False,
    help="Direct training run artefact UUID. If set, --id and --artefact-type are not required.",
)
@click.pass_context
def read_artefact(ctx, id, save_path, artefact_type, artefact_uuid):
    from ..licenses.sam3_license import is_sam3_artefact, require_acceptance, write_license_file

    _is_sam3 = artefact_uuid is not None and is_sam3_artefact(artefact_uuid)

    if _is_sam3:
        require_acceptance()

    client = ctx.obj["client"]

    if artefact_uuid is not None:
        artefact = get_training_run_artefact(
            client,
            artefact_uuid,
            download_file_url=True,
            file_url_save_path=save_path,
        )
    else:
        if id is None or artefact_type is None:
            raise click.ClickException("Provide either --artefact-uuid, or both --id and --artefact-type.")

        artefact = get_latest_training_run_artefact(
            client,
            id,
            download_file_url=True,
            file_url_save_path=save_path,
            filter_by_artefact_type=artefact_type,
        )
    assert artefact is not None

    if _is_sam3:
        license_path = write_license_file(save_path)
        click.echo(f"SAM License written to {license_path}")

    artefact_yaml_path = save_path.parent / f"{artefact.id}.yaml"
    artefact.dump_yaml(artefact_yaml_path)
    click.echo(json.dumps(artefact.model_dump(), cls=HLJSONEncoder))


@artefact_group.command("create")
@click.option(
    "-i",
    "--id",
    type=str,
    required=True,
    help="training run id",
)
@click.option(
    "-a",
    "--artefact-yaml",
    type=click.Path(file_okay=True, dir_okay=False, exists=True),
    required=True,
    callback=_to_pathlib,
    help=".yaml file base_model.py:TrainingRunArtefactType",
)
@click.pass_context
def create_artefact(ctx, id, artefact_yaml):
    """Create an artefact for a training run and upload artefact-yaml.file_url to s3.

    \b
    Fields in artefact.yaml:
    - file_url: should contain the absolute path to the checkpoint file you want to upload.
    - type: [REQUIRED] one of OnnxOpset11, OnnxOpset14, OnnxRuntimeAmd64, OnnxRuntimeArm, PyTorch, TorchScriptV1
    - checkpoint: path to checkpoint in file
    - inference_config: [REQUIRED] inference configuration in json format

    \b
    # artefact.yaml
    file_url: /home/users/rick/checkpoint.onnx14
    type: OnnxOpset14
    inference_config: {}
    """
    client = ctx.obj["client"]

    artefact: TrainingRunArtefactType = TrainingRunArtefactType.from_yaml(artefact_yaml)
    artefact.create(client, id)


@training_run_group.command("create")
@click.option(
    "-eid",
    "--evaluation-id",
    type=str,
    required=True,
)
@click.option(
    "-xid",
    "--experiment-id",
    type=str,
    required=True,
)
@click.option(
    "-mid",
    "--capability-id",
    type=str,
    required=True,
)
@click.option(
    "-pid",
    "--workflow-id",
    type=str,
    required=True,
)
@click.option(
    "-n",
    "--name",
    type=str,
    required=True,
)
@click.option(
    "-d",
    "--description",
    type=str,
    required=False,
)
@click.option(
    "--convert-description-to-html",
    is_flag=True,
    default=False,
    help="Convert description markdown to HTML on save.",
)
@click.option(
    "--source-code-url",
    type=str,
    required=False,
)
@click.option(
    "--source-code-commit-hash",
    type=str,
    required=False,
)
@click.option(
    "--training-log-archive",
    type=click.Path(dir_okay=False),
    callback=_to_pathlib,
    required=False,
)
@click.pass_context
def create_training_run(
    ctx,
    evaluation_id,
    experiment_id,
    capability_id,
    workflow_id,
    name,
    description,
    convert_description_to_html,
    source_code_url,
    source_code_commit_hash,
    training_log_archive,
):
    client = ctx.obj["client"]

    result = create_training_run_client(
        client,
        evaluation_id=int(evaluation_id),
        experiment_id=int(experiment_id),
        model_id=int(capability_id),
        workflow_id=int(workflow_id),
        name=name,
        description=description,
        convert_description_to_html=convert_description_to_html,
        source_code_url=source_code_url,
        source_code_commit_hash=source_code_commit_hash,
    )
    print(result.id)


@training_run_group.command("update")
@click.option(
    "-i",
    "--id",
    "training_run_id",
    type=int,
    required=True,
    help="Training run ID.",
)
@click.option(
    "-n",
    "--name",
    type=str,
    required=False,
)
@click.option(
    "-d",
    "--description",
    type=str,
    required=False,
)
@click.option(
    "--convert-description-to-html",
    is_flag=True,
    default=False,
    help="Convert description markdown to HTML on save.",
)
@click.option(
    "--source-code-url",
    type=str,
    required=False,
)
@click.option(
    "--source-code-commit-hash",
    type=str,
    required=False,
)
@click.option(
    "--trainer-config",
    type=str,
    required=False,
    help="Trainer config JSON object string.",
)
@click.pass_context
def update_training_run(
    ctx,
    training_run_id,
    name,
    description,
    convert_description_to_html,
    source_code_url,
    source_code_commit_hash,
    trainer_config,
):
    client = ctx.obj["client"]
    trainer_config_obj = None
    if trainer_config is not None:
        try:
            trainer_config_obj = json.loads(trainer_config)
        except json.JSONDecodeError as e:
            raise click.ClickException(f"Invalid --trainer-config JSON: {e}") from e

    result = update_training_run_client(
        client,
        training_run_id,
        name=name,
        description=description,
        convert_description_to_html=convert_description_to_html,
        source_code_url=source_code_url,
        source_code_commit_hash=source_code_commit_hash,
        trainer_config=trainer_config_obj,
    )
    print(result.id)


@training_run_group.command("delete")
@click.option("-i", "--id", type=str, required=True, help="The ID of the training run to delete")
@click.option("-y", "--yes", is_flag=True, help="Skip confirmation prompt")
@click.pass_context
def delete(ctx, id, yes):
    """Delete a training run"""
    if not yes:
        if not click.confirm(f"Are you sure you want to delete training run {id}?"):
            click.echo("Aborted.")
            return
    client = ctx.obj["client"]

    class DeletedTrainingRun(GQLBaseModel):
        id: int

    class DeleteTrainingRunPayload(GQLBaseModel):
        errors: List[str]
        training_run: Optional[DeletedTrainingRun] = None

    result = client.deleteTrainingRun(return_type=DeleteTrainingRunPayload, id=id).gql_dict()
    click.echo(json.dumps(result, cls=HLJSONEncoder))
