import msgpack
import redis
from aiohttp import web
from params_proto import Proto, ParamsProto, Flag

from zaku.base import Server
from zaku.interfaces import Job


class Redis(ParamsProto, prefix="redis", cli_parse=False):
    """Redis Configuration for the TaskServer class.

    We support both direct connection to a single redis server,
        Redis Cluster mode, and a high availability setup using Redis Sentinel.

    Single Redis Server Usage
    =========================

    .. code-block:: shell

        # Put this into an .env file
        REDIS_HOST=localhost
        REDIS_PORT=6379
        REDIS_PASSWORD=xxxxxxxxxxxxxxxxxxx
        REDIS_DB=0

    CLI Options::

        --redis.host      :str 'localhost'
        --redis.port      :int 6379
        --redis.password  :any None
        --redis.db        :any 0

    Redis Cluster Usage
    ===================

    .. code-block:: shell

        # Put this into an .env file
        REDIS_CLUSTER_HOSTS=host1:port1,host2:port2,host3:port3
        REDIS_PASSWORD=xxxxxxxxxxxxxxxxxxx

    CLI Options::

        --redis.cluster_hosts     :str ''
        --redis.password          :any None

    Sentinel Usage (High availability setup)
    ========================================

    .. code-block:: shell

        # Put this into an .env file
        SENTINEL_HOSTs=host1:port1,host2:port2  # <=== Notice it is plural.
                                                #   This is a comma separated list of hosts.
        SENTINEL_PASSWORD=xxxxxxxxxxxxxxxxxxx.  # <=== This is the password for the sentinels.

        sentinel_cluster_name=primary           # <=== This is the name of the redis cluster.
        SENTINEL_DB=0                           # <=== This is the logical database.

    CLI Options::

        --redis.sentinel_hosts     :str 'localhost'
        --redis.sentinel_password  :any None
        --redis.sentinel_db        :any 0

    Common Errors:
    ==============

    Setting keepalive: ConnectionError is due to timeout. See below for more infomation.

    https://devcenter.heroku.com/articles/ah-redis-stackhero#:~:text=The%20error%20%E2%80%9Credis.,and%20the%20connection%20closes%20automatically
    """

    host: str = Proto("localhost", env="REDIS_HOST")
    port: int = Proto(6379, env="REDIS_PORT")
    password: str = Proto(env="REDIS_PASSWORD")

    # Redis Cluster support
    cluster_hosts: str = Proto(
        env="REDIS_CLUSTER_HOSTS", help="comma separated list of redis cluster nodes. Example: host1:port1,host2:port2,host3:port3."
    )

    # Sentinel configuration
    sentinel_hosts: str = Proto(
        env="SENTINEL_HOSTS", help="comma separated list of redis hosts. Example: host1:port1,host2:port2,host3:port3."
    )
    sentinel_password: str = Proto(password, env="SENTINEL_PASSWORD")

    cluster_name = Proto("primary", env="SENTINEL_CLUSTER_NAME")
    shuffle: bool = Proto(False, env="REDIS_SHUFFLE", help="shuffle the sentinel hosts on init.")

    db: int = Proto(0, env="REDIS_DB", help="The logical redis database, from 0 - 15.")

    # connection arguments
    health_check_interval = 10
    socket_connect_timeout = 5
    retry_on_timeout = True
    socket_keepalive = True

    def __post_init__(self, _deps=None):
        if self.cluster_hosts:
            # Redis Cluster mode
            from redis.asyncio.cluster import RedisCluster

            startup_nodes = []
            for host_port in self.cluster_hosts.split(","):
                host, port = host_port.split(":")
                startup_nodes.append({"host": host.strip(), "port": int(port.strip())})
            
            self.connection = RedisCluster(
                startup_nodes=startup_nodes,
                password=self.password,
                skip_full_coverage_check=True,
                health_check_interval=self.health_check_interval,
                socket_connect_timeout=self.socket_connect_timeout,
                retry_on_timeout=self.retry_on_timeout,
                socket_keepalive=self.socket_keepalive,
            )

        elif self.sentinel_hosts:
            import random

            self.sentinel_hosts = [h.split(":") for h in self.sentinel_hosts.split(",")]

            if self.shuffle:
                self.sentinel_hosts = random.shuffle(self.sentinel_hosts, key=lambda x: hash(x))

            self.sentinel = redis.asyncio.sentinel.Sentinel(
                self.sentinel_hosts,
                # redis_class=RobustRedis,
                password=self.sentinel_password,
                health_check_interval=self.health_check_interval,
                socket_connect_timeout=self.socket_connect_timeout,
                retry_on_timeout=self.retry_on_timeout,
                socket_keepalive=self.socket_keepalive,
            )

            self.connection = self.sentinel.master_for(self.cluster_name, password=self.password, db=self.db)

        else:
            # self.connection = RobustRedis(password=self.password, db=self.db, host=self.host, port=self.port)
            self.connection = redis.asyncio.Redis(
                password=self.password,
                db=self.db,
                host=self.host,
                port=self.port,
                health_check_interval=self.health_check_interval,
                socket_connect_timeout=self.socket_connect_timeout,
                retry_on_timeout=self.retry_on_timeout,
                socket_keepalive=self.socket_keepalive,
            )


