import asyncio
import os
from json import JSONDecodeError, JSONDecoder, JSONEncoder, loads
from typing import Literal

import redis

from ..helpers import (
    apply_module_callbacks,
    get_legacy_responses,
    get_protocol_version,
    nativestr,
)
from .commands import FPHAType, JSONCommands
from .decoders import bulk_of_jsons, decode_list


class _JSONBase(JSONCommands):
    """
    Create a client for talking to json.

    :param decoder:
    :type json.JSONDecoder: An instance of json.JSONDecoder

    :param encoder:
    :type json.JSONEncoder: An instance of json.JSONEncoder
    """

    def __init__(
        self, client, version=None, decoder=JSONDecoder(), encoder=JSONEncoder()
    ):
        """
        Create a client for talking to json.

        :param decoder:
        :type json.JSONDecoder: An instance of json.JSONDecoder

        :param encoder:
        :type json.JSONEncoder: An instance of json.JSONEncoder
        """
        # Set the module commands' callbacks
        _MODULE_CALLBACKS = {
            "JSON.ARRPOP": self._decode,
            "JSON.DEBUG": self._decode,
            "JSON.GET": self._decode,
            "JSON.MERGE": lambda r: r and nativestr(r) == "OK",
            "JSON.MGET": bulk_of_jsons(self._decode),
            "JSON.MSET": lambda r: r and nativestr(r) == "OK",
            "JSON.RESP": self._decode,
            "JSON.SET": lambda r: r and nativestr(r) == "OK",
            "JSON.TOGGLE": self._decode,
        }

        _RESP2_MODULE_CALLBACKS = {
            "JSON.ARRAPPEND": self._decode,
            "JSON.ARRINDEX": self._decode,
            "JSON.ARRINSERT": self._decode,
            "JSON.ARRLEN": self._decode,
            "JSON.ARRTRIM": self._decode,
            "JSON.CLEAR": int,
            "JSON.DEL": int,
            "JSON.FORGET": int,
            "JSON.GET": self._decode,
            "JSON.NUMINCRBY": lambda r, **kwargs: self._decode(r),
            "JSON.NUMMULTBY": lambda r, **kwargs: self._decode(r),
            "JSON.OBJKEYS": self._decode,
            "JSON.STRAPPEND": self._decode,
            "JSON.OBJLEN": self._decode,
            "JSON.STRLEN": self._decode,
            "JSON.TOGGLE": self._decode,
        }

        _RESP3_MODULE_CALLBACKS = {}
        # RESP2 wire normalised to the unified shape from the reverted
        # unification PR: NUMINCRBY/NUMMULTBY are always arrays,
        # JSON.RESP uses native floats, and missing JSON.TYPE keys stay
        # ``None`` instead of becoming ``[None]``.
        _RESP2_UNIFIED_MODULE_CALLBACKS = {
            "JSON.CLEAR": int,
            "JSON.DEL": int,
            "JSON.FORGET": int,
            "JSON.NUMINCRBY": self._decode_json_numop,
            "JSON.NUMMULTBY": self._decode_json_numop,
            "JSON.RESP": self._decode_resp_command_unified,
            "JSON.TYPE": lambda r: [r] if r is not None else r,
        }
        _RESP3_UNIFIED_MODULE_CALLBACKS = {
            "JSON.TYPE": lambda r: None if r == [None] else r,
        }
        # RESP3 wire normalised back to today's RESP2 Python shapes:
        # keep ``nativestr`` for OBJKEYS, unwrap JSON.TYPE one level so
        # legacy paths return scalars and missing keys return ``None``,
        # and re-encode native floats inside JSON.RESP back to string
        # form. NUMINCRBY/NUMMULTBY use the command path captured at the
        # call site to unwrap legacy paths while preserving JSONPath arrays.
        _RESP3_TO_RESP2_LEGACY_MODULE_CALLBACKS = {
            "JSON.NUMINCRBY": self._decode_resp3_legacy_numop,
            "JSON.NUMMULTBY": self._decode_resp3_legacy_numop,
            "JSON.OBJKEYS": self._decode,
            "JSON.RESP": self._resp_floats_to_str_command,
            "JSON.TYPE": lambda r: r[0] if isinstance(r, list) and len(r) == 1 else r,
        }

        self.client = client
        self.execute_command = client.execute_command
        self.MODULE_VERSION = version

        self._MODULE_CALLBACKS = apply_module_callbacks(
            get_protocol_version(self.client),
            get_legacy_responses(self.client),
            common=_MODULE_CALLBACKS,
            resp2=_RESP2_MODULE_CALLBACKS,
            resp3=_RESP3_MODULE_CALLBACKS,
            resp2_unified=_RESP2_UNIFIED_MODULE_CALLBACKS,
            resp3_unified=_RESP3_UNIFIED_MODULE_CALLBACKS,
            resp3_to_resp2_legacy=_RESP3_TO_RESP2_LEGACY_MODULE_CALLBACKS,
        )

        for key, value in self._MODULE_CALLBACKS.items():
            self.client.set_response_callback(key, value)

        self.__encoder__ = encoder
        self.__decoder__ = decoder

    def _decode(self, obj):
        """Get the decoder."""
        if obj is None:
            return obj

        try:
            x = self.__decoder__.decode(obj)
            if x is None:
                raise TypeError
            return x
        except TypeError:
            try:
                return self.__decoder__.decode(obj.decode())
            except AttributeError:
                return decode_list(obj)
        except (AttributeError, JSONDecodeError):
            return decode_list(obj)

    def _decode_json_numop(self, obj, **kwargs):
        """Decode a JSON.NUMINCRBY / JSON.NUMMULTBY result and normalise
        it to the unified array form.

        RESP2 wire returns a JSON bulk string — a scalar for legacy
        paths and a JSON-encoded list for dollar paths. RESP3 wire
        returns a native list. The unified shape is always a list.
        """
        if obj is None:
            return obj
        try:
            result = self.__decoder__.decode(
                obj if isinstance(obj, str) else obj.decode()
            )
        except (AttributeError, JSONDecodeError):
            return obj
        if not isinstance(result, list):
            result = [result]
        return result

    def _decode_resp3_legacy_numop(self, obj, **kwargs):
        """Decode RESP3 JSON numeric operations back to legacy RESP2 shape.

        RedisJSON returns a RESP3 array for both legacy paths (``.foo``)
        and JSONPath paths (``$.foo``). The command builder records the
        path so default RESP3-wire legacy clients can keep v8 scalar
        results for legacy paths while leaving JSONPath results as lists.
        """
        result = self._decode(obj)
        path = kwargs.get("_json_path")
        if (
            isinstance(path, str)
            and not path.startswith("$")
            and isinstance(result, list)
            and len(result) == 1
        ):
            return result[0]
        return result

    def _decode_resp_command_unified(self, obj):
        """Decode JSON.RESP and lift string-encoded floats inside the
        nested response to native ``float`` values so the unified shape
        matches the RESP3 wire.
        """
        return self._convert_resp_floats(self._decode(obj))

    def _resp_floats_to_str_command(self, obj):
        """Decode JSON.RESP and re-encode native ``float`` values back
        to their string form so the legacy RESP2 shape is preserved
        when the wire is RESP3.
        """
        return self._resp_floats_to_str(self._decode(obj))

    @staticmethod
    def _convert_resp_floats(obj):
        """Recursively convert string-encoded JSON floats.

        RESP2 has no native double type, so JSON.RESP returns JSON floats as
        bulk strings. Any string/bytes leaf that parses as ``float`` is the
        unified representation of such a value; non-numeric strings are left
        untouched.
        """
        if isinstance(obj, list):
            return [_JSONBase._convert_resp_floats(item) for item in obj]
        if isinstance(obj, (str, bytes)):
            value = obj.decode() if isinstance(obj, bytes) else obj
            try:
                return float(value)
            except (ValueError, OverflowError):
                return obj
        return obj

    @staticmethod
    def _resp_floats_to_str(obj):
        """Recursively walk ``obj`` and convert native ``float`` values
        back to their string-encoded form. Lists are walked
        element-wise; non-float leaves are returned unchanged.
        """
        if isinstance(obj, list):
            return [_JSONBase._resp_floats_to_str(item) for item in obj]
        if isinstance(obj, float):
            return str(obj)
        return obj

    def _encode(self, obj):
        """Get the encoder."""
        return self.__encoder__.encode(obj)

    def pipeline(self, transaction=True, shard_hint=None):
        """Creates a pipeline for the JSON module, that can be used for executing
        JSON commands, as well as classic core commands.

        Usage example:

        r = redis.Redis()
        pipe = r.json().pipeline()
        pipe.jsonset('foo', '.', {'hello!': 'world'})
        pipe.jsonget('foo')
        pipe.jsonget('notakey')
        """
        if isinstance(self.client, redis.RedisCluster):
            p = ClusterPipeline(
                nodes_manager=self.client.nodes_manager,
                commands_parser=self.client.commands_parser,
                startup_nodes=self.client.nodes_manager.startup_nodes,
                result_callbacks=self.client.result_callbacks,
                cluster_response_callbacks=self.client.cluster_response_callbacks,
                cluster_error_retry_attempts=self.client.retry.get_retries(),
                read_from_replicas=self.client.read_from_replicas,
                reinitialize_steps=self.client.reinitialize_steps,
                lock=self.client._lock,
            )

        else:
            p = Pipeline(
                connection_pool=self.client.connection_pool,
                response_callbacks=dict(self.client.response_callbacks),
                transaction=transaction,
                shard_hint=shard_hint,
            )

        p._encode = self._encode
        p._decode = self._decode
        return p


