"""OmniLink Python API Client.

A single class — ``OmniLinkClient`` — that wraps every public OmniLink REST
endpoint authenticated by an Omni Key.

Covered endpoints
-----------------
* ``/api/chat``              — multi-turn AI chat
* ``/api/short-term-memory`` — read / write per-agent conversation memory
* ``/api/stt``               — speech-to-text (Whisper)
* ``/api/tts``               — text-to-speech (Google Chirp3-HD)
* ``/api/translate``         — text translation (GPT-5)

Quick start
-----------
::

    from omnilink.client import OmniLinkClient

    client = OmniLinkClient(omni_key="omk_...")

    # Chat
    reply = client.chat("What is 2 + 2?", agent_name="math-tutor")
    print(reply["text"])

    # Memory
    client.set_memory("math-tutor", [
        {"role": "user",  "parts": [{"text": "Hello"}]},
        {"role": "model", "parts": [{"text": "Hi there!"}]},
    ])
    mem = client.get_memory("math-tutor")

    # STT
    with open("audio.webm", "rb") as f:
        result = client.transcribe(f.read(), mime_type="audio/webm")
    print(result["text"])

    # TTS
    audio_bytes = client.synthesize_to_bytes("Hello world")
    with open("out.mp3", "wb") as f:
        f.write(audio_bytes)

    # Translate
    result = client.translate("Bonjour le monde", target_language="English")
    print(result["translation"])
"""

from __future__ import annotations

import base64
from typing import Any, Dict, List, Optional

try:
    import requests  # type: ignore[import-not-found]
except Exception:  # pragma: no cover
    requests = None  # type: ignore[assignment]


class OmniLinkAPIError(Exception):
    """Raised when OmniLink returns a non-2xx response.

    Attributes
    ----------
    status_code : int
        HTTP status code returned by the server.
    body : dict
        Parsed JSON error body (or ``{"raw": ...}`` if the body was not JSON).
    """

    def __init__(self, status_code: int, body: Dict[str, Any]) -> None:
        self.status_code = status_code
        self.body = body
        message = body.get("error") or body.get("message") or str(body)
        super().__init__(f"OmniLink API error [{status_code}]: {message}")


