# 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 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

# 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:
            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

    def _estimate_conversation_tokens(self) -> int:
        """Rough estimate of current conversation size in tokens (chars / 4)."""
        total_chars = 0
        for msg in self.conversation:
            content = msg.get("content", "")
            if isinstance(content, str):
                total_chars += len(content)
            elif isinstance(content, list):
                for block in content:
                    if isinstance(block, dict):
                        for v in block.values():
                            if isinstance(v, str):
                                total_chars += len(v)
        return total_chars // 4

    def compact_conversation(self, max_context_tokens: int) -> None:
        """Truncate old tool results and text to keep conversation within context limit.

        Preserves the system message, initial user prompt, and the most recent
        messages while truncating content in older messages. Never modifies
        thinking blocks (they have cryptographic signatures in Anthropic's API).

        Args:
            max_context_tokens: The model's maximum context window in tokens.
        """
        estimated = self._estimate_conversation_tokens()
        if estimated < int(max_context_tokens * 0.7):
            return

        keep_start = 1
        for i, msg in enumerate(self.conversation):
            if msg.get("role") in ("user", "human"):
                keep_start = i + 1
                break

        keep_recent = min(6, len(self.conversation))
        if len(self.conversation) <= keep_start + keep_recent:
            return

        cutoff = len(self.conversation) - keep_recent
        trunc_marker = "\n...[truncated to save context]...\n"
        max_len = 300

        for msg in self.conversation[keep_start:cutoff]:
            content = msg.get("content")
            if msg.get("role") == "tool" and isinstance(content, str) and len(content) > max_len:
                msg["content"] = content[:150] + trunc_marker + content[-100:]
                continue
            if isinstance(content, str) and len(content) > max_len:
                msg["content"] = content[:150] + trunc_marker + content[-100:]
            elif isinstance(content, list):
                for block in content:
                    if not isinstance(block, dict):
                        continue
                    btype = block.get("type", "")
                    if btype == "thinking":
                        continue
                    if btype == "tool_result":
                        c = block.get("content", "")
                        if isinstance(c, str) and len(c) > max_len:
                            block["content"] = c[:150] + trunc_marker + c[-100:]
                    elif btype == "text":
                        t = block.get("text", "")
                        if len(t) > max_len:
                            block["text"] = t[:150] + trunc_marker + t[-100:]

    # =========================================================================
    # 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"}
