"""Utilities for converting between request states."""

import logging
from copy import deepcopy
from typing import (
    Any,
    Callable,
    cast,
    Dict,
    List,
    Optional,
)

from galaxy.tool_util_models.parameters import (
    BooleanParameterModel,
    ConditionalParameterModel,
    ConditionalWhen,
    create_job_runtime_model,
    DataCollectionParameterModel,
    DataCollectionRequest,
    DataCollectionRequestInternal,
    DataColumnParameterModel,
    DataInternalJson,
    DataJobInternalT,
    DataParameterModel,
    DataRequestCollectionUri,
    DataRequestHda,
    DataRequestInternalHda,
    DataRequestInternalHdca,
    DataRequestUri,
    DatasetCollectionElementReference,
    DiscriminatorType,
    DrillDownParameterModel,
    FloatParameterModel,
    GenomeBuildParameterModel,
    HiddenParameterModel,
    IntegerParameterModel,
    RepeatParameterModel,
    SectionParameterModel,
    SelectParameterModel,
    TextParameterModel,
    ToolParameterBundle,
    ToolParameterT,
)
from galaxy.tool_util_models.tool_source import (
    JsonTestCollectionDefDict,
    JsonTestDatasetDefDict,
)
from .state import (
    JobInternalToolState,
    JobRuntimeToolState,
    LandingRequestInternalToolState,
    LandingRequestToolState,
    RelaxedRequestToolState,
    RequestInternalDereferencedToolState,
    RequestInternalToolState,
    RequestToolState,
    TestCaseToolState,
)
from .visitor import (
    Callback,
    validate_explicit_conditional_test_value,
    visit_input_values,
    VISITOR_NO_REPLACEMENT,
)

log = logging.getLogger(__name__)


DecodeFunctionT = Callable[[str], int]
EncodeFunctionT = Callable[[int], str]
DereferenceCallable = Callable[[DataRequestUri], DataRequestInternalHda]
DereferenceCollectionCallable = Callable[[DataRequestCollectionUri], DataRequestInternalHdca]
# interfaces for adapting test data dictionaries to tool request dictionaries
# e.g. {class: File, path: foo.bed} => {src: hda, id: ab1235cdfea3}
AdaptDatasets = Callable[[JsonTestDatasetDefDict], DataRequestHda]
AdaptCollections = Callable[[JsonTestCollectionDefDict], DataCollectionRequest]

OPENAPI_REF_TEMPLATE = "#/components/schemas/{model}"


def cwl_runtime_model(input_models: ToolParameterBundle):
    model = create_job_runtime_model(input_models)
    openapi_schema: Dict[str, Any] = {
        "openapi": "3.1.0",
        "info": {
            "title": "Custom API",
            "version": "1.0.0",
            "description": "Generated by Pydantic with OpenAPI customization.",
        },
        "paths": {},  # Define paths manually or generate them elsewhere
        "components": {},
    }
    schemas = model.model_json_schema(mode="serialization", ref_template=OPENAPI_REF_TEMPLATE)
    openapi_schema["components"]["schemas"] = schemas.pop("$defs", {})
    openapi_schema["components"]["schemas"]["inputs"] = schemas
    return openapi_schema


def decode(
    external_state: RequestToolState,
    input_models: ToolParameterBundle,
    decode_id: Callable[[str], int],
    name_base: Optional[str] = None,
) -> RequestInternalToolState:
    """Prepare an internal representation of tool state (request_internal) for storing in the database."""

    external_state.validate(input_models, f"{name_base} (request model)")
    decode_callback = _decode_callback_for(decode_id)
    internal_state_dict = visit_input_values(
        input_models,
        external_state,
        decode_callback,
    )

    internal_request_state = RequestInternalToolState(internal_state_dict)
    internal_request_state.validate(input_models, f"{name_base} (request internal model)")
    return internal_request_state


def encode(
    internal_state: RequestInternalToolState, input_models: ToolParameterBundle, encode_id: EncodeFunctionT
) -> RequestToolState:
    """Prepare an external representation of tool state (request) from persisted state in the database (request_internal)."""

    encode_callback = _encode_callback_for(encode_id)
    request_state_dict = visit_input_values(
        input_models,
        internal_state,
        encode_callback,
    )
    request_state = RequestToolState(request_state_dict)
    request_state.validate(input_models)
    return request_state


