import os
import base64
from typing import Optional, Union
from google.protobuf.message import DecodeError
from grpc import Channel, insecure_channel, intercept_channel
from queryservice_client.protobuf.common.configuration_pb2 import ProtobufConfiguration
from queryservice_client.protobuf.common.auth_pb2 import Auth
from queryservice_client.protobuf.rdbms.rdbms_pb2_grpc import RdbmsStub
from queryservice_client.protobuf.rdbms.query_pb2 import (
    StartSessionRequest,
    QueryRequest,
    CancelQueryRequest,
    EndSessionRequest,
)
from queryservice_client.protobuf.servicediscovery.servicediscovery_pb2 import Agent
from queryservice_client.protobuf.servicediscovery.servicediscovery_pb2_grpc import (
    ProxyServiceDiscoveryStub,
)
from queryservice_client.stream.segments.statement import Statement
from queryservice_client.stream.stream import Stream
from queryservice_client.header_manipulator_client_interceptor import header_adder_interceptor
from future.utils import native


class QueryClient:
    DEFAULT_MAX_INBOUND_IN_BYTES = 1024 * 1024 * 64  # 64MiB

    def __init__(self, channel: Channel):
        self.channel = channel
        self.inner = RdbmsStub(self.channel)

    @classmethod
    def from_connector_id(
        cls,
        connector_id: int,
        max_inbound_message_size: int = DEFAULT_MAX_INBOUND_IN_BYTES,
    ):
        host = os.environ.get("CONNECTOR_HOST", "localhost")
        return cls.from_target(
            f"{host}:{10000 + connector_id}", max_inbound_message_size
        )

    @classmethod
    def from_target(
        cls, target: str, max_inbound_message_size: int = DEFAULT_MAX_INBOUND_IN_BYTES
    ):
        return cls(
            cls.resolve(
                target=target, max_inbound_message_size=max_inbound_message_size
            )
        )

    @classmethod
    def from_agent(
        cls,
        connector_id: int,
        agent_id: str,
        max_inbound_message_size: int = DEFAULT_MAX_INBOUND_IN_BYTES,
    ):
        channel = cls.resolve(
            agent_id=agent_id, max_inbound_message_size=max_inbound_message_size
        )
        headers = cls.headers(connector_id, agent_id)
        interceptor = header_adder_interceptor(headers)
        channel = intercept_channel(channel, interceptor)
        return cls(channel)

    def new_session(
        self, configuration: ProtobufConfiguration, auth: Optional[Auth] = None
    ):
        from queryservice_client.querysession import QuerySession

        request = StartSessionRequest(configuration=configuration, authentication=auth)
        response = self.inner.StartSession(request)
        return QuerySession(response.sessionId, self)

    def new_session_from_b64(
        self, b64_encoded_configuration: str, b64_encoded_auth: Optional[str] = None
    ):
        configuration = self.configuration_from_b64_string(b64_encoded_configuration)
        auth = (
            self.auth_configuration_from_b64_string(b64_encoded_auth)
            if b64_encoded_auth
            else None
        )
        return self.new_session(configuration, auth)

    def reuse_session(self, session_id: int):
        from queryservice_client.querysession import QuerySession

        return QuerySession(session_id, self)

    def query(
        self,
        session_id: int,
        query: str,
        limit: Optional[int] = None,
        timeout_in_seconds: Optional[int] = None,
    ) -> None:
        request = QueryRequest(
            sessionId=session_id, query=query, limit=limit, timeout=timeout_in_seconds
        )
        return Statement(
            Stream(self.inner.Query(request)),
            lambda: self.cancel_query(session_id),
        )

    def cancel_query(self, session_id: int) -> None:
        request = CancelQueryRequest(sessionId=session_id)
        self.inner.CancelQuery(request)

    def end_session(self, session_id: int) -> None:
        request = EndSessionRequest(sessionId=session_id)
        self.inner.EndSession(request)

    @staticmethod
    def resolve(
        agent_id: Optional[str] = None,
        target: Optional[str] = None,
        max_inbound_message_size: int = DEFAULT_MAX_INBOUND_IN_BYTES,
    ) -> Channel:
        if agent_id:
            return QueryClient.resolve_agent(agent_id)
        return insecure_channel(
            target,
            options=[("grpc.max_receive_message_length", max_inbound_message_size)],
        )

    @staticmethod
    def resolve_agent(
        agent_id: str, max_inbound_message_size: int = DEFAULT_MAX_INBOUND_IN_BYTES
    ) -> Channel:
        channel = insecure_channel("servicediscovery.ocf-system:80")
        try:
            target = (
                ProxyServiceDiscoveryStub(channel)
                .ResolveProxiesFor(Agent(uuid=agent_id))
                .proxies[0]
            )
            return insecure_channel(
                target,
                options=[("grpc.max_receive_message_length", max_inbound_message_size)],
            )
        finally:
            channel.close()

    @staticmethod
    def headers(connector_id: int, agent_id: str) -> dict:
        return {"x-alation-agent": native(str(agent_id)), "x-alation-connector": native(str(connector_id))}

    @staticmethod
    def configuration_from_b64_string(b64_string: str) -> ProtobufConfiguration:
        try:
            return ProtobufConfiguration.FromString(base64.b64decode(b64_string))
        except DecodeError as e:
            raise ValueError("Invalid base64 configuration string") from e

    @staticmethod
    def auth_configuration_from_b64_string(b64_string: str) -> Auth:
        try:
            return Auth.FromString(base64.b64decode(b64_string))
        except DecodeError as e:
            raise ValueError("Invalid base64 auth string") from e

    def close(self) -> None:
        self.channel.close()

    @staticmethod
    def resolve_client(connector_metadata: dict[str, Union[int, str]]) -> "QueryClient":
        connector_id: int = connector_metadata.get("connector_id")
        agent_id: str = connector_metadata.get("agent_id")
        use_dev_deployment_address = os.environ.get("USE_DEV_DEPLOYMENT_ADDRESS", "false").lower() == "true"
        if use_dev_deployment_address:
            deployment_address = None
        else:
            deployment_address: str = connector_metadata.get("deployment_address")
        if agent_id != None:
            return QueryClient.from_agent(connector_id, agent_id)
        elif deployment_address != None:
            return QueryClient.from_target(deployment_address)
        return QueryClient.from_connector_id(connector_id)


def run_query(connector_metadata: dict[str, Union[int, str]], query: str, logger: any) -> (list, list):
    rows = []
    columns = None
    try:
        logger.debug("Connector metadata: ")
        for key, value in connector_metadata.items():
            if key != "protobuf_config":
                logger.debug(f"{key}: {value}")

        logger.info(f"Running query: \n {query}")
        client = QueryClient.resolve_client(connector_metadata)
        protobuf_config_str: str = connector_metadata.get("protobuf_config")
        session = client.new_session_from_b64(protobuf_config_str)
        statement = session.query(query, timeout_in_seconds=600)  # 10 minutes timeout
        logger.info("Query execution completed.")       
        for result in statement:
            try:
                columns = result.get_metadata().get_columns_name()
                while True:
                    rowstream = next(result)
                    row = []
                    for w in rowstream:
                        row.append(w.get())
                    if row:
                        rows.append(row)
            except StopIteration:
                logger.info("Completed reading results.")
        session.close()
    except Exception as e:
        logger.error(f"Query execution failed with exception: {e}")
        logger.error(f"Failed query: \n {query}")
        raise
    finally:
        client.close()
    return columns, rows
