import asyncio
import inspect
import json
import math
import time
import zlib
from collections import namedtuple
from contextvars import ContextVar
from datetime import datetime
from http import HTTPStatus
from typing import Any, Callable, Coroutine, Dict, Iterable, Iterator, List, Tuple, Type, TypeVar

import zstandard as zstd
from aiohttp import ClientConnectionError, ClientPayloadError, ClientResponseError, ClientSession

import config
import context
from server.apple.constants import US_MARKET
from server.core.constants import (APPLE, CONCURRENCY_LEVEL, DEFAULT_SEARCH_MARKETS, GLOBAL_MARKET, GLOBAL_MARKET_CODE,
                                   SEARCH_ID, SEARCH_ISRC, SPOTIFY, TRACK_SEARCH_SUPPORTED_ID_TYPES,
                                   TRACKS_SEARCH_SUPPORTED_DSP, CompressionLib)
from server.core.exceptions import APIError

COMPRESSION_LIBS = {CompressionLib.ZLIB: zlib, CompressionLib.ZSTD: zstd}


def iter_chunk(iterable: list, chunk_size: int = 1) -> Iterator:
    """Split collection of items on chunks.

    :param iterable: Collection of items.
    :param chunk_size: Size of chunks.

    :return: Iterator of chunks.
    """

    count = len(iterable)
    for i in range(0, count, chunk_size):
        yield iterable[i : min(i + chunk_size, count)]  # noqa


def deep_update(original_dict: dict, update_dict: dict) -> dict:
    """Extend data in original dict from update_dict.

    :param original_dict: Original dict to update.
    :param update_dict: Dict with new content to insert into original.

    :return: Updated original dict.
    """

    for key, value in update_dict.items():
        original_value = original_dict.get(key)
        if isinstance(original_value, list):
            original_dict[key] = original_value + value
        elif isinstance(original_value, dict):
            original_dict[key] = deep_update(original_value, value)
        else:
            original_dict[key] = value
    return original_dict


class ListResult:
    """Result class to represent list type return value.
    Use with request_in_chunks decorator.
    """

    result_cls = list

    @classmethod
    def update(cls, result: list, inner_result: list) -> list:
        result.extend(inner_result)
        return result


class DictResult:
    """Result class to represent dict type return value.
    Use with request_in_chunks decorator.
    """

    result_cls = dict

    @classmethod
    def update(cls, result: dict, inner_result: dict) -> dict:
        result = deep_update(result, inner_result)
        return result


ResultType = TypeVar("ResultType", ListResult, DictResult)


def request_in_chunks(
    chunk_size: int = None,
    chunks_count: int = CONCURRENCY_LEVEL,
    items_index: int = 1,
    result_type: ResultType = DictResult,
) -> Callable:
    """Execute API requests for chunks of items.
    Executes original request by synchronous chunks with inner asynchronous requests equal to chunk_count each
    of chunk_size size.

        :param chunk_size: A size of a chunk.
        :param chunks_count: A number of chunks to run asynchronously.
        :param items_index: Args list index.
        :param result_type: Wrapped function result type class
        :return: Wrapped function.
    """

    def wrapper(f: Callable):
        async def wrapped(*args, **kwargs):
            items = args[items_index]
            result = result_type.result_cls()
            sequential_chunk_size = chunks_count * (chunk_size or 1)

            start_time = datetime.utcnow()

            for sequential_chunk in iter_chunk(items, sequential_chunk_size):
                if chunk_size:
                    tasks = []
                    for concurrent_chunk in iter_chunk(sequential_chunk, chunk_size):
                        args_list = list(args)
                        args_list[items_index] = concurrent_chunk
                        tasks.append(f(*tuple(args_list), **kwargs))
                    chunk_inner_results = await asyncio.gather(*tasks)
                    chunk_result = result_type.result_cls()
                    for inner_result in chunk_inner_results:
                        chunk_result = result_type.update(chunk_result, inner_result)
                else:
                    args_list = list(args)
                    args_list[items_index] = sequential_chunk
                    chunk_result = await f(*tuple(args_list), **kwargs)
                result = result_type.update(result, chunk_result)

            if config.DEBUG_REQUEST_COUNT:
                count = math.ceil(len(items) / chunk_size if chunk_size else chunks_count)
                set_debug_request_count(count, start_time, datetime.utcnow())

            return result

        wrapped.__signature__ = inspect.signature(f)
        return wrapped

    return wrapper


def get_id(data: dict) -> str:
    """Get id from apple playlist/track data."""
    return data.get("id")


