"""
Main Reno API client implementation.
"""

import json
import time
from typing import Dict, List, Optional, Union, Generator, Any

import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

from .exceptions import (
    RenoError,
    RenoConnectionError,
    RenoTimeoutError,
    RenoValidationError,
)
from .models import Completion, StreamChunk

# Valid message roles accepted by the API
_VALID_ROLES = frozenset({"user", "assistant", "system"})


def _parse_retry_after(headers: Any) -> Optional[float]:
    """Extract ``Retry-After`` seconds from response headers, if present."""
    value = headers.get("Retry-After")
    if value is None:
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _raise_from_error_body(err: dict, retry_after: Optional[float] = None) -> None:
    """Raise a :class:`RenoError` built from an API error dict."""
    raise RenoError(
        code=err.get("code", 9001),
        message=err.get("message", "Unknown error"),
        details=err.get("details") or err.get("suggestion"),
        retry_after=retry_after,
    )


class Reno:
    """
    Official Reno AI client — chat completions with robust error handling.

    Args:
        api_key:     Your Reno API key (must be non-empty).
        base_url:    Base URL for the API.
        timeout:     Request timeout in seconds (default: 30).
        max_retries: Number of automatic retries on transient failures (default: 3).

    Example::

        client = Reno(api_key="reno_sk_xxx")

        # Simple single-turn usage
        answer = client.ask("What is Python?")
        print(answer)

        # Multi-turn chat
        response = client.chat([
            {"role": "system",    "content": "You are a helpful assistant."},
            {"role": "user",      "content": "Hello!"},
        ])
        print(response.text)
    """

    DEFAULT_BASE_URL = "http://127.0.0.1:8000/api/v1"
    DEFAULT_TIMEOUT = 30
    DEFAULT_MODEL = "gemma2:2b-instruct"

    def __init__(
        self,
        api_key: str,
        base_url: str = DEFAULT_BASE_URL,
        timeout: int = DEFAULT_TIMEOUT,
        max_retries: int = 3,
    ) -> None:
        if not api_key or not api_key.strip():
            raise ValueError("api_key must be a non-empty string.")

        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.session = self._create_session(max_retries)
        self.session.headers.update(
            {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
                "User-Agent": "renoai-python/0.1.0",
            }
        )

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

    def _create_session(self, max_retries: int) -> requests.Session:
        """Return a :class:`requests.Session` pre-configured with retry logic."""
        session = requests.Session()
        retry_strategy = Retry(
            total=max_retries,
            status_forcelist=[429, 500, 502, 503, 504],
            allowed_methods=["POST"],
            backoff_factor=1,
            raise_on_status=False,
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        session.mount("http://", adapter)
        session.mount("https://", adapter)
        return session

    @staticmethod
    def _validate_messages(messages: List[Dict[str, str]]) -> None:
        """
        Raise :class:`RenoValidationError` if *messages* contains obvious problems.

        Checks performed:
        - List is non-empty.
        - Every item has ``role`` and ``content`` keys.
        - ``role`` is one of ``user``, ``assistant``, ``system``.
        - ``content`` is a non-empty string.
        """
        if not messages:
            raise RenoValidationError("'messages' must contain at least one message.")

        for i, msg in enumerate(messages):
            if not isinstance(msg, dict):
                raise RenoValidationError(
                    f"messages[{i}] must be a dict, got {type(msg).__name__}."
                )
            if "role" not in msg:
                raise RenoValidationError(f"messages[{i}] is missing the 'role' key.")
            if "content" not in msg:
                raise RenoValidationError(f"messages[{i}] is missing the 'content' key.")
            if msg["role"] not in _VALID_ROLES:
                raise RenoValidationError(
                    f"messages[{i}] has invalid role {msg['role']!r}. "
                    f"Must be one of {sorted(_VALID_ROLES)}."
                )
            if not isinstance(msg["content"], str) or not msg["content"].strip():
                raise RenoValidationError(
                    f"messages[{i}] 'content' must be a non-empty string."
                )

    def _handle_error_response(
        self, resp: requests.Response
    ) -> None:
        """
        Parse an HTTP error response and raise the appropriate exception.

        Handles both JSON error bodies and plain-text / unexpected responses.
        Always raises — never returns.
        """
        retry_after = _parse_retry_after(resp.headers)

        # Try to decode a structured JSON error body
        try:
            body = resp.json()
        except (json.JSONDecodeError, ValueError):
            # Non-JSON body (e.g. plain HTML from a proxy)
            raise RenoError(
                code=9001 if resp.status_code < 500 else 6001,
                message=f"HTTP {resp.status_code}: {resp.text[:200]}",
                retry_after=retry_after,
            )

        if "error" in body:
            err = body["error"]
            _raise_from_error_body(err, retry_after=retry_after)

        # ----------------------------------------------------------------
        # BUG FIX: previously the code could fall through here silently
        # when the JSON body had no "error" key.  We now always raise.
        # ----------------------------------------------------------------
        raise RenoError(
            code=9001,
            message=f"HTTP {resp.status_code}: unexpected response body",
            details=str(body)[:200],
            retry_after=retry_after,
        )

    def _request(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        """
        POST to ``/chat/completions/`` and return the parsed response dict.

        Raises:
            RenoError:           API returned an error payload.
            RenoConnectionError: Network-level connection failure.
            RenoTimeoutError:    Request exceeded :attr:`timeout`.
        """
        try:
            resp = self.session.post(
                f"{self.base_url}/chat/completions/",
                json=payload,
                timeout=self.timeout,
            )

            if resp.status_code >= 400:
                self._handle_error_response(resp)  # always raises

            try:
                data = resp.json()
            except (json.JSONDecodeError, ValueError) as exc:
                raise RenoError(
                    code=9001,
                    message="Invalid JSON in successful response.",
                    details=str(exc),
                )

            # Error embedded in a 2xx response
            if "error" in data:
                _raise_from_error_body(data["error"])

            return data.get("data", data)

        except RenoError:
            raise  # preserve already-typed exceptions
        except requests.exceptions.Timeout:
            raise RenoTimeoutError(
                f"Request timed out after {self.timeout}s. "
                "Try increasing the timeout or shortening your prompt."
            )
        except requests.exceptions.ConnectionError as exc:
            raise RenoConnectionError(
                f"Could not connect to {self.base_url}: {exc}"
            )
        except requests.exceptions.RequestException as exc:
            raise RenoError(code=9001, message=f"Request failed: {exc}")

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def chat(
        self,
        messages: List[Dict[str, str]],
        model: str = DEFAULT_MODEL,
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        stream: bool = False,
        **kwargs: Any,
    ) -> Union[Completion, Generator[StreamChunk, None, None]]:
        """
        Send a chat completion request with full control over parameters.

        Args:
            messages:    Conversation history as a list of ``{"role": ..., "content": ...}`` dicts.
            model:       Model identifier (default: ``gemma2:2b-instruct``).
            temperature: Sampling temperature in ``[0, 2]`` (default: ``0.7``).
            max_tokens:  Maximum tokens to generate (``None`` = server default).
            stream:      If ``True``, return a generator of :class:`StreamChunk` objects.
            **kwargs:    Additional parameters forwarded to the API payload.

        Returns:
            A :class:`Completion` for normal requests, or a generator of
            :class:`StreamChunk` objects when *stream* is ``True``.

        Raises:
            RenoValidationError: If *messages* fails client-side validation.
            RenoError:           If the API returns an error.

        Example::

            response = client.chat([
                {"role": "system", "content": "You are helpful."},
                {"role": "user",   "content": "What is the speed of light?"},
            ])
            print(response.text)
        """
        self._validate_messages(messages)

        if not (0 <= temperature <= 2):
            raise RenoValidationError(
                f"temperature must be in [0, 2], got {temperature}."
            )
        if max_tokens is not None and (not isinstance(max_tokens, int) or max_tokens < 1):
            raise RenoValidationError(
                "max_tokens must be a positive integer."
            )

        payload: Dict[str, Any] = {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "stream": stream,
            **kwargs,
        }
        if max_tokens is not None:
            payload["max_tokens"] = max_tokens

        if stream:
            return self._stream(payload)

        return Completion(self._request(payload))

    def ask(
        self,
        prompt: str,
        system: Optional[str] = None,
        model: str = DEFAULT_MODEL,
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
    ) -> str:
        """
        Simplified single-turn helper — returns only the answer string.

        Args:
            prompt:      The user's question or prompt.
            system:      Optional system message prepended to the conversation.
            model:       Model identifier (default: ``gemma2:2b-instruct``).
            temperature: Sampling temperature in ``[0, 2]`` (default: ``0.7``).
            max_tokens:  Maximum tokens to generate.

        Returns:
            The generated text as a plain string.

        Example::

            answer = client.ask("What is 2 + 2?")
            # "2 + 2 equals 4."
        """
        if not prompt or not prompt.strip():
            raise RenoValidationError("'prompt' must be a non-empty string.")

        messages: List[Dict[str, str]] = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        return self.chat(
            messages,
            model=model,
            temperature=temperature,
            max_tokens=max_tokens,
        ).text  # type: ignore[union-attr]  # stream=False → always Completion

    def stream_text(
        self,
        prompt: str,
        system: Optional[str] = None,
        model: str = DEFAULT_MODEL,
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
    ) -> Generator[str, None, None]:
        """
        Convenience streaming helper — yields plain text deltas.

        Example::

            for token in client.stream_text("Tell me a story"):
                print(token, end="", flush=True)
        """
        messages: List[Dict[str, str]] = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        for chunk in self.chat(  # type: ignore[union-attr]
            messages,
            model=model,
            temperature=temperature,
            max_tokens=max_tokens,
            stream=True,
        ):
            if chunk.delta:
                yield chunk.delta

    # ------------------------------------------------------------------
    # Internal streaming implementation
    # ------------------------------------------------------------------

    def _stream(
        self,
        payload: Dict[str, Any],
    ) -> Generator[StreamChunk, None, None]:
        """
        Handle Server-Sent Events streaming from ``/chat/completions/``.

        Yields:
            :class:`StreamChunk` objects, each carrying a text *delta*.

        Raises:
            RenoError:           Stream contains an error payload.
            RenoConnectionError: Connection failed during streaming.
            RenoTimeoutError:    Read timed out during streaming.
        """
        try:
            with self.session.post(
                f"{self.base_url}/chat/completions/",
                json=payload,
                stream=True,
                timeout=self.timeout,
            ) as resp:

                if resp.status_code >= 400:
                    self._handle_error_response(resp)  # always raises

                for raw_line in resp.iter_lines():
                    if not raw_line:
                        continue

                    line = (
                        raw_line.decode("utf-8")
                        if isinstance(raw_line, bytes)
                        else raw_line
                    )

                    if not line.startswith("data: "):
                        continue

                    data_str = line[6:]  # strip "data: " prefix

                    if data_str.strip() == "[DONE]":
                        return

                    try:
                        chunk_data = json.loads(data_str)
                    except json.JSONDecodeError:
                        # Malformed SSE chunk — skip and continue
                        continue

                    if "error" in chunk_data:
                        _raise_from_error_body(chunk_data["error"])

                    yield StreamChunk(chunk_data)

        # ----------------------------------------------------------------
        # BUG FIX: previously the except block checked
        #   isinstance(e, (RenoError, ...))
        # but RenoError is NOT a subclass of requests.exceptions.RequestException,
        # so that guard was always False and the bare `raise` was dead code.
        #
        # Now we explicitly re-raise our own exceptions first so they can
        # never be caught by the requests-exception handlers below.
        # ----------------------------------------------------------------
        except (RenoError, RenoConnectionError, RenoTimeoutError):
            raise
        except requests.exceptions.Timeout:
            raise RenoTimeoutError(
                f"Stream timed out after {self.timeout}s."
            )
        except requests.exceptions.ConnectionError as exc:
            raise RenoConnectionError(
                f"Connection lost while streaming from {self.base_url}: {exc}"
            )
        except requests.exceptions.RequestException as exc:
            raise RenoError(code=9001, message=f"Stream request failed: {exc}")

    # ------------------------------------------------------------------
    # Context manager / lifecycle
    # ------------------------------------------------------------------

    def close(self) -> None:
        """Close the underlying HTTP session and release resources."""
        self.session.close()

    def __enter__(self) -> "Reno":
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()

    def __repr__(self) -> str:
        return f"Reno(base_url={self.base_url!r})"


# ---------------------------------------------------------------------------
# Conversation helper
# ---------------------------------------------------------------------------

class Conversation:
    """
    Stateful multi-turn conversation helper built on top of :class:`Reno`.

    Keeps the message history for you so you can focus on the dialogue.

    Example::

        client = Reno(api_key="reno_sk_xxx")
        conv = Conversation(client, system="You are a pirate.")

        print(conv.say("Hello!"))
        print(conv.say("What is your ship called?"))
        print(conv.history)   # full message list
    """

    def __init__(
        self,
        client: Reno,
        system: Optional[str] = None,
        model: str = Reno.DEFAULT_MODEL,
        temperature: float = 0.7,
    ) -> None:
        self.client = client
        self.model = model
        self.temperature = temperature
        self.history: List[Dict[str, str]] = []

        if system:
            self.history.append({"role": "system", "content": system})

    def say(self, message: str, **kwargs: Any) -> str:
        """
        Send *message* as the user, append both it and the reply to history,
        and return the assistant's reply text.

        Args:
            message: The user's input.
            **kwargs: Extra parameters forwarded to :meth:`Reno.chat`.

        Returns:
            The assistant's reply as a plain string.
        """
        self.history.append({"role": "user", "content": message})

        response: Completion = self.client.chat(  # type: ignore[assignment]
            self.history,
            model=self.model,
            temperature=self.temperature,
            **kwargs,
        )

        self.history.append(response.to_message())
        return response.text

    def reset(self, keep_system: bool = True) -> None:
        """
        Clear the conversation history.

        Args:
            keep_system: If ``True`` (default), the system message is preserved.
        """
        if keep_system and self.history and self.history[0]["role"] == "system":
            self.history = [self.history[0]]
        else:
            self.history = []

    def __repr__(self) -> str:
        return f"Conversation(turns={len(self.history)}, model={self.model!r})"
