import json
import logging
import tempfile
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Union
from uuid import UUID

import yaml
from pydantic import AliasChoices, BaseModel, BeforeValidator, ConfigDict, Field, field_validator
from typing_extensions import Annotated

from ..core import (
    DEPRECATED_CAPABILITY_IMPLEMENTATION_FILE,
    GQLBaseModel,
    snake_to_camel,
)
from ..datasets.cropping import CropArgs
from .aws_s3 import upload_file_to_s3
from .gql_client import HLClient
from .io import download_bytes

logger = logging.getLogger(__name__)

__all__ = [
    "EntityAttributeValueTypeEnum",
    "EntityAttributeValue",
    "EntityAttribute",
    "ModelOutputType",
    "ModelInputs",
    "ModelSchemaType",
    "DatasetType",
    "ArtefactInfo",
    "TaxonPart",
    "TrainingRunArtefactTypeEnum",
    "TrainingRunArtefactType",
    "TrainingRunType",
    "get_training_run",
    "get_training_run_artefacts",
    "get_training_run_artefact",
    "get_latest_training_run_artefact",
    "get_capability_implementation_file_url",
    "create_training_run",
    "create_training_run_artefact",
    "update_training_run",
]


def _get_now_str():
    return datetime.now().isoformat()


class TrainingRunArtefactTypeEnum(str, Enum):
    DeprecatedCapabilityImplementationFile = DEPRECATED_CAPABILITY_IMPLEMENTATION_FILE
    DeprecatedClassilvier = "DeprecatedClassilvier"
    DeprecatedMmpond = "DeprecatedMmpond"
    DeprecatedSilverclassify = "DeprecatedSilverclassify"
    OnnxOpset11 = "OnnxOpset11"
    OnnxOpset14 = "OnnxOpset14"
    OnnxRuntimeAmd64 = "OnnxRuntimeAmd64"
    OnnxRuntimeArm = "OnnxRuntimeArm"
    PyTorch = "PyTorch"
    TensorFlowV1 = "TensorFlowV1"
    TextEmbedder = "TextEmbedder"
    TorchScriptV1 = "TorchScriptV1"


def validate_uuid(v, key):
    validated = None
    if isinstance(v, UUID):
        validated = str(v)
    else:
        try:
            UUID(v)
            validated = v
        except:
            msg = f"'{key}' must be valid UUID, got: {v}"
            raise ValueError(msg)
    return validated


class EntityAttributeValueTypeEnum(str, Enum):
    boolean = "boolean"
    enum = "enum"
    integer = "integer"
    decimal = "decimal"
    string = "string"
    geometry = "geometry"
    array = "array"


class TaxonPart(BaseModel, extra="forbid"):
    entity_attribute_enum_id: Optional[Union[UUID, str]] = None
    entity_attribute_enum_value: Optional[str] = None
    entity_attribute_id: Union[UUID, str]
    entity_attribute_name: str
    entity_attribute_value_type: Optional[EntityAttributeValueTypeEnum] = None

    @field_validator("entity_attribute_enum_id")
    @classmethod
    def validate_entity_attribute_enum_id(cls, v):
        if v is None:
            return True
        else:
            return str(validate_uuid(v, "entity_attribute_enum_id"))

    @field_validator("entity_attribute_id")
    @classmethod
    def validate_entity_attribute_id(cls, v):
        return str(validate_uuid(v, "entity_attribute_id"))

    @property
    def value_type(self):
        return self.entity_attribute_value_type


