import json
import threading
import requests
from typing import List, Optional, Dict, Any, Callable

import grpc
from google.protobuf import json_format
from google.protobuf.struct_pb2 import Struct

from .proto import celestialtree_pb2 as pb2
from .proto import celestialtree_pb2_grpc as pb2_grpc


class Client:
    """
    Python client for CelestialTree HTTP API (and optional gRPC).
    """

    def __init__(
        self,
        host: str = "127.0.0.1",
        http_port: int = 7777,
        grpc_port: int = 7778,
        timeout: float = 5.0,
        grpc_secure: bool = False,
        transport: Optional[str] = None,
    ):
        self.http_addr = f"http://{host}:{http_port}"
        self.grpc_addr = f"{host}:{grpc_port}"

        self.timeout = timeout
        self.grpc_secure = grpc_secure

        self.transport = transport if transport in ("http", "grpc") else "http"

    def init_session(self):
        if hasattr(self, "session"):
            return

        self.session = requests.Session()
        self.session.headers.update(
            {
                "Content-Type": "application/json",
                "Accept": "application/json",
            }
        )

    def init_grpc(self):
        if hasattr(self, "grpc_channel") and hasattr(self, "grpc_stub"):
            return

        if self.grpc_secure:
            creds = grpc.ssl_channel_credentials()
            self.grpc_channel = grpc.secure_channel(self.grpc_addr, creds)
        else:
            self.grpc_channel = grpc.insecure_channel(self.grpc_addr)

        self.grpc_stub = pb2_grpc.CelestialTreeServiceStub(self.grpc_channel)

    def raise_for_status(self, r: requests.Response):
        if 200 <= r.status_code < 300:
            return

        try:
            j: dict = r.json()
            error = j.get("error", "request failed")
            detail = j.get("detail")
            raise RuntimeError(f"{error} ({detail})" if detail else error)
        except Exception:
            raise RuntimeError(f"request failed: HTTP {r.status_code}: {r.text[:300]}")

    # ---------- Core APIs ----------

    def emit(
        self,
        type_: str,
        parents: Optional[List[int]] = None,
        message: Optional[str] = None,
        payload: Optional[list | dict] = None,
    ) -> int:
        """
        Emit a new event into CelestialTree.
        """
        if self.transport == "grpc":
            return self.emit_grpc(type_, parents, message, payload)
        return self.emit_http(type_, parents, message, payload)

    def emit_http(
        self,
        type_: str,
        parents: Optional[List[int]] = None,
        message: Optional[str] = None,
        payload: Optional[list | dict] = None,
    ) -> int:
        """
        Emit a new event into CelestialTree.
        """
        self.init_session()

        body = {
            "type": type_,
            "parents": parents or [],
        }

        if message is not None:
            body["message"] = message

        if payload is not None:
            if isinstance(payload, (dict, list)):
                body["payload"] = payload
            else:
                raise TypeError("payload must be JSON-serializable")

        r = self.session.post(
            f"{self.http_addr}/emit",
            json=body,
            timeout=self.timeout,
        )

        self.raise_for_status(r)
        return r.json()["id"]

    def emit_grpc(
        self,
        type_: str,
        parents: Optional[List[int]] = None,
        message: Optional[str] = None,
        payload: Optional[list | dict] = None,
    ) -> int:
        """
        Emit a new event into CelestialTree via gRPC.
        Requires grpcio + generated pb2 / pb2_grpc.
        """
        self.init_grpc()

        req = pb2.EmitRequest(
            type=type_,
            message=message or "",
            parents=parents or [],
        )

        if payload is not None:
            if not isinstance(payload, (dict, list)):
                raise TypeError("payload must be JSON-serializable (dict or list)")

            # proto field is google.protobuf.Struct (object). Struct 本身不支持 list 作为根。
            # 所以：dict 直接塞；list 根的话用 {"_": payload} 包一层（服务端再原样存 JSON）
            if isinstance(payload, list):
                payload = {"_": payload}

            st = Struct()
            json_format.ParseDict(payload, st)
            req.payload.CopyFrom(st)

        try:
            resp = self.grpc_stub.Emit(req, timeout=self.timeout)
        except grpc.RpcError as e:
            # 给你一个更像 HTTP raise_for_status 的报错体验
            raise RuntimeError(
                f"grpc emit failed: {e.code().name}: {e.details()}"
            ) from e

        return int(resp.id)

    def get_event(self, event_id: int) -> Dict[str, Any]:
        self.init_session()

        r = self.session.get(
            f"{self.http_addr}/event/{event_id}",
            timeout=self.timeout,
        )

        self.raise_for_status(r)
        return r.json()

    def children(self, event_id: int) -> List[int]:
        self.init_session()

        r = self.session.get(
            f"{self.http_addr}/children/{event_id}",
            timeout=self.timeout,
        )

        self.raise_for_status(r)
        return r.json()

    def ancestors(self, event_id: int) -> List[int]:
        self.init_session()

        r = self.session.get(
            f"{self.http_addr}/ancestors/{event_id}",
            timeout=self.timeout,
        )

        self.raise_for_status(r)
        return r.json()

    def descendants(self, event_id: int, view: str = "struct") -> Dict[str, Any]:
        self.init_session()

        params = None
        if view and view != "struct":
            # 默认 struct 不传参，保持最干净也最兼容
            params = {"view": view}

        r = self.session.get(
            f"{self.http_addr}/descendants/{event_id}",
            params=params,
            timeout=self.timeout,
        )

        self.raise_for_status(r)
        return r.json()

    def descendants_batch(
        self, event_ids: List[int], view: str = "struct"
    ) -> List[Dict[str, Any]]:
        """
        Batch descendants.
        POST /descendants
        body: {"ids":[...], "view":"struct|meta"}
        response: [tree, tree, ...]
        """
        self.init_session()

        if not event_ids:
            raise ValueError("event_ids is required")

        body: Dict[str, Any] = {"ids": event_ids}
        if view and view != "struct":
            body["view"] = view

        r = self.session.post(
            f"{self.http_addr}/descendants",
            data=json.dumps(body),
            timeout=self.timeout,
        )

        self.raise_for_status(r)
        return r.json()

    def provenance(self, event_id: int, view: str = "struct") -> Dict[str, Any]:
        self.init_session()

        params = None
        if view and view != "struct":
            params = {"view": view}

        r = self.session.get(
            f"{self.http_addr}/provenance/{event_id}",
            params=params,
            timeout=self.timeout,
        )

        self.raise_for_status(r)
        return r.json()

    def provenance_batch(
        self, event_ids: List[int], view: str = "struct"
    ) -> List[Dict[str, Any]]:
        """
        Batch provenance (parents tree).
        POST /provenance
        body: {"ids":[...], "view":"struct|meta"}
        response: [tree, tree, ...]
        """
        self.init_session()

        if not event_ids:
            raise ValueError("event_ids is required")

        body: Dict[str, Any] = {"ids": event_ids}
        if view and view != "struct":
            body["view"] = view

        r = self.session.post(
            f"{self.http_addr}/provenance",
            data=json.dumps(body),
            timeout=self.timeout,
        )

        self.raise_for_status(r)
        return r.json()

    def heads(self) -> List[int]:
        self.init_session()

        r = self.session.get(
            f"{self.http_addr}/heads",
            timeout=self.timeout,
        )

        self.raise_for_status(r)
        return r.json()

    def health(self) -> bool:
        self.init_session()
        try:
            r = self.session.get(
                f"{self.http_addr}/healthz",
                timeout=self.timeout,
            )
            return r.status_code == 200
        except Exception:
            return False

    def version(self) -> Dict[str, Any]:
        self.init_session()

        r = self.session.get(
            f"{self.http_addr}/version",
            timeout=self.timeout,
        )

        self.raise_for_status(r)
        return r.json()

    # ---------- SSE Subscribe ----------

    def subscribe(
        self,
        on_event: Callable[[Dict[str, Any]], None],
        daemon: bool = True,
    ) -> threading.Thread:
        """
        Subscribe to SSE stream.
        on_event will be called for each emitted Event.
        """

        def _run():
            with self.session.get(
                f"{self.http_addr}/subscribe",
                stream=True,
                timeout=None,
            ) as r:
                r.raise_for_status()
                for line in r.iter_lines(decode_unicode=True):
                    if not line:
                        continue

                    if line.startswith("data:"):
                        data = line[len("data:") :].strip()
                        try:
                            ev = json.loads(data)
                            on_event(ev)
                        except Exception:
                            pass

        self.init_session()

        t = threading.Thread(target=_run, daemon=daemon)
        t.start()
        return t
