# Author: Koushik Sen (ksen@berkeley.edu)
# Contributors:
# Koushik Sen (ksen@berkeley.edu)
# add your name here

"""Abstract base class for LLM provider model implementations."""

import asyncio
import base64
import dataclasses
import inspect
import logging
import mimetypes
import threading
import types as types_module
from abc import ABC, abstractmethod
from collections.abc import Callable, Coroutine
from pathlib import Path
from typing import Any, Union, get_args, get_origin

logger = logging.getLogger(__name__)

# Type alias for the async token streaming callback.
TokenCallback = Callable[[str], Coroutine[Any, Any, None]]

SUPPORTED_MIME_TYPES = {
    "image/jpeg",
    "image/png",
    "image/gif",
    "image/webp",
    "application/pdf",
}


@dataclasses.dataclass
class Attachment:
    """A file attachment (image or document) to include in a prompt.

    Attributes:
        data: Raw file bytes.
        mime_type: MIME type string (e.g. "image/jpeg", "application/pdf").
    """

    data: bytes
    mime_type: str

    @staticmethod
    def from_file(path: str) -> "Attachment":
        """Create an Attachment from a file path.

        Args:
            path: Path to the file to attach.

        Returns:
            An Attachment with the file's bytes and detected MIME type.

        Raises:
            ValueError: If the MIME type is not supported.
            FileNotFoundError: If the file does not exist.
        """
        file_path = Path(path)
        mime_type, _ = mimetypes.guess_type(str(file_path))
        if mime_type is None:
            suffix = file_path.suffix.lower()
            mime_map = {
                ".jpg": "image/jpeg",
                ".jpeg": "image/jpeg",
                ".png": "image/png",
                ".gif": "image/gif",
                ".webp": "image/webp",
                ".pdf": "application/pdf",
            }
            mime_type = mime_map.get(suffix, "")
        if mime_type not in SUPPORTED_MIME_TYPES:
            raise ValueError(
                f"Unsupported MIME type '{mime_type}' for file '{path}'. "
                f"Supported: {sorted(SUPPORTED_MIME_TYPES)}"
            )
        return Attachment(data=file_path.read_bytes(), mime_type=mime_type)

    def to_base64(self) -> str:
        """Return the file data as a base64-encoded string."""
        return base64.b64encode(self.data).decode("ascii")

    def to_data_url(self) -> str:
        """Return a data: URL suitable for OpenAI image_url fields."""
        return f"data:{self.mime_type};base64,{self.to_base64()}"


_callback_helper_loop: asyncio.AbstractEventLoop | None = None
_callback_helper_ready = threading.Event()
_callback_helper_lock = threading.Lock()


def _get_callback_loop() -> asyncio.AbstractEventLoop:
    global _callback_helper_loop
    with _callback_helper_lock:
        if _callback_helper_loop is not None and not _callback_helper_loop.is_closed():
            return _callback_helper_loop

        def run_loop() -> None:
            global _callback_helper_loop
            _callback_helper_loop = asyncio.new_event_loop()
            asyncio.set_event_loop(_callback_helper_loop)
            _callback_helper_ready.set()
            _callback_helper_loop.run_forever()

        _callback_helper_ready.clear()
        t = threading.Thread(target=run_loop, daemon=True)
        t.start()
        _callback_helper_ready.wait(timeout=5)
        if _callback_helper_loop is None or _callback_helper_loop.is_closed():
            raise RuntimeError("Callback helper loop failed to start")
        return _callback_helper_loop