class ModelOutputType(BaseModel, extra="forbid"):
    head: int
    position: int
    conf_thresh: Annotated[float, Field(ge=0.0, lt=1.0)]
    taxon: List[TaxonPart]

    def _validate_one_taxon(self, property_name):
        if len(self.taxon) > 1:
            raise ValueError(
                f"Can only use ModelOutputType.{property_name} if there is a single taxa in self.taxon"
            )

    @property
    def entity_attribute_enum_value(self):
        self._validate_one_taxon("entity_attribute_enum_value")
        return self.taxon[0].entity_attribute_enum_value

    @property
    def entity_attribute_enum_id(self) -> Optional[Union[UUID, str]]:
        self._validate_one_taxon("entity_attribute_enum_id")
        return self.taxon[0].entity_attribute_enum_id

    @property
    def entity_attribute_id(self) -> Union[UUID, str]:
        self._validate_one_taxon("entity_attribute_id")
        return self.taxon[0].entity_attribute_id

    @property
    def entity_attribute_name(self) -> str:
        self._validate_one_taxon("entity_attribute_name")
        return self.taxon[0].entity_attribute_name

    @property
    def entity_attribute_value_type(self) -> Optional[EntityAttributeValueTypeEnum]:
        self._validate_one_taxon("entity_attribute_value_type")
        return self.taxon[0].entity_attribute_value_type


class EntityAttribute(BaseModel, extra="forbid"):
    entity_attribute_id: Union[UUID, str]
    entity_attribute_name: str

    @field_validator("entity_attribute_id")
    @classmethod
    def validate_entity_attribute_id(cls, v):
        return str(validate_uuid(v, "entity_attribute_id"))


class EntityAttributeValue(BaseModel, extra="forbid"):
    entity_attribute_enum_id: Optional[Union[UUID, str]] = None
    entity_attribute_enum_value: Optional[str] = None
    entity_attribute_id: Union[UUID, str]
    entity_attribute_name: str
    entity_attribute_value_type: EntityAttributeValueTypeEnum
    entity_attribute_value: Any = None

    @field_validator("entity_attribute_enum_id")
    @classmethod
    def validate_entity_attribute_enum_id(cls, v):
        if v is None:
            return v
        return str(validate_uuid(v, "entity_attribute_enum_id"))

    @field_validator("entity_attribute_id")
    @classmethod
    def validate_entity_attribute_id(cls, v):
        return str(validate_uuid(v, "entity_attribute_id"))

    @property
    def value_type(self):
        return self.entity_attribute_value_type

    @property
    def value(self):
        if self.value_type == EntityAttributeValueTypeEnum.enum:
            return self.entity_attribute_enum_id
        else:
            return self.entity_attribute_value


class ModelInputs(BaseModel, extra="forbid"):
    entity_attributes: List[EntityAttribute]
    filters: List[List[EntityAttributeValue]]


