import os
from pathlib import Path
from typing import Dict, List, Literal, Optional, Union

from pydantic import BaseModel
from sgqlc.operation import Operation

from ML_management.graphql import schema
from ML_management.graphql.schema import CodeJob, EnvParamInput
from ML_management.graphql.send_graphql_request import send_graphql_request
from ML_management.mlmanagement.git_info import get_git_info
from ML_management.mlmanagement.log_api import _open_pipe_send_request, _raise_error
from ML_management.mlmanagement.utils import calculate_hash_directory, check_exclude_include, hash_file
from ML_management.mlmanagement.visibility_options import VisibilityOptions
from ML_management.sdk.image import get_image_by_name
from ML_management.sdk.parameters import ResourcesForm
from ML_management.variables import get_log_service_url


class CodeMetaInfo(BaseModel):
    code_id: int
    experiment_name: str


def _get_code_hash(hash_code: str, visibility: VisibilityOptions) -> Optional[int]:
    op = Operation(schema.Query)
    op.get_code_with_hash(hash_code=hash_code, visibility=visibility.name)

    result = send_graphql_request(op=op, json_response=False)
    return result.get_code_with_hash


def add_custom_code_job(
    local_path: str,
    bash_commands: List[str],
    image_name: str,
    job_name: Optional[str] = None,
    exclude: list[str | Path] | None = None,
    include: list[str | Path] | None = None,
    env_variables: Optional[Dict[str, str]] = None,
    resources: Optional[ResourcesForm] = None,
    is_distributed: bool = False,
    experiment_name: str = "Default",
    experiment_visibility: Union[
        Literal[VisibilityOptions.PRIVATE, VisibilityOptions.PUBLIC], VisibilityOptions
    ] = VisibilityOptions.PRIVATE,
    visibility: Union[Literal["private", "public"], VisibilityOptions] = VisibilityOptions.PRIVATE,
    additional_system_packages: Optional[List[str]] = None,
    verbose: bool = True,
) -> CodeJob:
    """
    Create execution job from arbitrary code.

    The job created by this function can be executed on an arbitrary number of nodes greater than or equal to 1.
    The computing cluster distributes all resources equally among the allocated nodes.
    If such an equal distribution is not possible, the job will be rejected.

    Parameters
    ----------
    local_path: str
        Path to folder with code.
    bash_commands: List[str]
        Commands that will be run in the execution container.
    image_name: str
        The name of the base image on which the job will be executed.
    job_name: Optional[str] = None
        Name of the created job.
    exclude: list[str | pathlib.Path] | None = None
        List the relative paths of the files within the logging folder that should be excluded from logging.
        For example: ["unused.py", "artifacts/modules/file.py"].
        Regular expressions are not currently supported.
        Parameters exclude and include are mutually exclusive and only one of them can be set.
    include: list[str | pathlib.Path] | None = None
        List the relative paths of only those files within the logging folder that should be included in logging.
        For example: ["model.py", "artifacts/modules/code.py"].
        Regular expressions are not currently supported.
        Parameters exclude and include are mutually exclusive and only one of them can be set.
    env_variables: Optional[Dict[str, str]] = None
        Environment variables that will be set before starting the job.
    resources: ResourcesForm
        Resources required for job execution.
        They will be allocated on one node, if it is not possible, job will be rejected.
    is_distributed: bool = False
        Distributed mode.
    experiment_name: ExperimentParams
        Name of the experiment. Default: "Default"
    experiment_visibility: Union[Literal['private', 'public'], VisibilityOptions]
        Visibility of experiment if this is new. Default: PRIVATE.
    visibility: Union[Literal['private', 'public'], VisibilityOptions]
        Visibility of this job to other users. Default: PRIVATE.
    additional_system_packages: Optional[List[str]] = None
        List of system libraries for Debian family distributions that need to be installed in the job. Default: None
    verbose: bool = True
        Whether to disable the entire progressbar wrapper.

    Returns
    -------
    CodeJob
        Instance of the Job class.
    """
    local_path = os.path.abspath(local_path)
    image_id = get_image_by_name(image_name).id
    visibility = VisibilityOptions(visibility)
    experiment_visibility = VisibilityOptions(experiment_visibility)
    if not os.path.exists(local_path):
        raise FileNotFoundError(f"Path: {local_path} does not exist.")

    exclude, include = check_exclude_include(exclude, include)

    log_request = {
        "visibility": visibility,
        "experiment_name": experiment_name,
        "experiment_visibility": experiment_visibility,
    }

    hash_code = calculate_hash_directory(local_path) if os.path.isdir(local_path) else hash_file(local_path)
    code_id = _get_code_hash(hash_code, visibility)
    experiment_params = schema.ExperimentInput(experiment_name=experiment_name, visibility=experiment_visibility.name)

    if code_id is None:
        git_info = get_git_info(local_path)
        log_request["hash_code"] = hash_code
        log_request["git_info"] = git_info.model_dump() if git_info else None
        response = _open_pipe_send_request(
            local_path,
            log_request,
            url=get_log_service_url("log_job_code"),
            verbose=verbose,
            exclude=exclude,
            include=include,
        )
        _raise_error(response)
        result = response.json()
        info = CodeMetaInfo.model_validate(result)
        code_id = info.code_id

    if resources is None:
        resources = ResourcesForm()
    resources = schema.ResourcesInput(
        cpus=resources.cpus,
        memory_per_node=resources.memory_per_node,
        gpu_number=resources.gpu_number,
        gpu_type=resources.gpu_type,
    )
    op = Operation(schema.Mutation)
    mutation = op.add_custom_code_job(
        form=schema.JobCodeParameters(
            code_id=code_id,
            bash_commands=bash_commands,
            experiment_params=experiment_params,
            visibility=VisibilityOptions(visibility).name,
            additional_system_packages=additional_system_packages,
            job_name=job_name,
            image_id=image_id,
            resources=resources,
            is_distributed=is_distributed,
            env_variables=[EnvParamInput(key=key, value=value) for key, value in env_variables.items()]
            if env_variables
            else None,
        )
    )

    mutation.name()
    mutation.id()

    job = send_graphql_request(op, json_response=False)

    return job.add_custom_code_job