def landing_decode(
    external_state: LandingRequestToolState, input_models: ToolParameterBundle, decode_id: Callable[[str], int]
) -> LandingRequestInternalToolState:
    """Prepare an external representation of tool state (request) for storing in the database (request_internal)."""

    external_state.validate(input_models)
    decode_callback = _decode_callback_for(decode_id)
    internal_state_dict = visit_input_values(
        input_models,
        external_state,
        decode_callback,
    )

    internal_request_state = LandingRequestInternalToolState(internal_state_dict)
    internal_request_state.validate(input_models)
    return internal_request_state


def landing_encode(
    internal_state: LandingRequestInternalToolState, input_models: ToolParameterBundle, encode_id: EncodeFunctionT
) -> LandingRequestToolState:
    """Prepare an external representation of tool state (request) for storing in the database (request_internal)."""

    encode_callback = _encode_callback_for(encode_id)
    request_state_dict = visit_input_values(
        input_models,
        internal_state,
        encode_callback,
    )
    request_state = LandingRequestToolState(request_state_dict)
    request_state.validate(input_models)
    return request_state


def strictify(relaxed_state: RelaxedRequestToolState, input_models: ToolParameterBundle) -> RequestToolState:
    """Convert a relaxed request state into a strict request state by applying legacy behavior."""

    tool_state = deepcopy(relaxed_state.input_state)

    def _strictify_parameter(tool_state: Dict[str, Any], parameter: ToolParameterT) -> None:
        if isinstance(parameter, ConditionalParameterModel):
            conditional_state = _initialize_conditional_state(parameter, tool_state)

            test_parameter = parameter.test_parameter
            test_parameter_name = test_parameter.name

            explicit_test_value: Optional[DiscriminatorType] = (
                conditional_state[test_parameter_name] if test_parameter_name in conditional_state else None
            )
            test_value = validate_explicit_conditional_test_value(test_parameter_name, explicit_test_value)
            when = _select_which_when(parameter, test_value, conditional_state)
            _strictify_parameter(conditional_state, test_parameter)
            _strictify_parameters(conditional_state, when)
        elif isinstance(parameter, RepeatParameterModel):
            repeat_instances = _initialize_repeat_state(parameter, tool_state)
            for instance_state in repeat_instances:
                _strictify_parameters(instance_state, parameter)
        elif isinstance(parameter, SectionParameterModel):
            section_state = _initialize_section_state(parameter, tool_state)
            _fill_defaults(section_state, parameter)
        elif isinstance(parameter, TextParameterModel):
            parameter_name = parameter.name
            if parameter_name not in tool_state:
                if not parameter.optional:
                    # restore legacy behavior of allowing empty string implicit default
                    # for these non-optional inputs.
                    tool_state[parameter_name] = ""
                else:
                    tool_state[parameter_name] = None
            else:
                # legacy behavior of converting explicit None into implicit null. We should introduce
                # a layer somewhere to deal with this behavior further up the stack and clean up these models.
                if not parameter.optional and tool_state[parameter_name] is None:
                    tool_state[parameter_name] = ""

    def _strictify_parameters(tool_state: Dict[str, Any], input_models: ToolParameterBundle) -> None:
        for parameter in input_models.parameters:
            _strictify_parameter(tool_state, parameter)

    _strictify_parameters(tool_state, input_models)

    request_state = RequestToolState(tool_state)
    request_state.validate(input_models)
    return request_state


