from __future__ import annotations

import json
import os
import urllib.error
import urllib.request
from typing import Any, Dict, List, Optional, Mapping
from urllib.parse import urlencode

from pydantic import BaseModel, ConfigDict, Field

from aiel_sdk.errors import AielError


def _normalize_join(base_url: str, path: str) -> str:
    return base_url.rstrip("/") + "/" + path.lstrip("/")


def _try_get_keyring_token(service: str = "aiel", username: str = "default") -> str | None:
    try:
        import keyring  # type: ignore
    except Exception:
        return None
    try:
        return keyring.get_password(service, username)
    except Exception:
        return None


def _resolve_token(explicit: str | None) -> str | None:
    if explicit:
        return explicit
    env_token = os.getenv("AIEL_TOKEN")
    if env_token:
        return env_token
    kr = _try_get_keyring_token()
    if kr:
        return kr
    return None


class MemoryApiError(AielError):
    def __init__(self, status: int, message: str, body: Any | None = None):
        details = f"Memory API error ({status}): {message}"
        if body is not None:
            details += f" | body={body}"
        super().__init__(details)
        self.status = status
        self.body = body


# ----------------------------
# Models
# ----------------------------

class ChatMessage(BaseModel):
    model_config = ConfigDict(extra="allow")
    role: str
    content: str
    ts: Optional[str] = None
    meta: Dict[str, Any] = Field(default_factory=dict)

class ThreadAppendIn(BaseModel):
    model_config = ConfigDict(extra="allow")
    # used when thread_id == "auto"
    thread_key: Optional[str] = None
    messages: List[ChatMessage]
    ttl_seconds: Optional[int] = None
    scope = dict[str, Any]



class ThreadGetOut(BaseModel):
    model_config = ConfigDict(extra="allow")
    thread_id: str
    messages: List[Dict[str, Any]] = Field(default_factory=list)
    updated_at: Optional[str] = None
class ThreadScope(BaseModel):
    user_id: Optional[str] = None
    agent_id: Optional[str] = None

class ThreadStatePatchIn(BaseModel):
    scope: ThreadScope = Field(default_factory=ThreadScope)
    patch: dict = Field(default_factory=dict)
    ttl_seconds: Optional[int] = Field(default=None, ge=60, le=31536000)


class ThreadStateOut(BaseModel):
    thread_id: str
    state: dict
    updated_at: str


class CheckpointPutIn(BaseModel):
    model_config = ConfigDict(extra="allow")
    thread_key: Optional[str] = None
    state: Dict[str, Any] = Field(default_factory=dict)
    metadata: Dict[str, Any] = Field(default_factory=dict)
    ttl_seconds: Optional[int] = None


class RecallPutIn(BaseModel):
    model_config = ConfigDict(extra="allow")
    namespace: List[str]
    key: str
    value: Dict[str, Any] = Field(default_factory=dict)
    tags: Dict[str, Any] = Field(default_factory=dict)
    ttl_seconds: Optional[int] = None


class RecallGetIn(BaseModel):
    model_config = ConfigDict(extra="allow")
    namespace: List[str]
    key: str


class RecallSearchIn(BaseModel):
    model_config = ConfigDict(extra="allow")
    namespace: List[str]
    query: Dict[str, Any] = Field(default_factory=dict)
    limit: int = 20


# ----------------------------
# Client
# ----------------------------
class MemoryClient:
    def __init__(
        self,
        base_url: str,
        api_key: str | None = None,
        timeout: float = 30.0,
        extra_headers: Optional[Dict[str, str]] = None,
        use_x_api_token_fallback: bool = False,
    ) -> None:
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.extra_headers = extra_headers or {}
        self.use_x_api_token_fallback = use_x_api_token_fallback

        if not self.base_url:
            raise ValueError("base_url is required")

        self._token = _resolve_token(api_key)

    def _build_headers(self) -> Dict[str, str]:
        headers = {"Content-Type": "application/json", **self.extra_headers}
        if self._token:
            headers["Authorization"] = f"Bearer {self._token}"
            if self.use_x_api_token_fallback:
                headers["X-API-Token"] = self._token
        return headers

    def _request(
        self,
        method: str,
        path: str,
        *,
        json_body: Any | None = None,
        params: Mapping[str, Any] | None = None,
    ) -> Any:
        # ✅ safe query encoding
        if params:
            qs = urlencode({k: v for k, v in params.items() if v is not None}, doseq=True)
            if qs:
                path = f"{path}{'&' if '?' in path else '?'}{qs}"

        url = _normalize_join(self.base_url, path)
        headers = self._build_headers()

        data = None
        if json_body is not None:
            data = json.dumps(json_body).encode("utf-8")

        req = urllib.request.Request(url, data=data, headers=headers, method=method)
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                raw = resp.read()
                if not raw:
                    return None
                try:
                    return json.loads(raw.decode("utf-8"))
                except json.JSONDecodeError:
                    return raw.decode("utf-8")
        except urllib.error.HTTPError as e:
            body = None
            try:
                raw = e.read()
                if raw:
                    body = json.loads(raw.decode("utf-8"))
            except Exception:
                body = None
            raise MemoryApiError(e.code, e.reason, body) from e