class ModelSchemaType(BaseModel, extra="forbid"):
    model_config = ConfigDict(protected_namespaces=(), use_enum_values=True)

    model_inputs: ModelInputs
    model_outputs: List[ModelOutputType]

    @field_validator("model_outputs")
    @classmethod
    def validate_model_output_positions(cls, model_outputs):
        if len(model_outputs) == 0:
            return model_outputs

        model_outputs = sorted(model_outputs, key=lambda o: (o.head, o.position))
        if (model_outputs[0].position != 0) and (model_outputs[0].head != 0):
            raise ValueError("model_outputs ids must start at 0")

        num_heads = len(set([o.head for o in model_outputs]))
        for head_idx in range(num_heads):
            for pos_idx, o in enumerate([_o for _o in model_outputs if _o.head == head_idx]):
                if head_idx != o.head:
                    raise ValueError("model_output.head values must start at 0 and be contigious")
                if pos_idx != o.position:
                    raise ValueError("model_output.position values must start at 0 and be contigious")

        return model_outputs

    def get_input_select_attribute_ids(self) -> List[str]:
        return [str(i.entity_attribute_id) for i in self.model_inputs.entity_attributes]

    def get_input_filter_attribute_ids(self) -> List[str]:
        filter_attr_ids = []
        for input_filter in self.model_inputs.filters:
            filter_attr_ids.append([str(i.entity_attribute_id) for i in input_filter])
        return filter_attr_ids

    def get_input_filter_attribute_values(self) -> List[Any]:
        filter_values = []
        for input_filter in self.model_inputs.filters:
            filter_values.append([i.value for i in input_filter])

        filter_values = [v if len(v) > 0 else None for v in filter_values]
        return filter_values

    def get_head_model_outputs(self, head: int):
        model_outputs = [m for m in self.model_outputs if m.head == head]
        return model_outputs

    def get_head_output_attribute_ids(self, head: int) -> List[str]:
        head_model_outputs = self.get_head_model_outputs(head)
        return [str(m.entity_attribute_id) for m in head_model_outputs]

    def get_head_output_attribute_names(self, head: int) -> List[str]:
        head_model_outputs = self.get_head_model_outputs(head)
        return [str(m.entity_attribute_name) for m in head_model_outputs]

    def get_head_output_attribute_enum_ids(self, head: int) -> List[str]:
        head_model_outputs = self.get_head_model_outputs(head)
        return [str(m.entity_attribute_enum_id) for m in head_model_outputs]

    def get_head_output_attribute_enum_values(self, head: int) -> List[str]:
        head_model_outputs = self.get_head_model_outputs(head)
        return [str(m.entity_attribute_enum_value) for m in head_model_outputs]

    def get_head_output_value_type(self, head: int) -> Optional[EntityAttributeValueTypeEnum]:
        """Will return None if the entity_attribute_value_type is also None
        because it is an Optional parameter of ModelOutputType
        """
        head_model_outputs = self.get_head_model_outputs(head)
        head_output_value_types = {m.entity_attribute_value_type for m in head_model_outputs}
        assert len(head_output_value_types) == 1
        return head_output_value_types.pop()


def _validate_hash_signature_not_none(v, args) -> str:
    assert v is not None, (
        f"dataset({args.data['id']}) hash_signature is None. This is likely becuase it has not been locked"
    )
    return v


class DatasetType(BaseModel, extra="forbid"):
    id: int
    hash_signature: Annotated[str, BeforeValidator(_validate_hash_signature_not_none)]


class ArtefactInfo(GQLBaseModel):
    """camelCase dict keys are intended as these args will be passed to a graphql mutation"""

    type: TrainingRunArtefactTypeEnum
    training_run_id: int
    artefact_path: str
    checkpoint: str
    inference_config: Dict

    def to_dict(self) -> dict:
        returnDict = self.model_dump()
        returnDict["type"] = str(self.type.name)
        return returnDict


def enum_representer(dumper, data):
    return dumper.represent_scalar("tag:yaml.org,2002:str", data.value)


yaml.SafeDumper.add_representer(EntityAttributeValueTypeEnum, enum_representer)

_LEGACY_ENUM_TAG = (
    "tag:yaml.org,2002:python/object/apply:highlighter.client.training_runs.EntityAttributeValueTypeEnum"
)


def enum_constructor(loader, node):
    if isinstance(node, yaml.SequenceNode):
        items = loader.construct_sequence(node)
        return items[0] if items else None
    if isinstance(node, yaml.ScalarNode):
        return loader.construct_scalar(node)
    return loader.construct_object(node)


yaml.SafeLoader.add_constructor(_LEGACY_ENUM_TAG, enum_constructor)


