from typing import List, Dict, Any, Optional, Generator

import logging
import threading

logger: logging.Logger = logging.getLogger(__name__)

class ServerProxy:
    '''This class is used to transform model, database, storage and vector payloads into PayloadStructs before sending it to the RemoteEngineRunReactor via the runPixel endpoint'''
    
    def __init__(
        self,
    ):
        '''
        Initialize the ServerProxy instance.
        '''
        self.epoc = 0
        self.condition = threading.Condition()
        from .server_client import ServerClient
        self.server = ServerClient.da_server
        assert self.server is not None, "Please authenticate using your access and secret keys"
        logger.info(f"Server is set")

    def get_next_epoc(self) -> str:
        '''This method atomically increments the epoc count by one plus the current value.'''
        self.epoc = self.epoc + 1
        return f"py_{self.epoc}"

    def comm(
        self,    
        epoc:str,
        engine_type:str,
        engine_id:str,
        method_name:str,
        method_args:Optional[List[Any]] = [],
        method_arg_types:Optional[List[str]] = [],
        insight_id: Optional[str] = None
    ) -> None:
        """
        This method in responsible for:
            - converting the args into a PayloadStruct
            - adds itself to the clients monitor block
            - forwards the PayloadStruct to the client to be sent to the Tomcat Server

        Args:
            epoc (`str`): The epoc ID for the payload struct
            engine_type (`str`): The engine type that will be called from the tomcat server. Options are model, storage, database or vector and are set in NativePyEngineWorker.java
            engine_id (`str`): The unique identifier of the engine being called. This passed so the tomcat server can call Utility.java to find the engine
            method_name (`str`): The IEngine method name that is available in the engine_type
            method_args (`Optional[List[Any]]`): A list of object to be sent to the IEngine method as inputs
            method_arg_types (`Optional[List[str]]`): A list of Java class names that represent the method args types
            insight_id (`Optional[str]`): The unique identifier for the insight
        
        Returns:
            `None`
        """

        # converts this into a PayloadStruct
        payload = {
            "epoc": epoc,
            "response": False,
            "engineType": engine_type,
            "interim": False,
            "objId": engine_id,  # all the method stuff will come here and below
            "methodName": method_name,
            "payload": method_args,
            "payloadClassNames": method_arg_types,
            "insightId": insight_id if insight_id is not None else self.server.cur_insight,
            "operation": "ENGINE",
        }
        
        self.server.monitors.update({epoc: self.condition})  # adds itself to the monitor block
        logger.info(f"The server is RESTServer object")
        self.server.send_request(payload)

    def call(
        self,
        engine_type:str,
        engine_id:str,
        method_name:str = None,
        method_args:Optional[List[Any]] = [],
        method_arg_types:Optional[List[str]] = [],
        insight_id: Optional[str] = None
    ) -> List[Dict]:
        '''
        This method is responsible for initiating a synchronous communication with the server, which calls the `comm` method.
        
        Args:
            epoc (`str`): The epoc ID for the payload struct.
            engine_type (`str`): The engine type that will be called from the Tomcat server. Options are model, storage, database, or vector and are set in NativePyEngineWorker.java.
            engine_id (`str`): The unique identifier of the engine being called. This is passed so the Tomcat server can call Utility.java to find the engine.
            method_name (`str`): The IEngine method name that is available in the `engine_type`.
            method_args (`Optional[List[Any]]`): A list of objects to be sent to the IEngine method as inputs.
            method_arg_types (`Optional[List[str]]`): A list of Java class names that represent the method argument types.
            insight_id (`Optional[str]`): The unique identifier for the insight
            
        Returns:
            `List[Dict]`: A list that contains the response from the Tomcat server engine.
        '''
        epoc = self.get_next_epoc()
        self.comm(
            epoc=epoc,
            engine_type=engine_type,
            engine_id=engine_id,
            method_name=method_name,
            method_args=method_args,
            method_arg_types=method_arg_types,
            insight_id=insight_id,
        )
        
        new_payload_struct = self.server.monitors.pop(epoc)
        logger.info(new_payload_struct)
        
        if "ex" in new_payload_struct:
            # if exception, convert it to an Exception and raise it
            raise Exception(new_payload_struct["ex"])
        else:
            logger.info(f"answer is .. {new_payload_struct['payload']}")
            return new_payload_struct["payload"]
        
    def async_call(
        self,
        engine_type:str,
        engine_id:str,
        method_name:str = None,
        method_args:Optional[List[Any]] = [],
        method_arg_types:Optional[List[str]] = [],
        insight_id: Optional[str] = None
    ) -> str:
        '''
        This is really a call in a separate thread.
        
        This method is responsible for initiating a synchronous communication with the server, which calls the `comm` method.
        
        Args:
            epoc (`str`): The epoc ID for the payload struct.
            engine_type (`str`): The engine type that will be called from the Tomcat server. Options are model, storage, database, or vector and are set in NativePyEngineWorker.java.
            engine_id (`str`): The unique identifier of the engine being called. This is passed so the Tomcat server can call Utility.java to find the engine.
            method_name (`str`): The IEngine method name that is available in the `engine_type`.
            method_args (`Optional[List[Any]]`): A list of objects to be sent to the IEngine method as inputs.
            method_arg_types (`Optional[List[str]]`): A list of Java class names that represent the method argument types.
            insight_id (`Optional[str]`): The unique identifier for the insight
            
        Returns:
            `List[Dict]`: A list that contains the response from the Tomcat server engine.
        '''
        insight_id = insight_id if insight_id is not None else self.server.cur_insight
        epoc = self.get_next_epoc()
        threading.Thread(
            target=self.comm,
            kwargs={
                "epoc": epoc,
                "engine_type": engine_type,
                "engine_id": engine_id,
                "method_name": method_name,
                "method_args": method_args,
                "method_arg_types": method_arg_types,
                "insight_id": insight_id,
            },
        ).start()
        
        return insight_id
        
    def get_partial_responses(self, job_id:str) -> Generator:
        '''
        Use the client connection to make posts to the partial endpoint and create a generator.
        
        Args:
            job_id (`str`): The identifier of the job being run
            
        Returns:
            `Generator`: A generator that can be iterated through to get partials
        '''  
        yield from self.server.get_partial_responses(job_id)
        
    def test(self):
        epoc = self.get_next_epoc()
        thread = threading.Thread(
            target=self.comm,
            kwargs={
                "epoc": epoc,
                "engine_type": "ENGINE",
                "engine_id": "hello",
                "method_name": "method1",
                "method_args": ["cat", "bat", None],
                "method_arg_types": [
                    "java.lang.String",
                    "java.lang.String",
                    "java.lang.String",
                ],
            },
        )

        thread.start()
        thread.join()
        new_payload_struct = self.server.monitors.pop(epoc)
        # print(new_payload_struct)
        # if exception
        # convert exception and give back
        if "ex" in new_payload_struct:
            return new_payload_struct["ex"]
        else:
            # if we get to the point for json pickle we can do that
            # new_payload_struct = process_payload(new_payload_struct)
            return new_payload_struct["payload"]

    def process_payload(self, payload_struct):
        # try to see if the types are pickle
        # if so unpickle it
        import jsonpickle as jp

        payload_data = None
        if "payload" in payload_struct:
            payload_data = payload_struct["payload"]
        if payload_data is not None and isinstance(payload_data,List):
            for data in payload_data:
                index = payload_data.index(data)
                try:
                    orig_obj = data
                    obj = jp.loads(orig_obj)
                    payload_struct["payload"][index] = obj
                except Exception as e:
                    pass

        return payload_struct