async def handle_requests(
    requests: Iterable[Tuple[Callable, dict, Any, bool] or Tuple[Callable, dict]]
) -> Iterable[Any]:
    """Make multiple requests in parallel and return response or default

    Args:
        requests: Requests definition: function, kwargs, default value, condition.

    Returns:
        Response or default.

    """
    _requests = []
    for r in requests:
        r_len = len(r)
        if r_len == 5:
            _requests.append(r)
        elif r_len == 4:
            _requests.append((r[0], tuple(), r[1], r[2], r[3]))
        elif r_len == 3:
            _requests.append((r[0], r[1], r[2], None, True))
        elif r_len == 2:
            _requests.append((r[0], tuple(), r[1], None, True))
        else:
            raise ValueError(f"handle_requests got unsupported configuration: {r}")

    tasks = [i[0](*i[1], **i[2]) for i in _requests if i[4]]

    responses = await asyncio.gather(*tasks)

    result = []
    response_index = 0
    for index, request in enumerate(_requests):
        if request[4]:
            result.append(responses[response_index])
            response_index += 1
        else:
            result.append(request[3])

    return result


async def make_request(
    session: ClientSession,
    url: str,
    method: str = "GET",
    params: Dict = None,
    headers: Dict = None,
    auth: Type[namedtuple] = None,
    body: Dict or None = None,
    data: Dict or None = None,
    error_cls: Type[APIError] = APIError,
    error_message: str = None,
):
    """Make http request, handle errors.

    :param session: Http client session.
    :param url: Request URL.
    :param method: HTTP method.
    :param params: Request params.
    :param headers: Headers.
    :param auth: Auth tuple.
    :param body: Request body.
    :param data: Request data.
    :param error_cls: Error class.
    :param error_message: Custom error message.
    """
    result = None

    if not headers:
        headers = {}
    headers["Connection"] = "keep-alive"

    try:
        async with session.request(
            method, url, params=params, headers=headers, auth=auth, json=body, data=data, raise_for_status=False
        ) as resp:
            if resp.content:
                if resp.content_type == "application/json":
                    result = await resp.json()
                elif resp.content_type == "application/octet-stream":
                    result = await resp.text()
                    result = json.loads(result or "{}")
            try:
                resp.raise_for_status()
            except ClientResponseError as e:
                raise error_cls(
                    detail=error_message or e.message,
                    original_status_code=e.status,
                    original_response=result,
                    headers=e.headers,
                )
    except ClientConnectionError as e:
        raise error_cls(detail=e.strerror if hasattr(e, "strerror") else "Connection error")
    except ClientPayloadError as e:
        raise error_cls(detail=e.strerror if hasattr(e, "strerror") else "Payload error")
    except asyncio.CancelledError:
        raise error_cls(detail="Request was cancelled")
    return result


def retry(
    count: int = config.DEFAULT_RETRY_COUNT,
    wait_rate: int = config.DEFAULT_RETRY_WAIT,
    max_timeout: int = config.DEFAULT_RETRY_MAX_TIMEOUT,
    auth_handler: Callable[[Any, APIError], Coroutine] = None,
):
    """Retry decorator.

    :param count: Retry count.
    :param wait_rate: Wait multiplier (current attempt number * wait_rate seconds).
    :param max_timeout: Max wait timeout.
    :param auth_handler: Auth error handler.
    """

    def inner(f: Callable):
        async def wrapped(self, *args, **kwargs):
            wait_time = 0
            for i in range(count + 1):
                try:
                    result = await f(self, *args, **kwargs)
                    set_debug_retry_stats(i, wait_time)
                    return result
                except APIError as e:
                    # exit if it is the last try
                    if i == count:
                        raise

                    # handle rate limits
                    if e.get_status_code() == HTTPStatus.TOO_MANY_REQUESTS:
                        retry_after = int(e.headers.get("Retry-After", 1))
                        wait_time = wait_time + retry_after
                        # exit if it is needed to wait too much
                        if wait_time > max_timeout:
                            raise
                        await asyncio.sleep(retry_after)
                    # handle auth errors
                    elif e.get_status_code() in (HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN, 419):
                        # try to handle this only first time
                        if i > 0 or not auth_handler:
                            raise
                        await auth_handler(self, e)
                    # basic retry errors handling
                    elif e.get_status_code() in config.RETRY_FOR_STATUS_CODES:
                        # otherwise wait and retry
                        current_time = wait_rate * (i + 1)
                        wait_time = wait_time + current_time
                        if wait_time > max_timeout:
                            raise
                        await asyncio.sleep(current_time)
                    # exit if it is an error where retry doesn't help
                    else:
                        raise

        return wrapped

    return inner