class TrainingRunArtefactType(GQLBaseModel):
    model_config = ConfigDict(use_enum_values=True)

    id: Optional[str] = None
    file_url: Optional[str] = None
    type: TrainingRunArtefactTypeEnum
    updated_at: Optional[str] = Field(default_factory=_get_now_str)
    inference_config: Dict = {}  # Graphql expects at least an empty dict not None
    training_config: Dict = {}  # Graphql expects at least an empty dict not None
    supporting_files: Dict = {}  # Graphql expects at least an empty dict not None

    @field_validator("inference_config", "training_config", "supporting_files", mode="before")
    @classmethod
    def none_to_empty_dict(cls, v):
        if v is None:
            return {}
        return v

    @classmethod
    def from_yaml(cls, path: Union[Path, str]):
        path = Path(path)
        with path.open("r") as f:
            data = yaml.safe_load(f)
        return cls(**data)

    def dump_yaml(self, path: Union[Path, str]):
        path = Path(path)
        with path.open("w") as f:
            yaml.safe_dump(self.model_dump(), f)

    def create(self, client: HLClient, training_run_id: int):
        """Create and upload a training run artefact to Highlighter.

        This method handles the complete workflow of artefact creation including:
        - Uploading the artefact file to S3
        - Creating the artefact record in Highlighter via GraphQL
        - Setting up the associated capability for inference

        Args:
            client: Authenticated HLClient instance for GraphQL operations
            training_run_id: ID of the training run to associate this artefact with

        Returns:
            The created training run artefact with assigned ID

        Raises:
            Exception: If artefact creation fails or returns errors
        """
        from .capabilities import create_capability_for_artefact

        artefact_file_data = upload_file_to_s3(
            client,
            self.file_url,
            mimetype="application/octet-stream",
        )
        logger.info(f"Created file in s3:\n{artefact_file_data}")

        class _TrainingRunArtefactType(GQLBaseModel):
            id: str

        class CreateTrainingRunPayload(GQLBaseModel):
            errors: List[str]
            training_run_artefact: Optional[_TrainingRunArtefactType] = None

        artefact_result = client.createTrainingRunArtefact(
            return_type=CreateTrainingRunPayload,
            trainingRunId=training_run_id,
            type=self.type,
            fileData=artefact_file_data,
            trainingConfig={},
            inferenceConfig=self.inference_config,
        )
        logger.info(f"create artefact result: {artefact_result}")

        has_parameters = "parameters" in self.inference_config or "frame_parameters" in self.inference_config
        required_keys = ("machine_agent_type_id", "code")
        if not has_parameters or any(key not in self.inference_config for key in required_keys):
            logger.info(
                "Skipping capability creation for artefact %s: inference_config missing required keys",
                artefact_result.training_run_artefact.id,
            )
            return artefact_result.training_run_artefact

        create_capability_for_artefact(
            client=client,
            artefact_id=artefact_result.training_run_artefact.id,
            training_run_id=training_run_id,
            inference_config=self.inference_config,
        )
        return artefact_result.training_run_artefact


class TrainingRunType(GQLBaseModel):
    model_config = ConfigDict(use_enum_values=True, protected_namespaces=(), populate_by_name=True)

    id: int
    name: str
    description: Optional[str] = None
    updated_at: str
    trainer_config: Dict = {}
    trainer_config_file: Optional[str] = None
    input_output_schema: Optional[Any] = Field(default=None, alias="inputOutputSchema")
    data: Optional[Dict[str, List[DatasetType]]] = None
    research_plan_id: Optional[int] = Field(default=None, validation_alias=AliasChoices("evaluation_id"))
    experiment_id: Optional[int] = None

    head_idx: int = 0
    oversample: bool = False
    crop_args: Optional[Dict[str, Any]] = None

    @field_validator("trainer_config", mode="before")
    @classmethod
    def trainer_config_must_match_schema(cls, v):
        # API may return null for newly created runs before config is set.
        if v is None:
            return {}
        if isinstance(v, dict):
            return v
        raise ValueError("trainer_config must be a dict")

    @field_validator("input_output_schema", mode="before")
    @classmethod
    def input_output_schema_must_match_model_schema(cls, v):
        if v is None:
            return None
        if isinstance(v, ModelSchemaType):
            return v
        if isinstance(v, dict):
            return ModelSchemaType.model_validate(v)
        raise ValueError("input_output_schema must be a dict")

    @field_validator("crop_args", mode="before")
    @classmethod
    def crop_args_must_match_schema(cls, v):
        if v is None:
            return None
        if isinstance(v, CropArgs):
            return v
        if isinstance(v, dict):
            return CropArgs(**v).model_dump()
        raise ValueError("crop_args must be a dict")

    @classmethod
    def from_yaml(cls, path: Union[Path, str]):
        path = Path(path)
        with path.open("r") as f:
            data = yaml.safe_load(f)
        return cls(**data)

    def dump_yaml(self, path: Union[Path, str]):
        path = Path(path)
        with path.open("w") as f:
            yaml.safe_dump(self.model_dump(), f)

    @property
    def evaluation_id(self) -> Optional[int]:
        return self.research_plan_id


