#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   stream.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Streaming client for Vi SDK localhost server.
"""

import json
import time
from collections.abc import Callable, Iterator
from typing import Any

import requests
from openai.types.chat import ChatCompletion
from requests.exceptions import RequestException, Timeout


class ViStreamingClient:
    """Enhanced client for streaming responses from Vi SDK server."""

    def __init__(self, base_url: str = "http://localhost:8000/v1"):
        self.base_url = base_url.rstrip("/")
        self.session = requests.Session()

    def create(
        self,
        model: str,
        messages: list[dict[str, Any]] | None = None,
        *,
        source: str | None = None,
        user_prompt: str | None = None,
        temperature: float = 0.7,
        max_tokens: int | None = None,
        stream: bool = True,
        show_progress: bool = True,
        progress_callback: Callable[[str], None] | None = None,
        timeout: int = 300,
        **kwargs,
    ) -> ChatCompletion | Iterator[str]:
        """Create a chat completion.

        Args:
            model: Model name (e.g., "nvila")
            messages: OpenAI-style messages (optional, for compatibility)
            source: Image source path (Vi SDK format)
            user_prompt: User prompt text (Vi SDK format)
            temperature: Sampling temperature
            max_tokens: Maximum tokens to generate
            stream: Whether to stream response
            show_progress: Show progress indicators
            progress_callback: Custom progress callback
            timeout: Request timeout in seconds
            **kwargs: Additional parameters

        Returns:
            - Iterator of response chunks if stream=True
            - ChatCompletion object if stream=False

        """
        # Support both Vi SDK format and OpenAI format
        if source is not None and user_prompt is not None:
            # Vi SDK native format
            if stream:
                return self._stream(
                    model=model,
                    source=source,
                    user_prompt=user_prompt,
                    temperature=temperature,
                    max_tokens=max_tokens,
                    stream=stream,
                    show_progress=show_progress,
                    progress_callback=progress_callback,
                    timeout=timeout,
                    **kwargs,
                )
            else:
                return self._complete(
                    model=model,
                    source=source,
                    user_prompt=user_prompt,
                    temperature=temperature,
                    max_tokens=max_tokens,
                    timeout=timeout,
                    **kwargs,
                )

        if messages is not None:
            # Convert OpenAI format to Vi SDK format
            image_url = None
            text_content = ""

            for message in messages:
                if message.get("role") == "user":
                    content = message.get("content", [])
                    if isinstance(content, list):
                        for item in content:
                            if item.get("type") == "image_url":
                                image_url = item["image_url"]["url"]
                            elif item.get("type") == "text":
                                text_content = item["text"]

            if image_url and text_content:
                if stream:
                    return self._stream(
                        model=model,
                        source=image_url,
                        user_prompt=text_content,
                        temperature=temperature,
                        max_tokens=max_tokens,
                        stream=stream,
                        show_progress=show_progress,
                        progress_callback=progress_callback,
                        timeout=timeout,
                        **kwargs,
                    )
                else:
                    return self._complete(
                        model=model,
                        source=image_url,
                        user_prompt=text_content,
                        temperature=temperature,
                        max_tokens=max_tokens,
                        timeout=timeout,
                        **kwargs,
                    )

            raise ValueError("OpenAI format requires both image_url and text content")

        raise ValueError(
            "Either 'source'/'user_prompt' (Vi format) or 'messages' (OpenAI format) must be provided"
        )

    def _complete(
        self,
        model: str,
        source: str,
        user_prompt: str,
        *,
        temperature: float = 0.7,
        max_tokens: int | None = None,
        timeout: int = 300,
        **kwargs,
    ) -> str:
        """Get complete chat completion response (non-streaming)."""
        payload = {
            "model": model,
            "source": source,
            "user_prompt": user_prompt,
            "temperature": temperature,
            "stream": False,  # Explicitly set to False for non-streaming
        }

        if max_tokens:
            payload["max_tokens"] = max_tokens

        # Add any additional kwargs
        payload.update(kwargs)

        try:
            response = self.session.post(
                f"{self.base_url}/chat/completions",
                json=payload,
                timeout=timeout,
            )
            response.raise_for_status()
            data = response.json()

            return ChatCompletion(**data)

        except Timeout as e:
            raise ViStreamingError(
                f"Request timed out. The model might be busy: {e}"
            ) from e
        except ConnectionError as e:
            raise ViStreamingError(
                f"Could not connect to Vi server. Is it running? {e}"
            ) from e
        except RequestException as e:
            raise ViStreamingError(f"Request failed: {e}") from e

    def _stream(
        self,
        model: str,
        source: str,
        user_prompt: str,
        *,
        temperature: float = 0.7,
        max_tokens: int | None = None,
        stream: bool = True,
        show_progress: bool = True,
        progress_callback: Callable[[str], None] | None = None,
        timeout: int = 300,
    ) -> Iterator[str]:
        """Stream chat completion with better error handling and UX.

        Args:
            model: Model name (e.g., "nvila")
            source: Path to image file
            user_prompt: Text prompt
            temperature: Sampling temperature
            max_tokens: Maximum tokens to generate
            stream: Whether to stream response
            show_progress: Show progress indicator
            progress_callback: Custom progress callback
            timeout: Request timeout in seconds

        Yields:
            Text chunks as they arrive

        Raises:
            ViStreamingError: For streaming-specific errors
            requests.RequestException: For HTTP errors

        """
        payload = {
            "model": model,
            "source": source,
            "user_prompt": user_prompt,
            "temperature": temperature,
            "stream": stream,
        }

        if max_tokens:
            payload["max_tokens"] = max_tokens

        try:
            response = self.session.post(
                f"{self.base_url}/chat/completions",
                json=payload,
                stream=True,
                timeout=timeout,
            )
            response.raise_for_status()

            return self._parse_stream(
                response,
                show_progress=show_progress,
                progress_callback=progress_callback,
            )

        except Timeout as e:
            raise ViStreamingError(
                f"Request timed out. The model might be busy: {e}"
            ) from e
        except ConnectionError as e:
            raise ViStreamingError(
                f"Could not connect to Vi server. Is it running? {e}"
            ) from e
        except RequestException as e:
            raise ViStreamingError(f"Request failed: {e}") from e

    def _parse_stream(
        self,
        response: requests.Response,
        show_progress: bool = True,
        progress_callback: Callable[[str], None] | None = None,
    ) -> Iterator[str]:
        """Parse streaming response with error handling."""
        start_time = time.time()
        chunk_count = 0

        try:
            for line_bytes in response.iter_lines():
                if not line_bytes:
                    continue

                # Handle SSE format
                if line_bytes.startswith(b"data: "):
                    chunk_data = line_bytes[6:].decode("utf-8")

                    # End of stream
                    if chunk_data == "[DONE]":
                        if show_progress and progress_callback:
                            elapsed = time.time() - start_time
                            progress_callback(
                                f"\n✓ Complete ({chunk_count} chunks, {elapsed:.1f}s)"
                            )
                        break

                    try:
                        chunk = json.loads(chunk_data)

                        # Extract content
                        content = self._extract_content(chunk)
                        if content:
                            chunk_count += 1

                            # Show progress indicator
                            if show_progress and chunk_count % 10 == 0:
                                if progress_callback:
                                    progress_callback(".")
                                else:
                                    print(".", end="", flush=True)

                            yield content

                    except json.JSONDecodeError:
                        # Skip malformed chunks but log for debugging
                        continue

        except Exception as e:
            raise ViStreamingError(f"Error parsing stream: {e}") from e

    def _extract_content(self, chunk: dict[str, Any]) -> str | None:
        """Safely extract content from chunk."""
        try:
            return chunk["choices"][0]["delta"]["content"]
        except (KeyError, IndexError, TypeError):
            return None


class ViStreamingError(Exception):
    """Custom exception for Vi streaming errors."""
