import json
import logging
import sys
from pathlib import Path
from typing import List

import click

from highlighter.cli.cli_logging import ColourStr
from highlighter.cli.trainer_registry import select_trainer
from highlighter.client import HLJSONEncoder, TrainingRunType, json_tools
from highlighter.client.gql_client import HLClient
from highlighter.client.training_runs import (
    TrainingRunArtefactType,
    get_training_run,
    update_training_run,
)
from highlighter.core.utilities import change_dir
from highlighter.trainers import _scaffold
from highlighter.trainers.base_trainer import BaseTrainer

logger = logging.getLogger(__name__)


def _load_training_run(client: HLClient, training_run_dir: Path) -> TrainingRunType:
    training_run_yaml = training_run_dir / "training_run.yaml"
    if training_run_yaml.exists():
        training_run = TrainingRunType.from_yaml(training_run_yaml)
        update_training_run(
            client,
            training_run.id,
            name=training_run.name,
            trainer_config=training_run.trainer_config,
        )
        return training_run

    training_run_id = int(training_run_dir.stem)
    training_run = get_training_run(client, training_run_id)
    training_run.dump_yaml(training_run_yaml)
    return training_run


def _get_trainer(client: HLClient, training_run_dir: Path):
    training_run = _load_training_run(client, training_run_dir)

    # backwards compat
    trainer_type_file = training_run_dir / ".hl" / "trainer-type"
    if trainer_type_file.exists():
        with trainer_type_file.open("r") as f:
            trainer_type = f.readline().strip()
        trainer_name, trainer_cls = select_trainer(trainer_type)
    else:
        trainer_name, trainer_cls = select_trainer(training_run.trainer_config_file)

    if (training_run_dir / "trainer.py").exists():
        Trainer = _scaffold.load_trainer_module(training_run_dir).Trainer
    else:
        Trainer = trainer_cls

    if Trainer is None:
        logger.error(f"Unable to determine trainer from '{trainer_name}'")
        sys.exit(1)
    return Trainer(training_run_dir, training_run)


@click.group("train")
@click.pass_context
def train_group(ctx):
    pass


def _validate_training_run_dir(training_run_dir: Path):
    if not training_run_dir.exists():
        logger.error(f"training_run_dir {training_run_dir} does not exist.")
        sys.exit(1)

    try:
        int(training_run_dir.stem)
        return training_run_dir
    except ValueError:
        logger.error(
            f"Invalid training_run_dir {training_run_dir}. Should be 'ml_training/<TRAINING_RUN_ID>'"
        )
        sys.exit(1)


def _format_artefact_path_for_cli(artefact_path: Path) -> Path:
    return (
        artefact_path.relative_to(Path.cwd()) if artefact_path.is_relative_to(Path.cwd()) else artefact_path
    )


@train_group.command("start")
@click.argument("training-run-dir", required=False, default=".")
@click.pass_context
def train_start(ctx, training_run_dir):
    """Start model training and generate evaluation metrics.

    Executes the full training pipeline including data preparation, model training,
    and evaluation metric generation. Outputs structured evaluation metrics in JSON
    format for further processing.

    \b
    Args:
        training_run_dir: Directory containing training configuration (default: current directory)

    \b
    Returns:
        Prints training completion message and evaluation metrics in JSON format.
        Also provides command for manual artefact upload if needed.
    """

    client: HLClient = ctx.obj["client"]

    training_run_dir = Path(training_run_dir).absolute()
    training_run_dir = _validate_training_run_dir(training_run_dir)
    training_run_id = training_run_dir.stem

    trainer = _get_trainer(client, training_run_dir)

    combined_ds = trainer.get_datasets(client)

    if trainer.training_data_dir.exists():
        _scaffold.ask_to_remove(trainer.training_data_dir)

    with change_dir(training_run_dir):
        _, artefact_paths, _, eval_metric_results = trainer._train(combined_ds)

    click.echo(f"Training {training_run_id} complete")
    click.echo(
        f"{json.dumps({k: e.model_dump() for k, e in eval_metric_results.items()}, cls=json_tools.HLJSONEncoder)}"
    )
    click.echo("Next run:")
    for artefact_path in artefact_paths:
        apath = _format_artefact_path_for_cli(artefact_path)
        cmd = ColourStr.green(f"hl training-run artefact create -i {training_run_id} -a {apath}")
        click.echo(f"`{cmd}`")