def _download_training_run_artefact(
    training_run_artefact: TrainingRunArtefactType,
    file_url_save_path: Optional[str] = None,
):
    if file_url_save_path is None:
        file_url_save_path = f"{tempfile.mkdtemp()}/{training_run_artefact.id}"

    download_bytes(training_run_artefact.file_url, Path(file_url_save_path))

    training_run_artefact.file_url = str(file_url_save_path)
    return training_run_artefact


def get_training_run(
    hl_client: HLClient,
    training_run_id: int,
) -> TrainingRunType:
    training_run: TrainingRunType = hl_client.trainingRun(
        return_type=TrainingRunType,
        id=training_run_id,
    )
    return training_run


def create_training_run(
    hl_client: HLClient,
    *,
    evaluation_id: int,
    experiment_id: int,
    model_id: int,
    workflow_id: int,
    name: str,
    description: Optional[str] = None,
    convert_description_to_html: Optional[bool] = None,
    source_code_url: Optional[str] = None,
    source_code_commit_hash: Optional[str] = None,
    started_at: Optional[Union[datetime, str]] = None,
    finished_at: Optional[Union[datetime, str]] = None,
    status: Optional[Union[str, Enum]] = None,
) -> TrainingRunType:
    """Create a training run via GraphQL mutation."""

    class CreateTrainingRunResponse(GQLBaseModel):
        training_run: Optional[TrainingRunType] = None
        errors: List[Any]

    kwargs: Dict[str, Any] = {
        "evaluation_id": evaluation_id,
        "experiment_id": experiment_id,
        "model_id": model_id,
        "workflow_id": workflow_id,
        "name": name,
    }

    if description is not None:
        kwargs["description"] = description
    if convert_description_to_html is not None:
        kwargs["convert_description_to_html"] = convert_description_to_html
    if source_code_url is not None:
        kwargs["source_code_url"] = source_code_url
    if source_code_commit_hash is not None:
        kwargs["source_code_commit_hash"] = source_code_commit_hash
    if started_at is not None:
        kwargs["started_at"] = started_at.isoformat() if isinstance(started_at, datetime) else started_at
    if finished_at is not None:
        kwargs["finished_at"] = finished_at.isoformat() if isinstance(finished_at, datetime) else finished_at
    if status is not None:
        kwargs["status"] = status.value if isinstance(status, Enum) else status

    kwargs = {snake_to_camel(k): v for k, v in kwargs.items()}

    response = hl_client.createTrainingRun(
        return_type=CreateTrainingRunResponse,
        **kwargs,
    )
    if response.errors:
        raise ValueError(f"{response.errors}")
    if response.training_run is None:
        raise ValueError("No training_run returned from createTrainingRun")
    return response.training_run


