from typing import Any, Dict, List
from .parser import parser, PipelineComponent, CallSequence, ParallelCallSequences, SystemCall
from ..processors.registry import ProcessorRegistry
from ..utils import dict_topk

import asyncio

# TODO: probably should unify it...
_role_to_service = {
    'search': 'search', 'rerank': 'score', 'expander': 'decompose_query', 'merger': 'fuse'
}


class SearchPipeline:
    def __init__(
            self, 
            pipeline: PipelineComponent, 
            collection: str,
            runtime_kwargs: Dict[str, Dict[str, Any]] = None,
            verify: bool = True
        ):
        self.pipeline = pipeline
        self.collection = collection
        self.runtime_kwargs = runtime_kwargs or {}
        self.doc_content_cache = {}

        if verify: 
            self.verify()
        alias_not_found = set(self.runtime_kwargs.keys()) - set([ c.alias for c in self.pipeline.all_calls ])
        assert len(alias_not_found) == 0, f"Cannot find alias {alias_not_found}"

    def verify(self):
        assert ProcessorRegistry.has_service(self.collection, "content")
        for call in self.pipeline.all_calls:
            assert ProcessorRegistry.has_service(call.name, _role_to_service[call.role]), \
                f"Cannot find a {_role_to_service[call.role]} service under `{call.name}`"
        
    @classmethod
    def from_string(
            cls, 
            pipeline_string: str, 
            collection: str, 
            runtime_kwargs: Dict[str, Dict[str, Any]] = None, 
            verify: bool = True
        ) -> 'SearchPipeline':
        return cls(parser.parse(pipeline_string), collection, runtime_kwargs, verify)

    async def get_doc_content(self, doc_id: str):
        if doc_id not in self.doc_content_cache:
            ret = await ProcessorRegistry[self.collection]["content"].submit({'id': doc_id})
            self.doc_content_cache[doc_id] = ret['text']
        return self.doc_content_cache[doc_id]

    async def run(
            self, 
            query: str, 
            last_output: Dict[str, Any] = None,
            current_node: PipelineComponent = None, 
            scratch: Dict[str, Dict[str, Any]] = None
        ) -> List[Dict[str, float]]:
        current_node = current_node or self.pipeline
        last_output = last_output or {}
        scratch = scratch or {}

        if isinstance(current_node, CallSequence):
            for stage in current_node.stages:
                last_output = await self.run(
                    query, last_output, stage, scratch
                )
            return last_output
    
        if isinstance(current_node, ParallelCallSequences):
            expanded_queries = [query]
            if current_node.expander is not None:
                last_output = await self.run(
                    query, last_output, current_node.expander, scratch
                )
                expanded_queries = last_output['queries']
                
            concurrent_run_outputs = await asyncio.gather(*[
                self.run(q, last_output, seq, scratch)
                for seq in current_node.sequences for q in expanded_queries
            ])

            return await self.run(
                query, {'scores': [ o['scores'] for o in concurrent_run_outputs ], **last_output},
                current_node.merger, scratch
            )

        assert isinstance(current_node, SystemCall)

        payload = {
            'query': query, 
            **last_output, 
            **scratch,
            **self.runtime_kwargs.get(current_node.alias, {})
        }
        if current_node.limit is not None:
            payload['limit'] = current_node.limit

        processor = ProcessorRegistry[current_node.name][_role_to_service[current_node.role]]

        if current_node.role == 'search':
            ret = await processor.submit(payload)
            assert 'scores' in ret and isinstance(ret['scores'], dict)
        
        if current_node.role == 'merger':
            ret = await processor.submit(payload)
            assert 'scores' in ret and isinstance(ret['scores'], dict)

        if current_node.role == 'rerank':
            docid_to_rerank: List[str] = list(last_output['scores'].keys())
            doc_text_list = await asyncio.gather(*[self.get_doc_content(d) for d in docid_to_rerank])
            payload['passages'] = doc_text_list
            ret = await processor.submit(payload)
            assert 'scores' in ret and isinstance(ret['scores'], list)
            ret['scores'] = dict(zip(docid_to_rerank, ret['scores']))

        if current_node.role == 'expander':
            ret = await processor.submit(payload)
            assert 'queries' in ret and isinstance(ret['queries'], list)

        # apply limit here just to safe
        if current_node.limit is not None:
            if 'scores' in ret:
                ret['scores'] = dict_topk(ret['scores'], current_node.limit)
            elif 'queries' in ret:
                ret['queries'] = ret['queries'][:current_node.limit]

        scratch[(current_node.alias, current_node.role)] = ret
        return ret