class TaskServer(ParamsProto, Server):
    """TaskServer

    This is the server that maintains the Task Queue.

    Usage

    .. code-block:: python

        app = TaskServer()
        app.run()

    CLI Options

    .. code-block:: shell

        python -m zaku.server --help

        -h, --help            show this help message and exit
        --prefix            :str 'Zaku-task-queues'
        --queue-len         :int 100
        --host              :str 'localhost'
                              set to 0.0.0.0 to enable remote (not localhost) connections.
        --port              :int 9000
        --cors              :str 'https://vuer.ai,https://dash.ml,http://lo...
        --cert              :str None the path to the SSL certificate
        --key               :str None the path to the SSL key
        --ca-cert           :str None the trusted root CA certificates
        --REQUEST-MAX-SIZE  :int 100000000 the maximum packet size
        --free-port         :bool False
                              kill process squatting target port if True.
        --static-root       :str '.'
        --verbose           :bool False
                              show the list of configurations during launch if True.

    """

    prefix = "Zaku-task-queues"

    # Queue Parameters
    queue_len = 100  # use a max length to avoid the memory from blowing up.

    # Server Parameters
    host: str = Proto(
        "localhost",
        help="set to 0.0.0.0 to enable remote (not localhost) connections.",
    )
    port: int = 9000
    cors: str = "https://vuer.ai,https://dash.ml,http://localhost:8000,http://127.0.0.1:8000,*"

    # SSL Parameters
    cert: str = Proto(None, dtype=str, help="the path to the SSL certificate")
    key: str = Proto(None, dtype=str, help="the path to the SSL key")
    ca_cert: str = Proto(None, dtype=str, help="the trusted root CA certificates")

    REQUEST_MAX_SIZE: int = Proto(100_000_000, env="WEBSOCKET_MAX_SIZE", help="the maximum packet size")

    free_port: bool = Flag("kill process squatting target port if True.")
    static_root: str = "."
    verbose = True  # Flag("show the list of configurations during launch if True.")

    def print_info(self):
        print("========= Arguments =========")
        for k, v in vars(self).items():
            print(f" {k} = {v},")
        print("-----------------------------")

    def __post_init__(self, _deps=None):
        if self.verbose:
            self.print_info()

        Server.__post_init__(self)

        self.redis_wrapper = Redis(_deps)
        self.redis = self.redis_wrapper.connection

    async def create_queue(self, request: web.Request):
        data = await request.json()
        try:
            await Job.create_queue(self.redis, **data, prefix=self.prefix)
        except Exception as e:
            # the error handling here is not ideal.
            return web.Response(text="ERROR: " + str(e), status=200)

        return web.Response(text="OK")

    async def add_job(self, request: web.Request):
        msg = await request.read()
        data = msgpack.unpackb(msg)
        # print("==>", data)
        await Job.add(self.redis, prefix=self.prefix, **data)
        return web.Response(text="OK")

    async def publish_job(self, request: web.Request):
        msg = await request.read()
        data = msgpack.unpackb(msg)
        # print("==>", data)
        num_subscribers = await Job.publish(self.redis, prefix=self.prefix, **data)
        # todo: return the number of subscribers.
        return web.Response(text=str(num_subscribers), status=200)

    async def reset_handler(self, request: web.Request):
        data = await request.json()
        # print("==>", data)
        await Job.reset(self.redis, **data, prefix=self.prefix)
        return web.Response(text="OK")

    async def remove_handler(self, request: web.Request):
        data = await request.json()
        # print("remove ==>", data)
        await Job.remove(self.redis, **data, prefix=self.prefix)
        return web.Response(text="OK")

    async def count_files_handler(self, request):
        data = await request.json()

        try:
            counts = await Job.count_files(self.redis, **data, prefix=self.prefix)
        except redis.exceptions.ResponseError as e:
            if "no such index" in str(e):
                return web.Response(status=200)
                # return web.Response(text="no such index", status=404)
            raise e

        msg = msgpack.packb({"counts": counts}, use_bin_type=True)
        return web.Response(body=msg, status=200)

    async def take_handler(self, request):
        data = await request.json()

        try:
            job_id, payload = await Job.take(self.redis, **data, prefix=self.prefix)
        except redis.exceptions.ResponseError as e:
            if "no such index" in str(e):
                return web.Response(status=200)
                # return web.Response(text="no such index", status=404)
            raise e

        if payload:
            msg = msgpack.packb({"job_id": job_id, "payload": payload}, use_bin_type=True)
            return web.Response(body=msg, status=200)

        return web.Response(status=200)

    async def unstale_handler(self, request: web.Request):
        data = await request.json()
        # print("take ==> data", data)
        await Job.unstale_tasks(self.redis, **data, prefix=self.prefix)

        return web.Response(text="OK", status=200)

    async def subscribe_one_handler(self, request):
        data = await request.json()

        payload = await Job.subscribe(self.redis, **data, prefix=self.prefix)

        if payload:
            return web.Response(body=payload, status=200)

        return web.Response(status=200)

    async def subscribe_streaming_handler(self, request) -> web.StreamResponse:
        data = await request.json()

        async def stream_response(response):
            try:
                async for payload in Job.subscribe_stream(self.redis, **data, prefix=self.prefix):
                    # use msgpack.Unpacker to determin the end of message.
                    await response.write(payload)

                await response.write_eof()

            except ConnectionResetError:
                print("client disconnected.")
                return response

            return response

        response = web.StreamResponse(
            status=200,
            reason="OK",
            headers={"Content-Type": "text/plain"},
        )
        await response.prepare(request)
        await stream_response(response)
        return response

    def setup_server(self):
        import os

        # use the same endpoint for websocket and file serving.
        self._route("/queues", self.create_queue, method="PUT")
        self._route("/tasks", self.add_job, method="PUT")
        self._route("/tasks", self.take_handler, method="POST")
        self._route("/tasks/counts", self.count_files_handler, method="GET")
        self._route("/tasks/reset", self.reset_handler, method="POST")
        self._route("/tasks", self.remove_handler, method="DELETE")
        self._route("/tasks/unstale", self.unstale_handler, method="PUT")

        self._route("/publish", self.publish_job, method="PUT")
        self._route("/subscribe_one", self.subscribe_one_handler, method="POST")
        self._route("/subscribe_stream", self.subscribe_streaming_handler, method="POST")

        # serve local files via /static endpoint
        self._static("/static", self.static_root)
        print("Serving file://" + os.path.abspath(self.static_root), "at", "/static")

        return self.app

    def run(self, kill=None, *args, **kwargs):
        if kill or self.free_port:
            import time
            from killport import kill_ports

            kill_ports(ports=[self.port])
            time.sleep(0.01)

        self.setup_server()

        print("redis server is now connected")
        print(f"serving at http://localhost:{self.port}")
        super().run()


def entry_point():
    TaskServer().run()


async def get_app():
    """This is sued by gunicorn to run the server in worker mode."""
    return TaskServer().setup_server()


if __name__ == "__main__":
    entry_point()