def update_training_run(
    hl_client: HLClient,
    training_run_id: int,
    *,
    evaluation_id: Optional[int] = None,
    experiment_id: Optional[int] = None,
    requested_by_id: Optional[int] = None,
    model_id: Optional[int] = None,
    name: Optional[str] = None,
    description: Optional[str] = None,
    convert_description_to_html: Optional[bool] = None,
    source_code_url: Optional[str] = None,
    source_code_commit_hash: Optional[str] = None,
    started_at: Optional[Union[datetime, str]] = None,
    finished_at: Optional[Union[datetime, str]] = None,
    status: Optional[Union[str, Enum]] = None,
    trainer_config: Optional[Union[Dict[str, Any], str]] = None,
) -> TrainingRunType:
    """Update a training run via GraphQL mutation."""

    class UpdateTrainingRunResponse(GQLBaseModel):
        training_run: Optional[TrainingRunType] = None
        errors: List[Any]

    kwargs: Dict[str, Any] = {}
    if evaluation_id is not None:
        kwargs["evaluation_id"] = evaluation_id
    if experiment_id is not None:
        kwargs["experiment_id"] = experiment_id
    if requested_by_id is not None:
        kwargs["requested_by_id"] = requested_by_id
    if model_id is not None:
        kwargs["model_id"] = model_id
    if name is not None:
        kwargs["name"] = name
    if description is not None:
        kwargs["description"] = description
    if convert_description_to_html is not None:
        kwargs["convert_description_to_html"] = convert_description_to_html
    if source_code_url is not None:
        kwargs["source_code_url"] = source_code_url
    if source_code_commit_hash is not None:
        kwargs["source_code_commit_hash"] = source_code_commit_hash
    if started_at is not None:
        kwargs["started_at"] = started_at.isoformat() if isinstance(started_at, datetime) else started_at
    if finished_at is not None:
        kwargs["finished_at"] = finished_at.isoformat() if isinstance(finished_at, datetime) else finished_at
    if status is not None:
        kwargs["status"] = status.value if isinstance(status, Enum) else status
    if trainer_config is not None:
        if isinstance(trainer_config, dict):
            kwargs["trainer_config"] = json.dumps(trainer_config)
        else:
            kwargs["trainer_config"] = trainer_config

    kwargs = {snake_to_camel(k): v for k, v in kwargs.items()}

    response = hl_client.updateTrainingRun(
        return_type=UpdateTrainingRunResponse,
        id=training_run_id,
        **kwargs,
    )
    if response.errors:
        raise ValueError(f"{response.errors}")
    if response.training_run is None:
        raise ValueError("No training_run returned from updateTrainingRun")
    return response.training_run


def get_training_run_artefacts(
    hl_client: HLClient,
    training_run_id: int,
    filter_by_artefact_type: Optional[Union[str, TrainingRunArtefactTypeEnum]] = None,
) -> List[TrainingRunArtefactType]:
    """Return a list of artefacts for a training run
    sorted newest to oldest according to TrainingRunArtefactType.updated_at

    If `filter_by_artefact_type` is set then will only retrun artefacts with
    matching TrainingRunArtefactType.type
    """

    class TrainingRunTypeWithArtefacts(TrainingRunType):
        model_implementation_file_url: Optional[str] = None
        training_run_artefacts: List[TrainingRunArtefactType] = []

    training_run: TrainingRunTypeWithArtefacts = hl_client.trainingRun(
        return_type=TrainingRunTypeWithArtefacts,
        id=training_run_id,
    )

    if training_run.model_implementation_file_url is not None:
        # If TrainingRun.updated_at is None then we set it to
        # datetime.min for the purposes of sorting.
        capability_implementation_file_artefact: TrainingRunArtefactType
        capability_implementation_file_artefact = TrainingRunArtefactType(
            id=None,
            file_url=training_run.model_implementation_file_url,
            type=TrainingRunArtefactTypeEnum.DeprecatedCapabilityImplementationFile,
            updated_at=datetime.min.isoformat() if training_run.updated_at else training_run.updated_at,
            inference_config={},
            supporting_files={},
        )
        training_run.training_run_artefacts.append(capability_implementation_file_artefact)

    if filter_by_artefact_type is not None:
        artefact_type = TrainingRunArtefactTypeEnum(filter_by_artefact_type)
        training_run.training_run_artefacts = [
            a for a in training_run.training_run_artefacts if a.type == artefact_type
        ]

    return sorted(training_run.training_run_artefacts, key=lambda x: x.updated_at, reverse=True)


