#!/usr/bin/env python

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

@File    :   vi_backend.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Vi inference backend using Vi's inference stack.
"""

from collections.abc import AsyncIterator
from pathlib import Path
from typing import Any

from vi.deployment.local.backends.base import GenerationResult, InferenceBackend
from vi.deployment.local.utils import (
    extract_image_and_prompt_from_messages,
    handle_vi_backend_source_error,
    validate_vi_format_request,
)
from vi.inference import ViModel
from vi.inference.responses import Stream


class ViBackend(InferenceBackend):
    """Vi inference backend using the Vi inference stack.

    Works with both local models and Datature models. Uses Vi's existing
    inference implementation with support for streaming, vision inputs,
    and structured outputs.
    """

    def __init__(
        self,
        pretrained_model_name_or_path: str | None = None,
        run_id: str | None = None,
        secret_key: str | None = None,
        organization_id: str | None = None,
        model_name: str | None = None,
        **kwargs: Any,
    ):
        """Initialize the Vi backend.

        Args:
            pretrained_model_name_or_path: Path to local model or HuggingFace model ID
            run_id: Run ID of model trained in Datature Vi
            secret_key: Datature API secret key (required with run_id)
            organization_id: Datature organization ID (required with run_id)
            model_name: Custom model identifier to use in API requests.
                If not provided, will be auto-generated from the model path/run_id
            **kwargs: Additional arguments passed to ViModel

        """
        self.pretrained_model_name_or_path = pretrained_model_name_or_path
        self.run_id = run_id
        self._kwargs = kwargs
        self._custom_model_name = model_name

        # Store credentials
        self._secret_key = secret_key
        self._organization_id = organization_id

        # Eagerly load model during initialization
        self._model = self._load_model()

    def _load_model(self) -> ViModel:
        """Load the model during initialization."""
        if self.pretrained_model_name_or_path:
            return ViModel(
                pretrained_model_name_or_path=self.pretrained_model_name_or_path,
                **self._kwargs,
            )
        if self.run_id:
            return ViModel(
                run_id=self.run_id,
                secret_key=self._secret_key,
                organization_id=self._organization_id,
                **self._kwargs,
            )
        raise ValueError(
            "Either pretrained_model_name_or_path or run_id must be provided"
        )

    async def generate_vi_format(
        self,
        source: str | None = None,
        user_prompt: str | None = None,
        response_format: Any | None = None,
        generation_config: dict | None = None,
        **kwargs: Any,
    ) -> GenerationResult:
        """Generate a complete response using Vi SDK format (non-streaming).

        Args:
            source: Image source path (can be None for text-only requests)
            user_prompt: Text prompt (can be None)
            response_format: Optional parsed Pydantic model for structured output
            generation_config: Dictionary of generation parameters (temperature, max_tokens, etc.)
            **kwargs: Additional generation parameters

        Returns:
            GenerationResult with complete response and token counts

        """
        # Validate request - need at least a prompt or source
        validate_vi_format_request(source, user_prompt)

        # response_format is already a parsed Pydantic model
        pydantic_schema = response_format

        # Use the generation_config passed from server (already in correct format)
        final_generation_config = generation_config or {}

        # Build kwargs for ViModel
        gen_kwargs = {
            "generation_config": final_generation_config,
            "stream": False,
            **kwargs,  # Additional parameters
        }

        # Add response_format if provided
        if pydantic_schema:
            gen_kwargs["response_format"] = pydantic_schema

        # Run inference
        try:
            result, error = self._model(
                source=source,
                user_prompt=user_prompt,
                show_progress=False,
                **gen_kwargs,
            )

            if error:
                raise error
        except ValueError as e:
            # Provide more helpful error messages
            enhanced_error = handle_vi_backend_source_error(e, source, user_prompt)
            raise enhanced_error from e

        # Extract text from result
        if hasattr(result, "caption"):
            text = result.caption
        elif hasattr(result, "response"):
            text = result.response
        else:
            text = str(result)

        # Estimate token counts (rough approximation)
        prompt_tokens = len(user_prompt.split()) * 2 if user_prompt else 0
        completion_tokens = len(text.split()) * 2

        return GenerationResult(
            text=text,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            finish_reason="stop",
        )

    async def generate(
        self,
        messages: list[dict] | None,
        response_format: Any | None = None,
        generation_config: dict | None = None,
        **kwargs: Any,
    ) -> GenerationResult:
        """Generate a complete response (non-streaming).

        Args:
            messages: List of chat messages in OpenAI format
            response_format: Optional parsed Pydantic model for structured output
            generation_config: Dictionary of generation parameters (temperature, max_tokens, etc.)
            **kwargs: Additional generation parameters

        Returns:
            GenerationResult with complete response and token counts

        """
        image_path, prompt = extract_image_and_prompt_from_messages(messages)

        # Use Vi format internally
        return await self.generate_vi_format(
            source=image_path,
            user_prompt=prompt,
            response_format=response_format,
            generation_config=generation_config,
            **kwargs,
        )

    async def generate_stream_vi_format(
        self,
        source: str | None = None,
        user_prompt: str | None = None,
        response_format: Any | None = None,
        generation_config: dict | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[str]:
        """Generate a streaming response using Vi SDK format.

        Args:
            source: Image source path (can be None for text-only requests)
            user_prompt: Text prompt (can be None)
            response_format: Optional parsed Pydantic model for structured output
            generation_config: Dictionary of generation parameters (temperature, max_tokens, etc.)
            **kwargs: Additional generation parameters

        Yields:
            Tokens as strings

        """
        # Validate request - need at least a prompt or source
        validate_vi_format_request(source, user_prompt)

        # response_format is already a parsed Pydantic model
        pydantic_schema = response_format

        # Use the generation_config passed from server (already in correct format)
        final_generation_config = generation_config or {}

        # Build kwargs for ViModel
        gen_kwargs = {
            "generation_config": final_generation_config,
            "stream": True,
            **kwargs,  # Additional parameters
        }

        # Add response_format if provided
        if pydantic_schema:
            gen_kwargs["response_format"] = pydantic_schema

        # Run inference with streaming
        try:
            stream: Stream = self._model(
                source=source,
                user_prompt=user_prompt,
                show_progress=False,
                **gen_kwargs,
            )

            # Yield tokens from the stream with proper error handling
            try:
                for token in stream:
                    # Ensure we only yield strings
                    if isinstance(token, str):
                        yield token
                    elif token is not None:
                        # Convert non-string tokens to string
                        yield str(token)
                    # Skip None tokens silently
            except Exception as iter_error:
                # If iteration fails, raise a more informative error
                raise RuntimeError(
                    f"Error during streaming: {type(iter_error).__name__}: {iter_error}"
                ) from iter_error
        except ValueError as e:
            # Provide more helpful error messages
            enhanced_error = handle_vi_backend_source_error(e, source, user_prompt)
            raise enhanced_error from e

    async def generate_stream(
        self,
        messages: list[dict] | None,
        response_format: Any | None = None,
        generation_config: dict | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[str]:
        """Generate a streaming response.

        Args:
            messages: List of chat messages in OpenAI format
            response_format: Optional parsed Pydantic model for structured output
            generation_config: Dictionary of generation parameters (temperature, max_tokens, etc.)
            **kwargs: Additional generation parameters

        Yields:
            Tokens as strings

        """
        image_path, prompt = extract_image_and_prompt_from_messages(messages)

        # Use Vi format internally
        async for token in self.generate_stream_vi_format(
            source=image_path,
            user_prompt=prompt,
            response_format=response_format,
            generation_config=generation_config,
            **kwargs,
        ):
            yield token

    def available_models(self) -> list[str]:
        """Get list of available model identifiers.

        Returns:
            List containing the model ID

        """
        # Use custom model name if provided
        if self._custom_model_name:
            return [self._custom_model_name]

        # Otherwise, auto-generate from model source
        if self.pretrained_model_name_or_path:
            # For HuggingFace model IDs (contains /), return full ID
            # For local paths, return just the directory name
            if (
                "/" in self.pretrained_model_name_or_path
                and not self.pretrained_model_name_or_path.startswith("/")
            ):
                return [self.pretrained_model_name_or_path]
            return [Path(self.pretrained_model_name_or_path).name]
        if self.run_id:
            return [self.run_id]
        return ["unknown-model"]
