import json
import os
import threading
import time
import urllib.error
import urllib.request

DEFAULT_BASE_URL = "http://localhost:8080"
DEFAULT_API_VERSION = "v1"
DEFAULT_SCOPE_ORDER = [
    "SCOPE_AGENT",
    "SCOPE_TEAM",
    "SCOPE_PROJECT",
    "SCOPE_GLOBAL",
]


class APIError(Exception):
    def __init__(self, status_code, code, message, details, request_id):
        super().__init__(message or "STELE API error")
        self.status_code = status_code
        self.code = code
        self.message = message
        self.details = details
        self.request_id = request_id

    def __str__(self):
        if self.code:
            return f"{self.code}: {self.message}"
        return self.message or f"STELE API error (status {self.status_code})"


def _env_or_default(key, fallback):
    value = os.environ.get(key, "").strip()
    return value or fallback


class SteleClient:
    def __init__(
        self,
        base_url=None,
        api_version=None,
        agent_id="",
        session_id="",
        client_id="",
        embedder_override="",
        ops_auth_token="",
        timeout=15,
    ):
        if base_url is None or base_url == "":
            base_url = _env_or_default("STELE_BASE_URL", DEFAULT_BASE_URL)
        if api_version is None or api_version == "":
            api_version = _env_or_default("STELE_API_VERSION", DEFAULT_API_VERSION)
        if not agent_id:
            agent_id = os.environ.get("STELE_AGENT_ID", "").strip()
        if not session_id:
            session_id = os.environ.get("STELE_SESSION_ID", "").strip()
        if not client_id:
            client_id = os.environ.get("STELE_CLIENT_ID", "").strip()
        if not embedder_override:
            embedder_override = os.environ.get("STELE_EMBEDDER_OVERRIDE", "").strip()
        if not ops_auth_token:
            ops_auth_token = (
                os.environ.get("STELE_OPS_AUTH_TOKEN", "").strip()
                or os.environ.get("STELE_MCP_AUTH_TOKEN", "").strip()
            )
        self.base_url = (base_url or DEFAULT_BASE_URL).rstrip("/")
        self.api_version = api_version or DEFAULT_API_VERSION
        self.agent_id = agent_id
        self.session_id = session_id
        self.client_id = client_id
        self.embedder_override = embedder_override
        self.ops_auth_token = ops_auth_token
        self.timeout = timeout
        self._counter = 0
        self._lock = threading.Lock()

    def write(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.write", payload)

    def query(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        if not payload.get("scope_order"):
            payload["scope_order"] = list(DEFAULT_SCOPE_ORDER)
        return self._request("POST", "/memory.query", payload)

    def feedback(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.feedback", payload)

    def protect(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.protect", payload)

    def supersede(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.supersede", payload)

    def forget(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.forget", payload)

    def list_memories(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.list", payload)

    def get_memory(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.get", payload)

    def search_memories(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.search", payload)

    def update_memory(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.update", payload)

    def verify(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.verify", payload)

    def extract(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.extract", payload)

    def learn_share(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.learn.share", payload)

    def edge_add(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.edge.add", payload)

    def edge_delete(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.edge.delete", payload)

    def edge_list(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.edge.list", payload)

    def edge_traverse(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.edge.traverse", payload)

    def edge_suggest(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.edge.suggest", payload)

    def bulk_write(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.bulk.write", payload)

    def bulk_forget(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.bulk.forget", payload, is_ops=True)

    def export_ndjson(self, request=None):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request_raw("POST", "/memory.export", payload, is_ops=True)

    def export_to_file(self, output_path, request=None):
        data = self.export_ndjson(request)
        with open(output_path, "w", encoding="utf-8") as f:
            f.write(data)
        return {"ok": True, "output_path": output_path, "bytes": len(data.encode("utf-8"))}

    def import_memories(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.import", payload, is_ops=True)

    def schema(self, request=None):
        if request is None:
            return self._request("GET", "/memory.schema", None)
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("GET", "/memory.schema", payload)

    def metrics(self, request=None):
        if request is None:
            return self._request("GET", "/memory.metrics", None)
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("GET", "/memory.metrics", payload)

    def maintenance(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.maintenance", payload, is_ops=True)

    def reindex(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.reindex", payload, is_ops=True)

    def agent_get(self, request=None):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        if not payload.get("agent_id"):
            payload["agent_id"] = payload["meta"].get("agent_id") or self.agent_id
        return self._request("POST", "/memory.agent.get", payload)

    def agent_upsert(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        profile = dict(payload.get("profile") or {})
        if not profile.get("agent_id"):
            profile["agent_id"] = payload["meta"].get("agent_id") or self.agent_id
        payload["profile"] = profile
        return self._request("POST", "/memory.agent.upsert", payload)

    def access_mark_used(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.access.mark_used", payload)

    def decay(self, request=None):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.decay", payload, is_ops=True)

    def events(self, request=None):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        return self._request("POST", "/memory.events", payload)

    def stats(self, request=None):
        if request:
            payload = dict(request)
            payload["meta"] = self.build_meta(payload.get("meta"))
            return self._request("GET", "/memory.stats", payload)
        return self._request("GET", "/memory.stats", None)

    def membership_add(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        if not payload.get("agent_id"):
            payload["agent_id"] = payload["meta"].get("agent_id") or self.agent_id
        if payload.get("agent_id") and not payload["meta"].get("agent_id"):
            payload["meta"]["agent_id"] = payload["agent_id"]
        return self._request("POST", "/memory.membership.add", payload)

    def membership_remove(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        if not payload.get("agent_id"):
            payload["agent_id"] = payload["meta"].get("agent_id") or self.agent_id
        if payload.get("agent_id") and not payload["meta"].get("agent_id"):
            payload["meta"]["agent_id"] = payload["agent_id"]
        return self._request("POST", "/memory.membership.remove", payload)

    def membership_list(self, request):
        payload = dict(request or {})
        payload["meta"] = self.build_meta(payload.get("meta"))
        if payload.get("agent_id") and not payload["meta"].get("agent_id"):
            payload["meta"]["agent_id"] = payload["agent_id"]
        return self._request("GET", "/memory.membership.list", payload)

    def health(self, request=None):
        if request:
            payload = dict(request)
            payload["meta"] = self.build_meta(payload.get("meta"))
            return self._request("GET", "/memory.health", payload)
        return self._request("GET", "/memory.health", None)

    def build_meta(self, meta=None):
        data = dict(meta or {})
        if not data.get("api_version"):
            data["api_version"] = self.api_version
        if not data.get("request_id"):
            data["request_id"] = self._next_request_id()
        if not data.get("agent_id"):
            data["agent_id"] = self.agent_id
        if not data.get("session_id"):
            data["session_id"] = self.session_id
        if not data.get("client_id") and self.client_id:
            data["client_id"] = self.client_id
        if not data.get("embedder_override") and self.embedder_override:
            data["embedder_override"] = self.embedder_override
        if not data.get("timestamp_ms"):
            data["timestamp_ms"] = int(time.time() * 1000)
        return data

    def _next_request_id(self):
        with self._lock:
            self._counter += 1
            counter = self._counter
        return f"req-{int(time.time() * 1000)}-{counter}"

    def _request(self, method, path, payload, is_ops=False):
        url = f"{self.base_url}{path}"
        headers = {"Accept": "application/json"}
        if is_ops and self.ops_auth_token:
            headers["Authorization"] = f"Bearer {self.ops_auth_token}"
        data = None
        if payload is not None:
            data = json.dumps(payload).encode("utf-8")
            headers["Content-Type"] = "application/json"
        req = urllib.request.Request(url, data=data, headers=headers, method=method)
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as response:
                raw = response.read()
                if not raw:
                    return None
                return json.loads(raw.decode("utf-8"))
        except urllib.error.HTTPError as exc:
            raw = exc.read()
            raise self._api_error(exc.code, raw) from None

    def _request_raw(self, method, path, payload, is_ops=False):
        url = f"{self.base_url}{path}"
        headers = {"Accept": "*/*"}
        if is_ops and self.ops_auth_token:
            headers["Authorization"] = f"Bearer {self.ops_auth_token}"
        data = None
        if payload is not None:
            data = json.dumps(payload).encode("utf-8")
            headers["Content-Type"] = "application/json"
        req = urllib.request.Request(url, data=data, headers=headers, method=method)
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as response:
                raw = response.read()
                return raw.decode("utf-8") if raw else ""
        except urllib.error.HTTPError as exc:
            raw = exc.read()
            raise self._api_error(exc.code, raw) from None

    def _api_error(self, status_code, raw):
        message = f"STELE API error (status {status_code})"
        code = ""
        details = ""
        request_id = ""
        if raw:
            try:
                payload = json.loads(raw.decode("utf-8"))
                meta = payload.get("meta") or {}
                err = meta.get("error") or {}
                code = err.get("code") or ""
                message = err.get("message") or message
                details = err.get("details") or ""
                request_id = meta.get("request_id") or ""
                if not code and isinstance(payload, dict):
                    # Some endpoints return plain errors: {"code":"...","message":"...","details":"..."}.
                    code = payload.get("code") or code
                    message = payload.get("message") or message
                    details = payload.get("details") or details
                    request_id = payload.get("request_id") or request_id
            except (ValueError, TypeError, UnicodeDecodeError):
                pass
        return APIError(status_code, code, message, details, request_id)