def get_training_run_artefact(
    hl_client: HLClient,
    training_run_artefact_id: str,
    download_file_url: bool = False,
    file_url_save_path: Optional[str] = None,
) -> TrainingRunArtefactType:
    """Get the TrainingRunArtefact object for a given training_run_artefact_id.

    If `download_file_url` is `True` then download either to a temporary
    directory or `file_url_save_path`. Once downloaded TrainingRunArtefact.file_url
    will be updated to the location where is has been downloaded to.
    """

    training_run_artefact: TrainingRunArtefactType
    training_run_artefact = hl_client.trainingRunArtefact(
        return_type=TrainingRunArtefactType, id=training_run_artefact_id
    )

    if download_file_url:
        training_run_artefact = _download_training_run_artefact(
            training_run_artefact, file_url_save_path=file_url_save_path
        )

    return training_run_artefact


def get_latest_training_run_artefact(
    hl_client: HLClient,
    training_run_id: int,
    download_file_url: bool = False,
    file_url_save_path: Optional[str] = None,
    filter_by_artefact_type: Optional[Union[str, TrainingRunArtefactTypeEnum]] = None,
) -> Optional[TrainingRunArtefactType]:
    """Get the TrainingRunArtefact with the most recent updated_at value.

    If `download_file_url` is `True` then download either to a temporary
    directory or `file_url_save_path`. Once downloaded TrainingRunArtefact.file_url
    will be updated to the location where is has been downloaded to.
    """
    training_run_artefacts: List[TrainingRunArtefactType] = get_training_run_artefacts(
        hl_client, training_run_id, filter_by_artefact_type=filter_by_artefact_type
    )

    if len(training_run_artefacts) == 0:
        return None
    training_run_artefact: TrainingRunArtefactType = training_run_artefacts[0]

    if download_file_url:
        abs_file_url_save_path = str(Path(file_url_save_path).absolute())
        training_run_artefact = _download_training_run_artefact(
            training_run_artefact, file_url_save_path=abs_file_url_save_path
        )

    return training_run_artefact


def get_capability_implementation_file_url(
    hl_client: HLClient,
    training_run_id: int,
) -> Optional[str]:
    class _TrainingRunWithFileUrl(TrainingRunType):
        model_implementation_file_url: Optional[str] = None

    result = hl_client.trainingRun(
        return_type=_TrainingRunWithFileUrl,
        id=training_run_id,
    )
    return result.model_implementation_file_url


def create_training_run_artefact(
    hl_client: HLClient,
    training_run_id: int,
    artefact_path: str,
    training_run_artefact_type: TrainingRunArtefactTypeEnum,
    checkpoint_name: str,
    inference_config: Union[Path, str, Dict],
    data_source_uuid: Optional[str] = None,
) -> TrainingRunArtefactType:
    if not isinstance(inference_config, dict):
        inference_config_path = Path(inference_config)
        with inference_config_path.open("r") as f:
            loaders = {
                ".yaml": yaml.safe_load,
                ".yml": yaml.safe_load,
                ".json": json.load,
            }

            inference_config = loaders[inference_config_path.suffix](f)

    artefact_file_data = upload_file_to_s3(
        hl_client,
        artefact_path,
        data_source_uuid=data_source_uuid,
    )
    print(f"Created file in s3:\n{artefact_file_data}")

    class CreateTrainingRunPayload(GQLBaseModel):
        errors: List[str]
        training_run_artefact: Optional[TrainingRunArtefactType] = None

    result = hl_client.createTrainingRunArtefact(
        return_type=CreateTrainingRunPayload,
        trainingRunId=training_run_id,
        type=training_run_artefact_type,
        checkpoint=checkpoint_name,
        fileData=artefact_file_data,
        inferenceConfig=inference_config,
    )
    if result.errors:
        raise ValueError(f"{result.errors}")

    return result.training_run_artefact