@train_group.command("evaluate")
@click.argument("training-run-dir", required=False, default=".")
@click.argument("checkpoint", required=True, type=click.Path(dir_okay=False, exists=True))
@click.argument("config", required=True, type=click.Path(dir_okay=False))
@click.option("--create/--no-create", required=False, type=bool, default=False)
@click.pass_context
def train_evaluate(ctx, training_run_dir, checkpoint, config, create):
    """Evaluate a trained model checkpoint and output structured metrics.

    Evaluates a model checkpoint against validation data and generates evaluation
    metrics in JSON format. Does not upload results to Highlighter - use 'export'
    command for automated upload.

    \b
    Args:
        training_run_dir: Directory containing training configuration (default: current directory)
        checkpoint: Path to model checkpoint file for evaluation
        config: Path to YOLO configuration file
        create: If set will create the metrics results in the Highlighter Research Plan (default: --no-create)

    \b
    Returns:
        Prints evaluation metrics in JSON format
    """
    training_run_dir = Path(training_run_dir).absolute()
    training_run_dir = _validate_training_run_dir(training_run_dir)
    training_run_id = training_run_dir.stem

    client: HLClient = ctx.obj["client"]
    trainer = _get_trainer(client, training_run_dir)

    eval_metric_results = trainer.evaluate(checkpoint, cfg_path=config)
    click.echo(f"{json.dumps(eval_metric_results, indent=4, cls=HLJSONEncoder)}")

    if create:
        for e in eval_metric_results.values():
            e.create(client)


@train_group.command("export")
@click.argument("training-run-dir", required=False, default=".")
@click.argument("checkpoint", required=True, type=click.Path(dir_okay=False, exists=True))
@click.argument("config", required=True, type=click.Path(dir_okay=False))
@click.pass_context
def train_export(ctx, training_run_dir, checkpoint, config):
    """Export model and automatically upload artefact with evaluation metrics.

    This is the main command for finalizing a training run. It exports the model
    to the appropriate format, creates and uploads the training artefact to
    Highlighter, evaluates the model, and automatically uploads all evaluation
    metrics. This streamlined approach eliminates the need for separate manual
    artefact upload steps.

    \b
    Args:
        training_run_dir: Directory containing training configuration (default: current directory)
        checkpoint: Path to model checkpoint file to export
        config: Path to YOLO configuration file

    \b
    Returns:
        Prints artefact upload confirmation with URLs to view results in Highlighter
    """
    training_run_dir = Path(training_run_dir).absolute()
    training_run_dir = _validate_training_run_dir(training_run_dir)
    training_run_id = training_run_dir.stem

    client: HLClient = ctx.obj["client"]
    trainer = _get_trainer(client, training_run_dir)

    artefact_paths: List[Path] = trainer.export(Path(checkpoint))
    training_run = trainer.training_run
    created_artefact_ids = []
    for artefact_path in artefact_paths:
        artefact: TrainingRunArtefactType = TrainingRunArtefactType.from_yaml(artefact_path)
        create_artefact_result = artefact.create(client, training_run.id)
        created_artefact_ids.append(create_artefact_result.id)

    eval_metric_results = trainer.evaluate(checkpoint, cfg_path=config)

    # create the metric results in Highlighter
    for e in eval_metric_results.values():
        e.create(client)

    training_run_url = client.endpoint_url.replace("graphql", f"training_runs/{training_run.id}")
    evaluation_url = client.endpoint_url.replace("graphql", f"evaluations/{training_run.evaluation_id}")
    msg = ColourStr.green(
        f"Artefacts {created_artefact_ids} uploaded to {training_run_url}.\nSee metric results at {evaluation_url}"
    )
    click.echo(f"{msg}")