class Model(ABC):
    """Abstract base class for LLM provider implementations."""

    def __init__(
        self,
        model_name: str,
        model_config: dict[str, Any] | None = None,
        token_callback: TokenCallback | None = None,
    ):
        """Initialize a Model instance.

        Args:
            model_name: The name/identifier of the model.
            model_config: Optional dictionary of model configuration parameters.
            token_callback: Optional async callback invoked with each streamed text token.
        """
        self.model_name = model_name
        self.model_config = model_config or {}
        self.token_callback = token_callback
        self.usage_info_for_messages: str = ""
        self.conversation: list[Any] = []
        self.client: Any = None
        self._callback_loop: asyncio.AbstractEventLoop | None = None

    def _invoke_token_callback(self, token: str) -> None:
        """Invoke the async token_callback synchronously, preserving order."""
        if self.token_callback is None:
            return
        try:
            running_loop = asyncio.get_running_loop()
        except RuntimeError:
            logger.debug("Exception caught", exc_info=True)
            running_loop = None
        if running_loop and running_loop.is_running():
            helper_loop = _get_callback_loop()
            future = asyncio.run_coroutine_threadsafe(self.token_callback(token), helper_loop)
            future.result(timeout=30)
            return
        if self._callback_loop is None or self._callback_loop.is_closed():
            self._callback_loop = asyncio.new_event_loop()
        self._callback_loop.run_until_complete(self.token_callback(token))

    def close_callback_loop(self) -> None:
        """Close the per-instance event loop used for synchronous token callback invocation.

        Safe to call multiple times; subsequent calls are no-ops.
        """
        if self._callback_loop is not None and not self._callback_loop.is_closed():
            self._callback_loop.close()
        self._callback_loop = None

    def __del__(self) -> None:
        self.close_callback_loop()

    def __str__(self) -> str:
        return f"{self.__class__.__name__}(name={self.model_name})"

    __repr__ = __str__

    @abstractmethod
    def initialize(self, prompt: str, attachments: list[Attachment] | None = None) -> None:
        """Initializes the conversation with an initial user prompt.

        Args:
            prompt: The initial user prompt to start the conversation.
            attachments: Optional list of file attachments (images, PDFs) to include.
        """
        pass

    @abstractmethod
    def generate(self) -> tuple[str, Any]:
        """Generates content from prompt.

        Returns:
            tuple[str, Any]: A tuple of (generated_text, raw_response).
        """
        pass

    @abstractmethod
    def generate_and_process_with_tools(
        self, function_map: dict[str, Callable[..., Any]]
    ) -> tuple[list[dict[str, Any]], str, Any]:
        """Generates content with tools, processes the response, and adds it to conversation.

        Args:
            function_map: Dictionary mapping function names to callable functions.

        Returns:
            tuple[list[dict[str, Any]], str, Any]: A tuple of
                (function_calls, response_text, raw_response).
        """
        pass

    def add_function_results_to_conversation_and_return(
        self, function_results: list[tuple[str, dict[str, Any]]]
    ) -> None:
        """Adds function results to the conversation state.

        Matches results to tool calls by index from the last assistant message.

        Args:
            function_results: List of tuples containing (function_name, result_dict).
        """
        tool_calls: list[dict[str, str]] = []
        for msg in reversed(self.conversation):
            if msg.get("role") == "assistant" and msg.get("tool_calls"):
                tool_calls = [
                    {"name": tc["function"]["name"], "id": tc["id"]} for tc in msg["tool_calls"]
                ]
                break

        for i, (func_name, result_dict) in enumerate(function_results):
            result_content = result_dict.get("result", str(result_dict))
            if self.usage_info_for_messages:
                result_content = f"{result_content}\n\n{self.usage_info_for_messages}"

            if i < len(tool_calls):
                tool_call_id = tool_calls[i]["id"]
            else:
                tool_call_id = f"call_{func_name}_{i}"

            self.conversation.append(
                {
                    "role": "tool",
                    "tool_call_id": tool_call_id,
                    "content": result_content,
                }
            )

    def add_message_to_conversation(self, role: str, content: str) -> None:
        """Adds a message to the conversation state.

        Args:
            role: The role of the message sender (e.g., 'user', 'assistant').
            content: The message content.
        """
        if role == "user" and self.usage_info_for_messages:
            content = f"{content}\n\n{self.usage_info_for_messages}"
        self.conversation.append({"role": role, "content": content})

    @abstractmethod
    def extract_input_output_token_counts_from_response(
        self, response: Any
    ) -> tuple[int, int, int, int]:
        """Extracts token counts from an API response.

        Args:
            response: The raw API response object.

        Returns:
            tuple[int, int, int, int]: (input_tokens, output_tokens,
                cache_read_tokens, cache_write_tokens).
        """
        pass

    @abstractmethod
    def get_embedding(self, text: str, embedding_model: str | None = None) -> list[float]:
        """Generates an embedding vector for the given text.

        Args:
            text: The text to generate an embedding for.
            embedding_model: Optional model name to use for embedding generation.

        Returns:
            list[float]: The embedding vector as a list of floats.
        """
        pass

    def set_usage_info_for_messages(self, usage_info: str) -> None:
        """Sets token information to append to messages sent to the LLM.

        Args:
            usage_info: The usage information string to append.
        """
        self.usage_info_for_messages = usage_info

    # =========================================================================
    # Helper methods for building tool schemas (shared across implementations)
    # =========================================================================

    def _build_openai_tools_schema(
        self, function_map: dict[str, Callable[..., Any]]
    ) -> list[dict[str, Any]]:
        """Builds the OpenAI-compatible tools schema from a function map.

        Args:
            function_map: Dictionary mapping function names to callable functions.

        Returns:
            list[dict[str, Any]]: A list of tool schemas in OpenAI format.
        """
        tools = []
        for func in function_map.values():
            tool_schema = self._function_to_openai_tool(func)
            tools.append(tool_schema)
        return tools

    def _function_to_openai_tool(self, func: Callable[..., Any]) -> dict[str, Any]:
        """Converts a Python function to an OpenAI tool schema.

        Args:
            func: The Python function to convert.

        Returns:
            dict[str, Any]: The tool schema in OpenAI format.
        """
        sig = inspect.signature(func)
        doc = inspect.getdoc(func) or ""

        # Parse docstring for parameter descriptions
        param_descriptions = self._parse_docstring_params(doc)

        # Build parameters schema
        properties: dict[str, Any] = {}
        required: list[str] = []

        for param_name, param in sig.parameters.items():
            param_type = param.annotation
            param_schema = self._python_type_to_json_schema(param_type)

            # Add description from docstring if available
            if param_name in param_descriptions:
                param_schema["description"] = param_descriptions[param_name]

            properties[param_name] = param_schema

            # Check if parameter is required (no default value)
            if param.default is inspect.Parameter.empty:
                required.append(param_name)

        # Get first line of docstring as function description
        description = doc.split("\n")[0] if doc else f"Function {func.__name__}"

        return {
            "type": "function",
            "function": {
                "name": func.__name__,
                "description": description,
                "parameters": {
                    "type": "object",
                    "properties": properties,
                    "required": required,
                },
            },
        }

    def _parse_docstring_params(self, docstring: str) -> dict[str, str]:
        """Parses parameter descriptions from a docstring.

        Args:
            docstring: The docstring to parse.

        Returns:
            dict[str, str]: A dictionary mapping parameter names to descriptions.
        """
        param_descriptions: dict[str, str] = {}
        lines = docstring.split("\n")
        in_args_section = False

        for line in lines:
            stripped = line.strip()
            if stripped.lower().startswith("args:"):
                in_args_section = True
                continue
            elif stripped.lower().startswith(("returns:", "raises:", "example:")):
                in_args_section = False
                continue

            if in_args_section and ":" in stripped:
                # Parse "param_name: description" or "param_name (type): description"
                parts = stripped.split(":", 1)
                if len(parts) == 2:
                    param_part = parts[0].strip()
                    desc_part = parts[1].strip()
                    # Handle "param_name (type)" format
                    if "(" in param_part:
                        param_name = param_part.split("(")[0].strip()
                    else:
                        param_name = param_part
                    param_descriptions[param_name] = desc_part

        return param_descriptions

    def _python_type_to_json_schema(self, python_type: Any) -> dict[str, Any]:
        """Converts a Python type annotation to a JSON schema type.

        Args:
            python_type: The Python type annotation to convert.

        Returns:
            dict[str, Any]: The JSON schema type definition.
        """
        if python_type is inspect.Parameter.empty:
            return {"type": "string"}

        origin = get_origin(python_type)
        args = get_args(python_type)

        # Handle Union types (including Optional which is Union[X, None])
        if origin is Union or origin is types_module.UnionType:
            # Filter out NoneType
            non_none_args = [a for a in args if a is not type(None)]
            if len(non_none_args) == 1:
                return self._python_type_to_json_schema(non_none_args[0])
            # Multiple types - use anyOf
            return {"anyOf": [self._python_type_to_json_schema(a) for a in non_none_args]}

        # Handle list/List types
        if origin is list:
            if args:
                return {
                    "type": "array",
                    "items": self._python_type_to_json_schema(args[0]),
                }
            return {"type": "array"}

        # Handle dict/Dict types
        if origin is dict:
            return {"type": "object"}

        # Handle basic types
        type_mapping: dict[type, dict[str, str]] = {
            str: {"type": "string"},
            int: {"type": "integer"},
            float: {"type": "number"},
            bool: {"type": "boolean"},
            type(None): {"type": "null"},
        }

        if python_type in type_mapping:
            return type_mapping[python_type]

        # Default to string for unknown types
        return {"type": "string"}
