import json
import logging
from collections import defaultdict
from typing import List
from typing import Union

from twisted.internet.defer import Deferred
from vortex.DeferUtil import deferToThreadWrapWithLogger
from vortex.Payload import Payload
from vortex.TupleSelector import TupleSelector
from vortex.handler.TupleDataObservableHandler import TuplesProviderABC

from peek_plugin_diagram._private.client.controller.BranchIndexCacheController import (
    BranchIndexCacheController,
)
from peek_plugin_diagram._private.storage.branch.BranchIndexEncodedChunk import (
    BranchIndexEncodedChunk,
)
from peek_plugin_diagram._private.tuples.branch.BranchTuple import BranchTuple
from peek_plugin_diagram._private.worker.tasks.branch._BranchIndexCalcChunkKey import (
    makeChunkKeyForBranchIndex,
)

logger = logging.getLogger(__name__)


class BranchTupleProvider(TuplesProviderABC):
    def __init__(self, cacheHandler: BranchIndexCacheController):
        self._cacheHandler = cacheHandler

    @deferToThreadWrapWithLogger(logger)
    def makeVortexMsg(
        self, filt: dict, tupleSelector: TupleSelector
    ) -> Union[Deferred, bytes]:
        modelSetKey = tupleSelector.selector["modelSetKey"]
        coordSetId = tupleSelector.selector["coordSetId"]
        keys = tupleSelector.selector["keys"]

        keysByChunkKey = defaultdict(list)

        results: List[BranchTuple] = []

        for key in keys:
            keysByChunkKey[makeChunkKeyForBranchIndex(modelSetKey, key)].append(
                key
            )

        for chunkKey, branchKeys in keysByChunkKey.items():
            chunk: BranchIndexEncodedChunk = self._cacheHandler.encodedChunk(
                chunkKey
            )

            if not chunk:
                logger.warning(
                    "BranchIndex chunk %s is missing from cache", chunkKey
                )
                continue

            resultsByKeyStr = (
                Payload().fromEncodedPayload(chunk.encodedData).tuples[0]
            )
            resultsByKey = json.loads(resultsByKeyStr)

            for branchKey in branchKeys:
                if branchKey not in resultsByKey:
                    logger.warning(
                        "Branch %s is missing from index, chunkKey %s",
                        branchKey,
                        chunkKey,
                    )
                    continue

                packedJsons = resultsByKey[branchKey]

                for packedJson in packedJsons:
                    result = BranchTuple.loadFromJson(packedJson, None, None)
                    if result.coordSetId == coordSetId or coordSetId is None:
                        results.append(result)

        # Create the vortex message
        return Payload(filt, tuples=results).makePayloadEnvelope().toVortexMsg()