class ClusterPipeline(JSONCommands, redis.cluster.ClusterPipeline):
    """Cluster pipeline for the module."""


class Pipeline(JSONCommands, redis.client.Pipeline):
    """Pipeline for the module."""


class JSON(_JSONBase):
    _is_async_client: Literal[False] = False


class AsyncJSON(_JSONBase):
    _is_async_client: Literal[True] = True

    async def set_file(
        self,
        name: str,
        path: str,
        file_name: str,
        nx: bool | None = False,
        xx: bool | None = False,
        decode_keys: bool | None = False,
        fpha: FPHAType | str | None = None,
    ) -> bool | None:
        """
        Set the JSON value at key ``name`` under the ``path`` to the content
        of the json file ``file_name``.

        This runs the blocking file read in a thread pool to avoid blocking
        the event loop.
        """

        def _read_file(fp: str) -> dict:
            with open(fp) as f:
                return loads(f.read())

        file_content = await asyncio.to_thread(_read_file, file_name)
        return await self.set(
            name,
            path,
            file_content,
            nx=nx,
            xx=xx,
            decode_keys=decode_keys,
            fpha=fpha,
        )

    async def set_path(
        self,
        json_path: str,
        root_folder: str,
        nx: bool | None = False,
        xx: bool | None = False,
        decode_keys: bool | None = False,
        fpha: FPHAType | str | None = None,
    ) -> dict[str, bool]:
        """
        Iterate over ``root_folder`` and set each JSON file to a value
        under ``json_path`` with the file name as the key.

        This method runs blocking filesystem operations (os.walk and file reads)
        in a thread pool to avoid blocking the event loop.
        """

        def _walk_directory(folder: str) -> list[str]:
            """Walk directory and return list of file paths (runs in thread pool)."""
            file_paths = []
            for root, dirs, files in os.walk(folder):
                for file in files:
                    file_paths.append(os.path.join(root, file))
            return file_paths

        set_files_result = {}

        # Run blocking os.walk in thread pool
        file_paths = await asyncio.to_thread(_walk_directory, root_folder)

        for file_path in file_paths:
            try:
                # TODO: rsplit(".") splits on all dots, mishandling paths
                # with dots in directories (e.g. /data/v1.2/file.json).
                # Should be rsplit(".", 1) — fix in a separate PR.
                file_name = file_path.rsplit(".")[0]
                await self.set_file(
                    file_name,
                    json_path,
                    file_path,
                    nx=nx,
                    xx=xx,
                    decode_keys=decode_keys,
                    fpha=fpha,
                )
                set_files_result[file_path] = True
            except JSONDecodeError:
                set_files_result[file_path] = False

        return set_files_result