def dereference(
    internal_state: RequestInternalToolState,
    input_models: ToolParameterBundle,
    dereference: DereferenceCallable,
    dereference_collection: DereferenceCollectionCallable,
) -> RequestInternalDereferencedToolState:

    def dereference_dict(src_dict: dict):
        src = src_dict.get("src")
        clazz = src_dict.get("class")
        if src == "url":
            data_request_uri: DataRequestUri = DataRequestUri.model_validate(src_dict)
            data_request_hda: DataRequestInternalHda = dereference(data_request_uri)
            return data_request_hda.model_dump()
        elif clazz == "Collection":
            data_request_collection_from_uri: DataRequestCollectionUri = DataRequestCollectionUri.model_validate(
                src_dict
            )
            data_request_hdca: DataRequestInternalHdca = dereference_collection(data_request_collection_from_uri)
            return data_request_hdca.model_dump()
        else:
            return src_dict

    def dereference_callback(parameter: ToolParameterT, value: Any):
        if isinstance(parameter, DataParameterModel):
            if value is None:
                return VISITOR_NO_REPLACEMENT
            if parameter.multiple and isinstance(value, list):
                return list(map(dereference_dict, value))
            else:
                assert isinstance(value, dict), str(value)
                return dereference_dict(value)
        elif isinstance(parameter, DataCollectionParameterModel):
            if value is None:
                return VISITOR_NO_REPLACEMENT
            assert isinstance(value, dict), str(value)
            return dereference_dict(value)
        else:
            return VISITOR_NO_REPLACEMENT

    request_state_dict = visit_input_values(
        input_models,
        internal_state,
        dereference_callback,
    )
    request_state = RequestInternalDereferencedToolState(request_state_dict)
    request_state.validate(input_models)
    return request_state


def encode_test(
    test_case_state: TestCaseToolState,
    input_models: ToolParameterBundle,
    adapt_datasets: AdaptDatasets,
    adapt_collections: AdaptCollections,
):

    def encode_callback(parameter: ToolParameterT, value: Any):
        if isinstance(parameter, DataParameterModel):
            if value is not None:
                if parameter.multiple:
                    assert isinstance(value, list), str(value)
                    test_datasets = cast(List[JsonTestDatasetDefDict], value)
                    return [d.model_dump() for d in map(adapt_datasets, test_datasets)]
                else:
                    assert isinstance(value, dict), str(value)
                    test_dataset = cast(JsonTestDatasetDefDict, value)
                    return adapt_datasets(test_dataset).model_dump()
        elif isinstance(parameter, DataCollectionParameterModel):
            if value is not None:
                assert isinstance(value, dict), str(value)
                test_collection = cast(JsonTestCollectionDefDict, value)
                return adapt_collections(test_collection).model_dump()
        elif isinstance(parameter, SelectParameterModel):
            if parameter.multiple and value is not None and isinstance(value, (str,)):
                return [v.strip() for v in value.split(",")]
            else:
                return VISITOR_NO_REPLACEMENT
        elif isinstance(parameter, DrillDownParameterModel):
            if parameter.multiple and value is not None:
                return [v.strip() for v in value.split(",")]
            else:
                return VISITOR_NO_REPLACEMENT
        elif isinstance(parameter, DataColumnParameterModel):
            if parameter.multiple and value is not None and isinstance(value, (str,)):
                return [int(v.strip()) for v in value.split(",")]
            else:
                return VISITOR_NO_REPLACEMENT

        return VISITOR_NO_REPLACEMENT

    request_state_dict = visit_input_values(
        input_models,
        test_case_state,
        encode_callback,
    )
    request_state = RequestToolState(request_state_dict)
    request_state.validate(input_models)
    return request_state


def fill_static_defaults(
    tool_state: Dict[str, Any], input_models: ToolParameterBundle, profile: float, partial: bool = True
) -> Dict[str, Any]:
    """If additional defaults might stem from Galaxy runtime, partial should be true.

    Setting partial to True, prevents runtime validation.
    """
    _fill_defaults(tool_state, input_models)

    if not partial:
        internal_state = JobInternalToolState(tool_state)
        internal_state.validate(input_models)
    return tool_state


def _fill_defaults(tool_state: Dict[str, Any], input_models: ToolParameterBundle) -> None:
    for parameter in input_models.parameters:
        _fill_default_for(tool_state, parameter)