class OmniLinkClient:
    """Full Python client for the OmniLink REST API.

    All methods are synchronous and thread-safe (each call creates its own
    ``requests.Session``).  Every method raises :class:`OmniLinkAPIError` on a
    non-2xx response and ``requests.RequestException`` on a network failure.

    Parameters
    ----------
    omni_key:
        Bearer token (Omni Key) used to authenticate every request.
        Omni Keys look like ``omk_...``.
    base_url:
        Root URL of the OmniLink deployment.
        Defaults to ``https://www.omnilink-agents.com``.
    timeout:
        HTTP request timeout in seconds (default: 30).

    Example
    -------
    ::

        client = OmniLinkClient(omni_key="omk_...", base_url="https://www.omnilink-agents.com")
        print(client.chat("Hello!", agent_name="assistant")["text"])
    """

    DEFAULT_BASE_URL: str = "https://www.omnilink-agents.com"

    def __init__(
        self,
        omni_key: str,
        base_url: str = DEFAULT_BASE_URL,
        *,
        timeout: int = 30,
    ) -> None:
        if requests is None:  # pragma: no cover
            raise ImportError(
                "The 'requests' package is required. Install it with: pip install requests"
            )
        if not omni_key or not isinstance(omni_key, str):
            raise ValueError("omni_key must be a non-empty string.")
        self._omni_key = omni_key
        self._base_url = base_url.rstrip("/")
        self._timeout = timeout

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _headers(self) -> Dict[str, str]:
        return {
            "Authorization": f"Bearer {self._omni_key}",
            "Content-Type": "application/json",
        }

    def _url(self, path: str) -> str:
        return f"{self._base_url}/{path.lstrip('/')}"

    def _raise_for_response(self, response: Any) -> None:
        if not response.ok:
            try:
                body = response.json()
            except Exception:
                body = {"raw": response.text}
            raise OmniLinkAPIError(response.status_code, body)

    def _post(self, path: str, body: Dict[str, Any]) -> Dict[str, Any]:
        response = requests.post(
            self._url(path),
            json=body,
            headers=self._headers(),
            timeout=self._timeout,
        )
        self._raise_for_response(response)
        return response.json()

    def _get(
        self, path: str, params: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        response = requests.get(
            self._url(path),
            params=params,
            headers=self._headers(),
            timeout=self._timeout,
        )
        self._raise_for_response(response)
        return response.json()

    def _put(self, path: str, body: Dict[str, Any]) -> Dict[str, Any]:
        response = requests.put(
            self._url(path),
            json=body,
            headers=self._headers(),
            timeout=self._timeout,
        )
        self._raise_for_response(response)
        return response.json()

    # ------------------------------------------------------------------
    # Chat  (/api/chat)
    # ------------------------------------------------------------------

    def chat(
        self,
        prompt: Optional[str] = None,
        *,
        agent_name: str,
        messages: Optional[List[Dict[str, Any]]] = None,
        engine: str = "g2-engine",
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        system_instruction: Optional[Dict[str, Any]] = None,
        system_instruction_suffix: Optional[str] = None,
        use_prompt_pipeline: bool = False,
        **extra: Any,
    ) -> Dict[str, Any]:
        """Send a chat request and return the parsed response.

        Either *prompt* (a plain string) **or** *messages* (an OpenAI-style
        list) must be provided; if both are given, *messages* takes precedence.

        Parameters
        ----------
        prompt:
            Shorthand for a single user message.
        agent_name:
            Target agent name. Required.
        messages:
            Full conversation array, e.g.
            ``[{"role": "user", "content": "Hi"}]``.
        engine:
            AI engine:
            ``"g1-engine"`` (Gemini), ``"g2-engine"`` (GPT-5, **default**),
            ``"g3-engine"`` (Grok), ``"g4-engine"`` (Claude).
        temperature:
            Sampling temperature forwarded to the engine.
        max_tokens:
            Token limit forwarded to the engine.
        system_instruction:
            Structured system-instruction object forwarded to the engine.
        system_instruction_suffix:
            Additional text appended to the system instruction.
        use_prompt_pipeline:
            When ``True`` the server composes system instructions
            automatically from agent context.
        **extra:
            Any additional fields to merge into the request body (e.g.
            ``agentPersona``, ``knowledgeMode``, ``customInstructions``).

        Returns
        -------
        dict
            On success: ``{"ok": True, "text": "...", "requestId": "...", ...}``.

        Raises
        ------
        ValueError
            When neither *prompt* nor *messages* is provided.
        OmniLinkAPIError
            On a non-2xx server response.
        """
        if messages is None and prompt is None:
            raise ValueError("Either 'prompt' or 'messages' must be provided.")

        body: Dict[str, Any] = {
            "agentName": agent_name,
            "engine": engine,
            "usePromptPipeline": use_prompt_pipeline,
        }

        if messages is not None:
            body["messages"] = messages
        else:
            body["prompt"] = prompt

        if temperature is not None:
            body["temperature"] = temperature
        if max_tokens is not None:
            body["max_tokens"] = max_tokens
        if system_instruction is not None:
            body["systemInstructionRequest"] = system_instruction
        if system_instruction_suffix is not None:
            body["systemInstructionSuffix"] = system_instruction_suffix

        body.update(extra)
        return self._post("/api/chat", body)

    # ------------------------------------------------------------------
    # Agent profiles  (/api/agent-profiles)
    # ------------------------------------------------------------------

    def list_profiles(self) -> List[Dict[str, Any]]:
        """Return all agent profiles for the authenticated user.

        Returns
        -------
        list
            List of profile dicts with ``id``, ``name``, ``settings``,
            ``appearance``, ``plan_locked``, and ``updated_at``.

        Raises
        ------
        OmniLinkAPIError
            On a non-2xx server response.
        """
        result = self._get("/api/agent-profiles")
        return result.get("profiles", [])

    def create_profile(
        self,
        name: str,
        *,
        settings: Optional[Dict[str, Any]] = None,
        appearance: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Create a new agent profile.

        Parameters
        ----------
        name:
            Display name for the profile (e.g. ``"door-controller"``).
        settings:
            Optional JSONB settings object for the profile.
        appearance:
            Optional JSONB appearance object for the profile.

        Returns
        -------
        dict
            The created profile with ``id``, ``name``, ``settings``,
            ``appearance``, ``plan_locked``, and ``updated_at``.

        Raises
        ------
        OmniLinkAPIError
            On a non-2xx server response (including plan limit errors).
        """
        body: Dict[str, Any] = {"name": name}
        if settings is not None:
            body["settings"] = settings
        if appearance is not None:
            body["appearance"] = appearance
        result = self._post("/api/agent-profiles", body)
        return result.get("profile", result)

    def update_profile(
        self,
        profile_id: str,
        *,
        name: Optional[str] = None,
        settings: Optional[Dict[str, Any]] = None,
        appearance: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Update an existing agent profile.

        Parameters
        ----------
        profile_id:
            UUID of the profile to update.
        name:
            New display name (optional).
        settings:
            New settings object (optional).
        appearance:
            New appearance object (optional).

        Returns
        -------
        dict
            The updated profile.

        Raises
        ------
        OmniLinkAPIError
            On a non-2xx server response.
        """
        body: Dict[str, Any] = {"id": profile_id}
        if name is not None:
            body["name"] = name
        if settings is not None:
            body["settings"] = settings
        if appearance is not None:
            body["appearance"] = appearance
        result = self._put("/api/agent-profiles", body)
        return result.get("profile", result)

    def delete_profile(self, profile_id: str) -> str:
        """Delete an agent profile.

        Parameters
        ----------
        profile_id:
            UUID of the profile to delete.

        Returns
        -------
        str
            The ID of the deleted profile.

        Raises
        ------
        OmniLinkAPIError
            On a non-2xx server response.
        """
        response = requests.request(
            "DELETE",
            self._url("/api/agent-profiles"),
            json={"id": profile_id},
            headers=self._headers(),
            timeout=self._timeout,
        )
        self._raise_for_response(response)
        return response.json().get("deleted", profile_id)

    # ------------------------------------------------------------------
    # Short-term memory  (/api/short-term-memory)
    # ------------------------------------------------------------------

    def list_agents(self) -> Dict[str, Any]:
        """Return all agent profiles and their short-term memory status.

        Returns
        -------
        dict
            ``{"agents": [...], "memoryAgents": [...]}``.

            * ``agents`` — agent profiles with ``id``, ``name``,
              ``updated_at``, and a boolean ``hasMemory`` flag.
            * ``memoryAgents`` — raw ``agent_name`` values found in
              short-term memories together with ``updated_at``.

        Raises
        ------
        OmniLinkAPIError
            On a non-2xx server response.
        """
        return self._get("/api/short-term-memory")

    def get_memory(self, agent_name: str) -> Optional[List[Dict[str, Any]]]:
        """Return the stored conversation for *agent_name*, or ``None``.

        Parameters
        ----------
        agent_name:
            Name (or profile ID) of the agent whose memory to retrieve.

        Returns
        -------
        list or None
            List of conversation entries, or ``None`` if no memory exists.

        Raises
        ------
        OmniLinkAPIError
            On a non-2xx server response.
        """
        result = self._get(
            "/api/short-term-memory", params={"agentName": agent_name}
        )
        return result.get("conversation")

    def set_memory(
        self,
        agent_name: str,
        conversation: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """Upsert the conversation memory for *agent_name*.

        Parameters
        ----------
        agent_name:
            Target agent name.
        conversation:
            List of conversation entries. Each entry must have a ``role``
            (``"user"`` or ``"model"``) and ``parts`` (list of
            ``{"text": "..."}`` dicts).  Example::

                [
                    {"role": "user",  "parts": [{"text": "Hello"}]},
                    {"role": "model", "parts": [{"text": "Hi!"}]},
                ]

        Returns
        -------
        list
            The sanitized conversation that was persisted.

        Raises
        ------
        OmniLinkAPIError
            On a non-2xx server response.
        """
        result = self._post(
            "/api/short-term-memory",
            {"agentName": agent_name, "conversation": conversation},
        )
        return result.get("conversation", [])

    def clear_memory(self, agent_name: str) -> None:
        """Delete the conversation memory for *agent_name* by writing an empty list.

        Parameters
        ----------
        agent_name:
            Target agent name.

        Raises
        ------
        OmniLinkAPIError
            On a non-2xx server response.
        """
        self.set_memory(agent_name, [])

    # ------------------------------------------------------------------
    # Speech-to-text  (/api/stt)
    # ------------------------------------------------------------------

    def transcribe(
        self,
        audio: bytes,
        *,
        mime_type: str = "audio/webm",
        language: str = "",
        duration_sec: float = 0.0,
    ) -> Dict[str, Any]:
        """Transcribe *audio* bytes to text using Whisper.

        Parameters
        ----------
        audio:
            Raw audio bytes (e.g. read directly from a file).
        mime_type:
            MIME type of the audio, e.g. ``"audio/webm"`` (default),
            ``"audio/mp4"``, ``"audio/wav"``.
        language:
            Optional ISO-639-1 hint (e.g. ``"en"``, ``"fr"``).  Leave
            empty to let Whisper auto-detect.
        duration_sec:
            Optional audio duration in seconds.  Used for usage tracking
            only; does not affect transcription quality.

        Returns
        -------
        dict
            Whisper response, typically ``{"text": "transcribed text..."}``.

        Raises
        ------
        OmniLinkAPIError
            On a non-2xx server response.
        """
        audio_b64 = base64.b64encode(audio).decode("ascii")
        body: Dict[str, Any] = {
            "audioBase64": audio_b64,
            "mimeType": mime_type,
        }
        if language:
            body["language"] = language
        if duration_sec > 0:
            body["durationSec"] = duration_sec
        return self._post("/api/stt", body)

    # ------------------------------------------------------------------
    # Text-to-speech  (/api/tts)
    # ------------------------------------------------------------------

    def synthesize(
        self,
        text: str,
        *,
        language_code: str = "en-US",
        voice_name: Optional[str] = None,
        audio_encoding: str = "MP3",
        speaking_rate: float = 1.0,
        sample_rate_hertz: Optional[int] = None,
    ) -> Dict[str, Any]:
        """Convert *text* to speech using Google Chirp3-HD.

        Parameters
        ----------
        text:
            The text to synthesize.
        language_code:
            BCP-47 language code, e.g. ``"en-US"`` (default), ``"fr-FR"``.
        voice_name:
            Chirp3-HD voice name. When ``None`` the server selects a
            suitable default for *language_code*.
        audio_encoding:
            Audio format: ``"MP3"`` (default), ``"LINEAR16"``,
            ``"OGG_OPUS"``, ``"MULAW"``, ``"ALAW"``.
        speaking_rate:
            Speaking speed multiplier (default: ``1.0``).
        sample_rate_hertz:
            Optional sample rate.  Omit to let the server choose.

        Returns
        -------
        dict
            ``{"audioContent": "<base64>", "audioMimeType": "audio/mpeg",
            "meta": {...}}``.

        Raises
        ------
        OmniLinkAPIError
            On a non-2xx server response.
        """
        voice: Dict[str, Any] = {"languageCode": language_code}
        if voice_name is not None:
            voice["name"] = voice_name

        audio_config: Dict[str, Any] = {
            "audioEncoding": audio_encoding,
            "speakingRate": speaking_rate,
        }
        if sample_rate_hertz is not None:
            audio_config["sampleRateHertz"] = sample_rate_hertz

        body: Dict[str, Any] = {
            "input": {"text": text},
            "voice": voice,
            "audioConfig": audio_config,
        }
        return self._post("/api/tts", body)

    def synthesize_to_bytes(
        self,
        text: str,
        **kwargs: Any,
    ) -> bytes:
        """Convenience wrapper — synthesize *text* and return raw audio bytes.

        All keyword arguments are forwarded to :meth:`synthesize`.

        Returns
        -------
        bytes
            Raw audio bytes (MP3 by default).

        Raises
        ------
        OmniLinkAPIError
            On a non-2xx server response.
        ValueError
            When the server returns no ``audioContent`` field.
        """
        result = self.synthesize(text, **kwargs)
        audio_b64 = result.get("audioContent")
        if not isinstance(audio_b64, str):
            raise ValueError(
                f"Unexpected TTS response — no audioContent field: {result}"
            )
        return base64.b64decode(audio_b64)

    # ------------------------------------------------------------------
    # Translation  (/api/translate)
    # ------------------------------------------------------------------

    def translate(
        self,
        text: str,
        *,
        target_language: str = "English",
        source_language: Optional[str] = None,
        model: str = "gpt-5",
        max_tokens: Optional[int] = None,
    ) -> Dict[str, Any]:
        """Translate *text* into *target_language*.

        Parameters
        ----------
        text:
            The text to translate.
        target_language:
            Target language name or code, e.g. ``"English"`` (default),
            ``"French"``, ``"Spanish"``, ``"ar"``.
        source_language:
            Optional source language.  When omitted the engine auto-detects.
        model:
            GPT model to use for translation (default: ``"gpt-5"``).
        max_tokens:
            Optional upper bound on the translation length.

        Returns
        -------
        dict
            ``{"translation": "translated text"}``.

        Raises
        ------
        OmniLinkAPIError
            On a non-2xx server response.
        """
        body: Dict[str, Any] = {
            "text": text,
            "targetLanguage": target_language,
            "model": model,
        }
        if source_language is not None:
            body["sourceLanguage"] = source_language
        if max_tokens is not None:
            body["max_tokens"] = max_tokens
        return self._post("/api/translate", body)
