# Copyright 2022-2023 XProbe Inc.
# derived from copyright 1999-2021 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
import itertools
import logging
from collections import defaultdict
from typing import Dict, List, Optional

import xoscar as mo

from ....lib.aio import alru_cache
from ...meta.api import MetaAPI
from ...storage.api import StorageAPI
from ..errors import TileableNotTracked

logger = logging.getLogger(__name__)


class LifecycleTrackerActor(mo.Actor):
    _meta_api: MetaAPI

    def __init__(self, session_id: str):
        self._session_id = session_id
        self._tileable_key_to_chunk_keys = dict()
        self._tileable_ref_counts = defaultdict(lambda: 0)
        self._chunk_ref_counts = defaultdict(lambda: 0)

        self._meta_api: Optional[MetaAPI] = None

    async def __post_create__(self):
        self._meta_api = await MetaAPI.create(self._session_id, self.address)

    async def __pre_destroy__(self):
        chunk_keys = [
            chunk_key
            for chunk_key, ref_count in self._chunk_ref_counts.items()
            if ref_count > 0
        ]
        # remove all chunks
        await self._remove_chunks(chunk_keys)

    @alru_cache
    async def _get_task_api(self):
        from ...task.api import TaskAPI

        return await TaskAPI.create(self._session_id, self.address)

    @staticmethod
    def gen_uid(session_id):
        return f"{session_id}_lifecycle_tracker"

    def _track(self, tileable_key: str, chunk_keys: List[str]):
        if tileable_key not in self._tileable_key_to_chunk_keys:
            self._tileable_key_to_chunk_keys[tileable_key] = []
        chunk_keys_set = set(self._tileable_key_to_chunk_keys[tileable_key])
        incref_chunk_keys = []
        tileable_ref_count = self._tileable_ref_counts.get(tileable_key, 0)
        for chunk_key in chunk_keys:
            if chunk_key in chunk_keys_set:
                continue
            if tileable_ref_count > 0:
                incref_chunk_keys.extend([chunk_key] * tileable_ref_count)
            self._tileable_key_to_chunk_keys[tileable_key].append(chunk_key)
        if incref_chunk_keys:
            self._incref_chunks(incref_chunk_keys)

    @mo.extensible
    async def track(self, tileable_key: str, chunk_keys: List[str]):
        return await asyncio.to_thread(self._track, tileable_key, chunk_keys)

    @classmethod
    def _check_ref_counts(cls, keys: List[str], ref_counts: List[int]):
        if ref_counts is not None and len(keys) != len(ref_counts):
            raise ValueError(
                f"`ref_counts` should have same size as `keys`, expect {len(keys)}, got {len(ref_counts)}"
            )

    def _incref_chunks(self, chunk_keys: List[str], counts: List[int] = None):
        counts = counts if counts is not None else itertools.repeat(1)
        for chunk_key, count in zip(chunk_keys, counts):
            self._chunk_ref_counts[chunk_key] += count

    async def incref_chunks(self, chunk_keys: List[str], counts: List[int] = None):
        self._check_ref_counts(chunk_keys, counts)
        return await asyncio.to_thread(self._incref_chunks, chunk_keys, counts=counts)

    def _get_remove_chunk_keys(self, chunk_keys: List[str], counts: List[int] = None):
        to_remove_chunk_keys = []
        counts = counts if counts is not None else itertools.repeat(1)
        for chunk_key, count in zip(chunk_keys, counts):
            ref_count = self._chunk_ref_counts[chunk_key]
            ref_count -= count
            assert ref_count >= 0, f"chunk key {chunk_key} will have negative ref count"
            self._chunk_ref_counts[chunk_key] = ref_count
            if ref_count == 0:
                # remove
                to_remove_chunk_keys.append(chunk_key)
        return to_remove_chunk_keys

    async def decref_chunks(self, chunk_keys: List[str], counts: List[int] = None):
        self._check_ref_counts(chunk_keys, counts)
        to_remove_chunk_keys = await asyncio.to_thread(
            self._get_remove_chunk_keys, chunk_keys, counts=counts
        )
        # make _remove_chunks release actor lock so that multiple `decref_chunks` can run concurrently.
        yield self._remove_chunks(to_remove_chunk_keys)

    async def _remove_chunks(self, to_remove_chunk_keys: List[str]):
        # get meta
        logger.debug(
            "Remove chunks %.500s with a refcount of zero", to_remove_chunk_keys
        )
        get_metas = []
        for to_remove_chunk_key in to_remove_chunk_keys:
            get_metas.append(
                self._meta_api.get_chunk_meta.delay(
                    to_remove_chunk_key, fields=["bands"], error="ignore"
                )
            )
        metas = await self._meta_api.get_chunk_meta.batch(*get_metas)

        # filter chunks that not exist
        new_to_remove_chunk_keys = []
        new_metas = []
        for to_remove_chunk_key, meta in zip(to_remove_chunk_keys, metas):
            if meta is not None:
                new_to_remove_chunk_keys.append(to_remove_chunk_key)
                new_metas.append(meta)
        to_remove_chunk_keys = new_to_remove_chunk_keys
        metas = new_metas

        all_bands = [meta["bands"] for meta in metas]
        key_to_addresses = dict()
        for to_remove_chunk_key, bands in zip(to_remove_chunk_keys, all_bands):
            key_to_addresses[to_remove_chunk_key] = bands

        # remove data via storage API
        storage_api_to_deletes = defaultdict(list)
        for key, bands in key_to_addresses.items():
            for band in bands:
                # storage API is cached for same arguments
                storage_api = await StorageAPI.create(
                    self._session_id, band[0], band[1]
                )
                storage_api_to_deletes[storage_api].append(
                    storage_api.delete.delay(key, error="ignore")
                )
        await asyncio.gather(
            *[
                storage_api.delete.batch(*deletes)
                for storage_api, deletes in storage_api_to_deletes.items()
            ]
        )

        # delete meta
        delete_metas = []
        for to_remove_chunk_key in to_remove_chunk_keys:
            delete_metas.append(
                self._meta_api.del_chunk_meta.delay(to_remove_chunk_key)
            )
        await self._meta_api.del_chunk_meta.batch(*delete_metas)

    def get_chunk_ref_counts(self, chunk_keys: List[str]) -> List[int]:
        return [self._chunk_ref_counts[chunk_key] for chunk_key in chunk_keys]

    def get_all_chunk_ref_counts(self) -> Dict[str, int]:
        result = dict()
        for chunk_key, ref_count in self._chunk_ref_counts.items():
            if ref_count > 0:
                result[chunk_key] = ref_count
        return result

    def _incref_tileables(self, tileable_keys: List[str], counts: List[int] = None):
        counts = counts if counts is not None else itertools.repeat(1)
        for tileable_key, count in zip(tileable_keys, counts):
            if tileable_key not in self._tileable_key_to_chunk_keys:
                raise TileableNotTracked(f"tileable {tileable_key} not tracked before")
            self._tileable_ref_counts[tileable_key] += count
            incref_chunk_keys = self._tileable_key_to_chunk_keys[tileable_key]
            # incref chunks for this tileable
            logger.debug(
                "Incref chunks %.500s while increfing tileable %s",
                incref_chunk_keys,
                tileable_key,
            )
            chunk_counts = None if count == 1 else [count] * len(incref_chunk_keys)
            self._incref_chunks(incref_chunk_keys, counts=chunk_counts)

    async def incref_tileables(
        self, tileable_keys: List[str], counts: List[int] = None
    ):
        self._check_ref_counts(tileable_keys, counts)
        return await asyncio.to_thread(
            self._incref_tileables, tileable_keys, counts=counts
        )

    def _get_decref_chunk_keys(
        self, tileable_keys: List[str], counts: List[int] = None
    ) -> Dict[str, int]:
        decref_chunk_keys = dict()
        counts = counts if counts is not None else itertools.repeat(1)
        for tileable_key, count in zip(tileable_keys, counts):
            if tileable_key not in self._tileable_key_to_chunk_keys:
                raise TileableNotTracked(f"tileable {tileable_key} not tracked before")
            self._tileable_ref_counts[tileable_key] -= count

            for chunk_key in self._tileable_key_to_chunk_keys[tileable_key]:
                if chunk_key not in decref_chunk_keys:
                    decref_chunk_keys[chunk_key] = count
                else:
                    decref_chunk_keys[chunk_key] += count
        logger.debug(
            "Decref chunks %.500s while decrefing tileables %s",
            decref_chunk_keys,
            tileable_keys,
        )
        return decref_chunk_keys

    async def decref_tileables(
        self, tileable_keys: List[str], counts: List[int] = None
    ):
        self._check_ref_counts(tileable_keys, counts)
        decref_chunk_key_to_counts = await asyncio.to_thread(
            self._get_decref_chunk_keys, tileable_keys, counts=counts
        )
        to_remove_chunk_keys = await asyncio.to_thread(
            self._get_remove_chunk_keys,
            list(decref_chunk_key_to_counts),
            counts=list(decref_chunk_key_to_counts.values()),
        )
        to_remove_tileable_keys = await asyncio.to_thread(
            list, (key for key in tileable_keys if self._tileable_ref_counts[key] <= 0)
        )
        coros = []
        if to_remove_chunk_keys:
            coros.append(self._remove_chunks(to_remove_chunk_keys))
        if to_remove_tileable_keys:
            task_api = await self._get_task_api()
            coros.append(task_api.remove_tileables(to_remove_tileable_keys))
        if coros:
            # release actor lock
            yield asyncio.gather(*coros)

    def get_tileable_ref_counts(self, tileable_keys: List[str]) -> List[int]:
        return [
            self._tileable_ref_counts[tileable_key] for tileable_key in tileable_keys
        ]