def _fill_default_for(tool_state: Dict[str, Any], parameter: ToolParameterT) -> None:
    parameter_name = parameter.name
    if isinstance(parameter, BooleanParameterModel):
        if parameter_name not in tool_state:
            # even optional parameters default to false if not in the body of the request :_(
            # see test_tools.py -> expression_null_handling_boolean or test cases for gx_boolean_optional.xml
            tool_state[parameter_name] = parameter.value or False

    if isinstance(parameter, (IntegerParameterModel, FloatParameterModel, HiddenParameterModel)):
        if parameter_name not in tool_state:
            tool_state[parameter_name] = parameter.value
    elif isinstance(parameter, GenomeBuildParameterModel):
        if parameter_name not in tool_state and parameter.optional:
            tool_state[parameter_name] = None
    elif isinstance(parameter, SelectParameterModel):
        # don't fill in dynamic parameters - wait for runtime to specify the default
        if parameter.dynamic_options:
            return

        if parameter_name not in tool_state:
            if not parameter.multiple:
                tool_state[parameter_name] = parameter.default_value
            else:
                default_values = parameter.default_values
                tool_state[parameter_name] = default_values if default_values else None
    elif isinstance(parameter, DrillDownParameterModel):
        if parameter_name not in tool_state:
            if parameter.multiple:
                options = parameter.default_options
                if options is not None:
                    tool_state[parameter_name] = options
            else:
                option = parameter.default_option
                if option is not None:
                    tool_state[parameter_name] = option
    elif isinstance(parameter, ConditionalParameterModel):
        conditional_state = _initialize_conditional_state(parameter, tool_state)

        test_parameter = parameter.test_parameter
        test_parameter_name = test_parameter.name

        explicit_test_value: Optional[DiscriminatorType] = (
            conditional_state[test_parameter_name] if test_parameter_name in conditional_state else None
        )
        test_value = validate_explicit_conditional_test_value(test_parameter_name, explicit_test_value)
        when = _select_which_when(parameter, test_value, conditional_state)
        _fill_default_for(conditional_state, test_parameter)
        _fill_defaults(conditional_state, when)
    elif isinstance(parameter, RepeatParameterModel):
        repeat_instances = _initialize_repeat_state(parameter, tool_state)
        for instance_state in repeat_instances:
            _fill_defaults(instance_state, parameter)
    elif isinstance(parameter, SectionParameterModel):
        section_state = _initialize_section_state(parameter, tool_state)
        _fill_defaults(section_state, parameter)
    elif isinstance(parameter, DataCollectionParameterModel):
        collection_parameter = parameter
        if parameter_name not in tool_state and collection_parameter.optional:
            tool_state[parameter_name] = None
    elif isinstance(parameter, TextParameterModel):
        if parameter_name not in tool_state:
            if not parameter.optional:
                # restore legacy behavior of allowing empty string implicit default
                # for these non-optional inputs.
                tool_state[parameter_name] = parameter.default_value or ""
            else:
                tool_state[parameter_name] = parameter.default_value or None
        else:
            # legacy behavior of converting explicit None into implicit null. We should introduce
            # a layer somewhere to deal with this behavior further up the stack and clean up these models.
            if not parameter.optional and tool_state[parameter_name] is None:
                tool_state[parameter_name] = parameter.default_value or ""


def _initialize_section_state(parameter: SectionParameterModel, tool_state: Dict[str, Any]) -> Dict[str, Any]:
    parameter_name = parameter.name
    if parameter_name not in tool_state:
        tool_state[parameter_name] = {}
    section_state = cast(Dict[str, Any], tool_state[parameter_name])
    return section_state


def _initialize_conditional_state(parameter: ConditionalParameterModel, tool_state: Dict[str, Any]) -> Dict[str, Any]:
    parameter_name = parameter.name
    if parameter_name not in tool_state:
        tool_state[parameter_name] = {}

    raw_conditional_state = tool_state[parameter_name]
    assert isinstance(raw_conditional_state, dict)
    conditional_state = cast(Dict[str, Any], raw_conditional_state)
    return conditional_state


def _initialize_repeat_state(parameter: RepeatParameterModel, tool_state: Dict[str, Any]) -> List[Dict[str, Any]]:
    parameter_name = parameter.name
    if parameter_name not in tool_state:
        tool_state[parameter_name] = []
    repeat_instances = cast(List[Dict[str, Any]], tool_state[parameter_name])
    if parameter.min:
        while len(repeat_instances) < parameter.min:
            repeat_instances.append({})
    return repeat_instances


def _select_which_when(
    conditional: ConditionalParameterModel, test_value: Optional[DiscriminatorType], conditional_state: Dict[str, Any]
) -> ConditionalWhen:
    for when in conditional.whens:
        if test_value is None and when.is_default_when:
            return when
        elif test_value == when.discriminator:
            return when
    else:
        raise Exception(
            f"Invalid conditional test value ({test_value}) for parameter ({conditional.test_parameter.name})"
        )