def _get_page(offset: int, limit: int, offset_index: int, limit_index: int, f: Callable, self, *args, **kwargs):
    """Get single page.

    :param offset: Page offset.
    :param limit: Page limit.
    :param offset_index: Offset args index.
    :param limit_index: Limit args index.
    """
    modified_args = list(args)
    modified_args[offset_index] = offset
    modified_args[limit_index] = limit
    return f(self, *tuple(modified_args), **kwargs)


def spotify_pagination(page_size: int, offset_index: int = -2, limit_index: int = -1):
    """Get full response from spotify.

    :param limit_index: Limit args index.
    :param offset_index: Offset args index.
    :param page_size: Page size.
    """

    def inner(f: Callable):
        async def wrapped(self, *args, **kwargs) -> list or dict:
            if args[limit_index] or args[offset_index]:
                return await f(self, *args, **kwargs)

            page_first = await _get_page(0, page_size, offset_index, limit_index, f, self, *args, **kwargs)

            items = page_first["items"]
            total = page_first["total"]

            if total > page_size:
                page_count = math.ceil((total - page_size) / page_size)

                tasks = [
                    _get_page((i + 1) * page_size, page_size, offset_index, limit_index, f, self, *args, **kwargs)
                    for i in range(page_count)
                ]

                responses = await asyncio.gather(*tasks)

                for page in responses:
                    items = items + page["items"]

            return items

        wrapped.__signature__ = inspect.signature(f)
        return wrapped

    return inner


def make_timestamp(datetime_value: datetime or None) -> int:
    """Datetime to timestamp.

    Args:
        datetime_value: Datetime.

    Returns:
        Timestamp.

    """
    return math.floor(datetime_value.timestamp() * 100) if datetime_value and datetime_value.year > 1970 else 0


def set_debug_retry_stats(count: int = 0, wait_time: int = 0, force_replace: bool = True) -> str:
    """Set debug retry stats.

    Args:
        count: Retry count.
        wait_time: Retry wait seconds.
        force_replace: Replace value if not empty or not.

    Returns:
        Retry stats.

    """
    data = context.RETRY_STATS.get()
    if config.DEBUG_RETRY_STATS and (force_replace or not data):
        data = json.dumps({"count": count, "time": wait_time})
        context.RETRY_STATS.set(data)
    return data


def set_debug_request_count(
    count: int = 0, start_time: datetime = None, end_time: datetime = None, force_replace: bool = True
) -> str:
    """Set debug requests count.

    Args:
        count: Made API reqiests count.
        start_time: Requests started at.
        end_time: Requests were finished at.
        force_replace: Replace non empty stats or not.

    Returns:
        Requests count stats.

    """
    data = context.REQUEST_COUNT.get()
    if config.DEBUG_REQUEST_COUNT and (force_replace or not data):
        data = json.dumps(
            {
                "count": count,
                "from": make_timestamp(start_time if start_time else datetime.utcnow()),
                "to": make_timestamp(end_time if end_time else datetime.utcnow()),
            }
        )
        context.REQUEST_COUNT.set(data)
    return data


def set_debug_data_source(
    source: str = "no",
    min_created_at: datetime = None,
    max_created_at: datetime = None,
    is_single: bool = False,
    force_replace: bool = True,
) -> str:
    """Set debug data source info.

    Args:
        source: Data source.
        min_created_at: Min created at.
        max_created_at: Max created at.
        is_single: Is single cache record request.
        force_replace: Replace non empty stats or not.

    """
    data = context.DATA_SOURCE.get()
    if config.DEBUG_DATA_SOURCE and (force_replace or not data):
        raw_data = {"source": source, "ts": make_timestamp(min_created_at)}
        if not is_single and min_created_at != max_created_at:
            raw_data["max_ts"] = make_timestamp(max_created_at)
        data = json.dumps(raw_data)
        context.DATA_SOURCE.set(data)
    return data


def count_time(context_var: ContextVar):
    """Count execution time.

    Args:
        context_var: Time container.

    """

    def inner(f: Callable):
        async def wrapped(*args, **kwargs):
            start_time = time.time()
            result = await f(*args, **kwargs)
            if config.DEBUG_EXECUTION_TIME:
                context_var.set(round(time.time() - start_time, 3))
            return result

        return wrapped

    return inner


