from typing import Iterator, Optional, Literal, Callable, Tuple

import time
import contextlib
import grpc
import asyncio
import gymnasium as gym
import numpy as np

from gymnasium.wrappers import HumanRendering
from rich.align import Align
from rich.live import Live

from rlmesh.client.console import RLMeshConsole, RLMeshStatus
from rlmesh.client.print import (
    print_envs,
    generate_live_env_panel,
    final_results_panel,
)
from rlmesh.env import RemoteEnv, RemoteVectorEnv
from rlmesh.ping import ping_gcp_regions_async
from rlmesh.proto import cp_public_pb2, cp_public_pb2_grpc, session_pb2
from rlmesh.utils import config

DEFAULT_CONTROLPLANE_ENDPOINT = "34.130.104.52:4510"


class RLMeshClient:
    def __init__(
        self,
        api_key: str | None = None,
        *,
        endpoint: str | None = DEFAULT_CONTROLPLANE_ENDPOINT,
        insecure: bool = True,
    ):
        self._addr = endpoint or config.api_url
        self._token = api_key
        self._locations = None

        self._chan = grpc.insecure_channel(self._addr) if insecure else None
        if self._chan is None:
            raise NotImplementedError("TLS not wired yet")

        self._stub = cp_public_pb2_grpc.PublicControlPlaneAPIStub(self._chan)
        self._metadata = []

        if self._token:
            self._metadata.append(("x-rlmesh-token", self._token))

        self._env = None

        self._console = RLMeshConsole()
        self._console.display_title()

    @contextlib.contextmanager
    def channel(self):
        try:
            yield self._chan
        finally:
            if self._chan is not None:
                self._chan.close()

    def ping_regions(self, status: RLMeshStatus | None = None):
        if status:
            status.update("Pinging regions...")

        regions = list(self.list_regions().regions)
        if not regions:
            return None

        results = asyncio.run(ping_gcp_regions_async(regions))
        sorted_results = sorted(results.items(), key=lambda x: x[1])

        table = self._console.table()
        table.add_column("Region", justify="left", style="cyan")
        table.add_column("Latency (ms)", justify="right", style="green")

        for region, latency in sorted_results:
            table.add_row(region, f"{latency:.1f}")

        self._console.print(
            self._console.panel(
                table, title="[bold white]Region Latency Results[/bold white]"
            )
        )

        return results

    def list_regions(self, deadline_s: float = 5.0) -> cp_public_pb2.ListRegionsRes:
        req = cp_public_pb2.ListRegionsReq()
        return self._stub.ListRegions(req, timeout=deadline_s, metadata=self._metadata)

    def list_environments(
        self, deadline_s: float = 5.0
    ) -> cp_public_pb2.ListEnvironmentsRes:
        req = cp_public_pb2.ListEnvironmentsReq()
        return self._stub.ListEnvironments(
            req, timeout=deadline_s, metadata=self._metadata
        )

    def create_interactive_session(
        self, gym_id: str, *, num_envs: int = 1, deadline_s: float = 5.0
    ) -> session_pb2.Session:
        if self._locations is None:
            self._locations = self.ping_regions()

        req = cp_public_pb2.CreateSessionReq(
            interactive=session_pb2.InteractiveSessionSpec(
                gym_id=gym_id, num_envs=num_envs, locations=self._locations
            )
        )
        try:
            res = self._stub.CreateSession(
                req, timeout=deadline_s, metadata=self._metadata
            )
            return res
        except grpc.RpcError as e:
            if e.code() == grpc.StatusCode.NOT_FOUND:
                try:
                    details = e.details() if hasattr(e, "details") else str(e)
                except Exception:
                    details = str(e)
                print(f"[rlmesh] CreateSession failed: {details}")
                print("[rlmesh] Available environments:")
                self.print_envs()
            raise

    def create_external_benchmark_session(
        self,
        comp_id: str,
        user_token: str,
        status: RLMeshStatus,
        deadline_s: float = 5.0,
    ) -> session_pb2.Session:
        if self._locations is None:
            self._locations = self.ping_regions(status)

        status.update("Creating session...")
        req = cp_public_pb2.CreateSessionReq(
            external_benchmark=session_pb2.ExternalBenchmarkSessionSpec(
                comp_id=comp_id, user_token=user_token, locations=self._locations
            )
        )
        try:
            res = self._stub.CreateSession(
                req, timeout=deadline_s, metadata=self._metadata
            )
            return res
        except grpc.RpcError as e:
            if e.code() == grpc.StatusCode.NOT_FOUND:
                try:
                    details = e.details() if hasattr(e, "details") else str(e)
                except Exception:
                    details = str(e)
                print(f"[rlmesh] CreateSession failed: {details}")
                print("[rlmesh] Available environments:")
                self.print_envs()
            raise

    def watch_session(
        self, session_id: str, *, deadline_s: float = 120.0
    ) -> Iterator[session_pb2.Session]:
        req = cp_public_pb2.WatchReq(session_id=session_id)
        stream = self._stub.WatchSession(
            req, timeout=deadline_s, metadata=self._metadata
        )
        for evt in stream:
            yield evt

    def wait_until_ready(
        self,
        session_id: str,
        status: RLMeshStatus | None = None,
        deadline_s: float = 600.0,
    ) -> session_pb2.Session:
        if status:
            status.update("Waiting for environment to be ready...")

        ready_evt: Optional[session_pb2.Session] = None
        for evt in self.watch_session(session_id, deadline_s=deadline_s):
            if evt.state == session_pb2.SessionState.READY:
                if status:
                    status.update("Environment ready!")
                ready_evt = evt
                break

            if evt.state == session_pb2.SessionState.QUEUED:
                if status:
                    status.update(
                        f"Waiting in queue ({f'{evt.queue_position} ahead' if evt.queue_position > 0 else 'next up'})..."
                    )

            if evt.state == session_pb2.SessionState.PROVISIONING:
                if status:
                    status.update("Waiting for environment to be ready...")

        if not ready_evt:
            raise RuntimeError("session stream ended without READY")
        return ready_evt

    def describe(
        self, session_id: str, *, deadline_s: float = 5.0
    ) -> cp_public_pb2.DescribeRes:
        req = cp_public_pb2.DescribeReq(session_id=session_id)
        res: cp_public_pb2.DescribeRes = self._stub.DescribeSession(
            req, timeout=deadline_s, metadata=self._metadata
        )

        details = self._console.table(box=None, show_header=False, expand=False)
        details.add_row("[cyan]ID:[/cyan]", session_id)
        details.add_row("[cyan]Gym Id:[/cyan]", res.env.gym_id)
        details.add_row("[cyan]Package:[/cyan]", res.env.package)
        details.add_row("[cyan]Num Envs:[/cyan]", str(res.num_envs))

        panel = self._console.panel(
            Align.left(details),
            title="[bold white]Session Details[/bold white]",
        )

        self._console.print(panel)
        return res

    def get_results(
        self, session_id: str, *, deadline_s: float = 20.0
    ) -> cp_public_pb2.GetResultsRes:
        req = cp_public_pb2.GetResultsReq(session_id=session_id)
        return self._stub.GetResults(req, timeout=deadline_s, metadata=self._metadata)

    def _connect_env(
        self,
        *,
        session: session_pb2.Session,
        desc: cp_public_pb2.DescribeRes,
        user_token: str | None = None,
    ) -> RemoteEnv:
        node_addr = getattr(session, "address")
        if not node_addr:
            raise ValueError("session is missing node address")

        metadata = self._metadata.copy()
        metadata.append(("x-rlmesh-session", session.id))

        if user_token is not None:
            metadata.append(("x-rlmesh-usertoken", user_token))

        return RemoteEnv(
            desc.env,
            transport_addr=node_addr,
            transport_metadata=metadata,
        )

    def _connect_vec_env(
        self,
        *,
        session: session_pb2.Session,
        desc: cp_public_pb2.DescribeRes,
        user_token: str | None = None,
    ) -> RemoteVectorEnv:
        node_addr = getattr(session, "address")
        if not node_addr:
            raise ValueError("session is missing node address")

        metadata = self._metadata.copy()
        metadata.append(("x-rlmesh-session", session.id))

        if user_token is not None:
            metadata.append(("x-rlmesh-usertoken", user_token))

        return RemoteVectorEnv(
            desc.env,
            num_envs=desc.num_envs,
            transport_addr=node_addr,
            transport_metadata=metadata,
        )

    def _create_and_connect(self, gym_id: str) -> Tuple[RemoteEnv, session_pb2.Session]:
        if self._env:
            try:
                self._env.close()
            finally:
                self._env = None

        sess = self.create_interactive_session(gym_id, num_envs=1)
        sess = self.wait_until_ready(sess.id)
        desc = self.describe(sess.id)

        env = self._connect_env(session=sess, desc=desc)
        env.render_mode = "rgb_array"
        self._env = env
        return env, sess

    def _create_and_connect_vec(
        self, gym_id: str, num_envs: int = 2
    ) -> Tuple[RemoteVectorEnv, session_pb2.Session]:
        if self._env:
            try:
                self._env.close()
            finally:
                self._env = None

        sess = self.create_interactive_session(gym_id, num_envs=num_envs)
        desc = self.describe(sess.id)
        sess = self.wait_until_ready(sess.id)

        env = self._connect_vec_env(session=sess, desc=desc)
        env.render_mode = "rgb_array"
        self._env = env
        return env, sess

    def _create_and_connect_benchmark(
        self, comp_id: str, user_token: str, status: RLMeshStatus
    ) -> Tuple[RemoteVectorEnv, session_pb2.Session, cp_public_pb2.DescribeRes]:
        if self._env:
            try:
                self._env.close()
            finally:
                self._env = None

        sess = self.create_external_benchmark_session(
            comp_id=comp_id, user_token=user_token, status=status
        )
        desc = self.describe(sess.id)
        sess = self.wait_until_ready(sess.id, status)

        env = self._connect_vec_env(session=sess, desc=desc, user_token=user_token)
        env.render_mode = "rgb_array"
        self._env = env
        return env, sess, desc

    def make(
        self, gym_id: str, render_mode: Literal["human", "rgb_array"] = "rgb_array"
    ) -> gym.Env:
        env, _ = self._create_and_connect(gym_id)

        if render_mode == "human":
            env = HumanRendering(env)

        self._env = env
        return env

    def evaluate(
        self, comp_id: str, user_token: str, get_action: Callable | None = None
    ):
        with self._console.status(
            f"Creating evaluation session for `{comp_id}`..."
        ) as status:
            envs, sess, desc = self._create_and_connect_benchmark(
                comp_id, user_token, status
            )

            if status:
                status.update("Starting benchmark...")
                status.stop()

            perf = {}
            steps = 0
            t0 = time.time()

            obs, _ = envs.reset()

            terminated = np.zeros(desc.num_envs, bool)
            truncated = np.zeros(desc.num_envs, bool)
            rewards = np.zeros(desc.num_envs, dtype=np.float64)

            with Live(refresh_per_second=4, console=self._console.console) as live:
                while not np.all(terminated | truncated):
                    steps += 1

                    actions = (
                        envs.action_space.sample()
                        if get_action is None
                        else get_action(obs)
                    )
                    obs, r, terms, truns, infos = envs.step(actions)

                    elapsed = time.time() - t0
                    fps = steps / elapsed if elapsed > 0 else 0

                    alive = ~(terminated | truncated)
                    rewards[alive] += r[alive]

                    terminated |= terms
                    truncated |= truns

                    if "perf" in infos:
                        perf = infos["perf"]

                    panel = generate_live_env_panel(
                        self._console,
                        steps,
                        num_envs=desc.num_envs,
                        active_envs=np.sum(~(terminated | truncated)),
                        score=np.mean(rewards),
                        fps=fps,
                        elapsed=elapsed,
                        perf=perf,
                    )
                    live.update(panel)

            envs.close()
            results = self.get_results(sess.id)
            results = {
                "score": results.score,
                "duration": results.duration,
                "timestamp": results.timestamp.ToJsonString(),
            }

            self._console.print(final_results_panel(self._console, results))
            return results

    def print_envs(self):
        res = self.list_environments()
        return print_envs(res)