def _encode_callback_for(encode_id: EncodeFunctionT) -> Callback:

    def encode_src_dict(src_dict: dict):
        if "id" in src_dict:
            encoded_dict = src_dict.copy()
            encoded_dict["id"] = encode_id(src_dict["id"])
            return encoded_dict
        else:
            return src_dict

    def encode_element(element: dict):
        if element.get("__class__") == "Batch":
            encoded = element.copy()
            values = encoded.pop("values")
            encoded["values"] = list(map(encode_src_dict, values))
            return encoded
        else:
            return encode_src_dict(element)

    def encode_callback(parameter: ToolParameterT, value: Any):
        if isinstance(parameter, DataParameterModel):
            if parameter.multiple and isinstance(value, list):
                return list(map(encode_element, value))
            else:
                assert isinstance(value, dict), str(value)
                return encode_element(value)
        elif isinstance(parameter, DataCollectionParameterModel):
            assert isinstance(value, dict), str(value)
            return encode_element(value)
        else:
            return VISITOR_NO_REPLACEMENT

    return encode_callback


def _decode_callback_for(decode_id: DecodeFunctionT) -> Callback:

    def decode_src_dict(src_dict: dict):
        if "id" in src_dict:
            decoded_dict = src_dict.copy()
            decoded_dict["id"] = decode_id(src_dict["id"])
            return decoded_dict
        else:
            return src_dict

    def decode_element(element: dict):
        if element.get("__class__") == "Batch":
            decoded = element.copy()
            values = decoded.pop("values")
            decoded["values"] = list(map(decode_src_dict, values))
            return decoded
        else:
            return decode_src_dict(element)

    def decode_callback(parameter: ToolParameterT, value: Any):
        if isinstance(parameter, DataParameterModel):
            if value is None:
                return VISITOR_NO_REPLACEMENT
            if parameter.multiple and isinstance(value, list):
                return list(map(decode_element, value))
            else:
                assert isinstance(value, dict), str(value)
                return decode_element(value)
        elif isinstance(parameter, DataCollectionParameterModel):
            if value is None:
                return VISITOR_NO_REPLACEMENT
            assert isinstance(value, dict), str(value)
            return decode_element(value)
        else:
            return VISITOR_NO_REPLACEMENT

    return decode_callback


DatasetToRuntimeJson = Callable[[DataJobInternalT], DataInternalJson]
CollectionToRuntimeJson = Callable[[DataCollectionRequestInternal, Optional[str]], Any]


def runtimeify(
    internal_state: JobInternalToolState,
    input_models: ToolParameterBundle,
    adapt_dataset: DatasetToRuntimeJson,
    adapt_collection: CollectionToRuntimeJson,
) -> JobRuntimeToolState:

    def adapt_dict(value: dict):
        assert isinstance(value, dict), str(value)
        src = value.get("src")
        if src == "dce":
            dce_ref = DatasetCollectionElementReference(**value)
            as_json = adapt_dataset(dce_ref).model_dump(by_alias=True)
        else:
            data_request_internal_hda = DataRequestInternalHda(**value)
            as_json = adapt_dataset(data_request_internal_hda).model_dump(by_alias=True)
        return as_json

    def to_runtime_callback(parameter: ToolParameterT, value: Any):
        if isinstance(parameter, DataParameterModel):
            if parameter.multiple and isinstance(value, list):
                return list(map(adapt_dict, value))
            else:
                return adapt_dict(value)
        elif isinstance(parameter, DataCollectionParameterModel):
            if value is None:
                return VISITOR_NO_REPLACEMENT
            assert isinstance(value, dict), str(value)
            collection_request = DataCollectionRequestInternal(**value)
            result = adapt_collection(collection_request, parameter.collection_type)
            return result.model_dump(by_alias=True)
        else:
            return VISITOR_NO_REPLACEMENT

    runtime_state_dict = visit_input_values(
        input_models,
        internal_state,
        to_runtime_callback,
    )
    runtime_state = JobRuntimeToolState(runtime_state_dict)
    runtime_state.validate(input_models)
    return runtime_state