def compress(lib: str or None, data: Any) -> Any:
    """Compress string data to decrease mongodb storage size.

    Args:
        lib: Compression lib name.
        data: Data to compress.

    Returns:
        Compressed data if compression is enabled else original value.

    Raises:
        NotImplementedError: Unknown compression lib name.

    """
    if lib:
        if lib in COMPRESSION_LIBS:
            return COMPRESSION_LIBS[lib].compress(json.dumps(data).encode())
        else:
            raise NotImplementedError()
    return data


def decompress(lib: str or None, data: Any) -> Any:
    """Decompress data from mongo using specific lib.

    Args:
        lib: Compression lib name.
        data: Data to decompress.

    Returns:
        Decompressed data if it was compressed before else data.

    Raises:
        NotImplementedError: Unknown compression lib name.

    """
    if lib:
        if lib in COMPRESSION_LIBS:
            return json.loads(COMPRESSION_LIBS[lib].decompress(data).decode())
        else:
            raise NotImplementedError()
    return data


def get_tracks_search_market(
    dsp: str, market: str = None, use_defaults_for_unsupported: bool = False, dsp_to_available_markets: dict = None
) -> str:
    is_global = market in (GLOBAL_MARKET, GLOBAL_MARKET_CODE)
    if is_global:
        return DEFAULT_SEARCH_MARKETS.get(dsp)
    elif market and (not dsp_to_available_markets or market in dsp_to_available_markets[dsp]):
        return market
    elif use_defaults_for_unsupported:
        return DEFAULT_SEARCH_MARKETS.get(dsp)


def get_another_value(v: Any, values: Iterable[Any]) -> Any:
    for _v in values:
        if _v != v:
            return _v


def get_track_search_id_type(track: dict, search_id_types: Iterable[str] = None) -> str or None:
    search_id_types = search_id_types or config.DSP_TO_TRACK_SEARCH_ID_TYPES[track["dsp"]]
    for t in search_id_types:
        if t in track:
            return t


def _get_spotify_track_search_args_kwargs(search_id_type):
    def f(market, tracks):
        return None, {
            "track_ids": [t[search_id_type] for t in tracks],
            "market": market,
            "by_isrc": search_id_type == SEARCH_ISRC,
        }

    return f


def _get_apple_track_search_by_isrc_args_kwargs(market, tracks):
    return (
        ([t[SEARCH_ISRC] for t in tracks],),
        {"storefront": market, "include": "artists"},
    )


def _get_apple_track_search_by_id_args_kwargs(market, tracks):
    return (
        None,
        {"song_ids": [t[SEARCH_ID] for t in tracks], "storefront": market, "include": "artists"},
    )


async def search_tracks_extended(
    dsp_to_id_type_to_market_to_tracks: Dict[str, Dict[str, Dict[str, List[dict]]]], apple_api, spotify_api
):
    dsp_to_id_type_to_call_and_args_kwargs = {
        APPLE: {
            SEARCH_ID: (apple_api.songs_by_equivalent_id, _get_apple_track_search_by_id_args_kwargs),
            SEARCH_ISRC: (apple_api.songs_by_isrc, _get_apple_track_search_by_isrc_args_kwargs),
        },
        SPOTIFY: {
            SEARCH_ID: (spotify_api.tracks, _get_spotify_track_search_args_kwargs(SEARCH_ID)),
            SEARCH_ISRC: (spotify_api.tracks, _get_spotify_track_search_args_kwargs(SEARCH_ISRC)),
        },
    }

    requests = []
    # we do it this way to keep for requests the same ordering as the input dict has, so we can parse the results
    # regarding the input
    for dsp, id_type_to_market_to_tracks in dsp_to_id_type_to_market_to_tracks.items():
        for id_type, market_to_tracks in id_type_to_market_to_tracks.items():
            for market, tracks in market_to_tracks.items():
                call_func, args_kwargs_getter_func = dsp_to_id_type_to_call_and_args_kwargs[dsp][id_type]
                args, kwargs = args_kwargs_getter_func(market, tracks)
                requests.append((call_func, args or tuple(), kwargs))
    return await handle_requests(requests)


def get_missing_items(
    input_items: List[dict], result_items: List[dict], id_getter, id_result_getter=None
) -> List[dict]:
    if len(input_items) == len(result_items):
        return []
    id_result_getter = id_result_getter or id_getter
    result_items_keys = {id_result_getter(i) for i in result_items}
    return [i for i in input_items if id_getter(i) not in result_items_keys]


DSP_TO_TRACKS_DATA_KEY = {APPLE: "data", SPOTIFY: "tracks"}


DSP_TO_TRACK_ID_GETTER = {
    APPLE: lambda t: t.get("attributes", {})["isrc"],
    SPOTIFY: lambda t: t.get("external_ids", {})["isrc"],
}