# ---- Thread messages ----

    def thread_append_messages(
        self,
        workspace_id: str,
        project_id: str,
        *,
        thread_id: str,
        payload: ThreadAppendIn,
    ) -> ThreadGetOut:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/thread/{thread_id}/messages:append"
        data = self._request("POST", path, json_body=payload.model_dump(exclude_none=True))
        return ThreadGetOut.model_validate(data)

    def thread_get_messages(
        self,
        workspace_id: str,
        project_id: str,
        *,
        thread_id: str,
        payload: Any,
    ) -> ThreadGetOut:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/thread/{thread_id}/messages?limit=50"
        data = self._request("POST", path, json_body=payload.model_dump(exclude_none=True))
        return ThreadGetOut.model_validate(data)
    
     # ---- Thread state ----
    def thread_get_state(
        self,
        workspace_id: str,
        project_id: str,
        *,
        thread_id: str,
        payload: ThreadStateOut,
    ) -> ThreadGetOut:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/thread/{thread_id}/state"
        data = self._request("POST", path, json_body=payload.model_dump(exclude_none=True))
        return ThreadGetOut.model_validate(data)
    
    def thread_patch_state(
        self,
        workspace_id: str,
        project_id: str,
        *,
        thread_id: str,
        payload: ThreadStatePatchIn,
    ) -> ThreadGetOut:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/thread/{thread_id}/state"
        data = self._request("PUT", path, json_body=payload.model_dump(exclude_none=True))
        return ThreadGetOut.model_validate(data)
    
    def thread_get_state(
        self,
        workspace_id: str,
        project_id: str,
        *,
        thread_id: str,
        thread_key: Optional[str] = None,
    ) -> ThreadStateOut:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/thread/{thread_id}/state"
        data = self._request("GET", path, params={"thread_key": thread_key})
        return ThreadStateOut.model_validate(data)
    
    # ---- Checkpoints ----

    def checkpoint_put(
        self,
        workspace_id: str,
        project_id: str,
        *,
        thread_id: str,
        payload: CheckpointPutIn,
    ) -> Dict[str, Any]:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/thread/{thread_id}/checkpoint"
        return self._request("POST", path, json_body=payload.model_dump(exclude_none=True)) or {}

    def checkpoint_get_latest(
        self,
        workspace_id: str,
        project_id: str,
        *,
        thread_id: str,
        thread_key: Optional[str] = None,
    ) -> Dict[str, Any]:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/thread/{thread_id}/checkpoint/latest"
        return self._request("GET", path, params={"thread_key": thread_key}) or {}

    # ---- Recall ----

    def recall_put(self, workspace_id: str, project_id: str, *, payload: RecallPutIn) -> Dict[str, Any]:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/recall:put"
        return self._request("POST", path, json_body=payload.model_dump(exclude_none=True)) or {}

    def recall_get(self, workspace_id: str, project_id: str, *, payload: RecallGetIn) -> Dict[str, Any]:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/recall:get"
        return self._request("POST", path, json_body=payload.model_dump(exclude_none=True)) or {}

    def recall_search(self, workspace_id: str, project_id: str, *, payload: RecallSearchIn) -> Dict[str, Any]:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/recall:search"
        return self._request("POST", path, json_body=payload.model_dump(exclude_none=True)) or {}

    def recall_forget(self, workspace_id: str, project_id: str, *, payload: Dict[str, Any]) -> Dict[str, Any]:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/recall:forget"
        return self._request("POST", path, json_body=payload) or {}