"""Language model interface and base implementation for EDSL.

This module contains the LanguageModel abstract base class, which defines the interface
for all language models in the EDSL framework. The LanguageModel class handles common
functionality like caching, response parsing, token usage tracking, and cost calculation,
while specific model implementations (like GPT, Claude, etc.) implement the actual API calls.

Key concepts and terminology:

- raw_response: The complete JSON response returned directly from the model API.
  Contains all model metadata and response information.

- edsl_augmented_response: The raw model response augmented with EDSL-specific
  information, such as cache keys, token usage statistics, and cost data.

- generated_tokens: The actual text output generated by the model in response
  to the prompt. This is the content displayed to the user.

- edsl_answer_dict: The standardized, parsed response from the model in the format
  either {'answer': ...} or {'answer': ..., 'comment': ...} that EDSL uses internally.

The LanguageModel class handles these different representations and provides a
consistent interface regardless of which model provider is being used.
"""

from __future__ import annotations
from functools import wraps
import asyncio
import json
import os
import warnings
from abc import ABC, abstractmethod

from typing import (
    Any,
    Union,
    List,
    Optional,
    TYPE_CHECKING,
)

from .exceptions import LanguageModelValueError

from ..data_transfer_models import (
    ModelResponse,
    ModelInputs,
    EDSLOutput,
    AgentResponseDict,
)

if TYPE_CHECKING:
    from .price_manager import ResponseCost
    from ..caching import Cache
    from ..scenarios import FileStore
    from ..questions import QuestionBase
    from ..key_management import KeyLookup


from ..utilities import sync_wrapper, jupyter_nb_handler, remove_edsl_version, dict_hash
from ..base import PersistenceMixin, RepresentationMixin, HashingMixin
from ..key_management import KeyLookupCollection

from .registry import RegisterLanguageModelsMeta
from .raw_response_handler import RawResponseHandler


def handle_key_error(func):
    """Decorator to catch and provide user-friendly error messages for KeyError exceptions.

    This decorator gracefully handles KeyError exceptions that may occur when parsing
    model responses with unexpected structures, providing a clear error message to
    help users understand what went wrong.

    Args:
        func: The function to decorate

    Returns:
        Decorated function that catches KeyError exceptions
    """

    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
            # Unreachable code
        except KeyError as e:
            return f"""KeyError occurred: {e}. This is most likely because the model you are using 
            returned a JSON object we were not expecting."""

    return wrapper


class classproperty:
    """Descriptor that combines @classmethod and @property behaviors.

    This descriptor allows defining properties that work on the class itself
    rather than on instances, making it possible to have computed attributes
    at the class level.

    Usage:
        class MyClass:
            @classproperty
            def my_prop(cls):
                return cls.__name__
    """

    def __init__(self, method):
        """Initialize with the decorated method.

        Args:
            method: The class method to be accessed as a property
        """
        self.method = method

    def __get__(self, instance, cls):
        """Return the result of calling the method on the class.

        Args:
            instance: The instance (if called on an instance)
            cls: The class (always provided)

        Returns:
            The result of calling the method with the class as argument
        """
        return self.method(cls)


class LanguageModel(
    PersistenceMixin,
    RepresentationMixin,
    HashingMixin,
    ABC,
    metaclass=RegisterLanguageModelsMeta,
):
    """Abstract base class for all language model implementations in EDSL.

    This class defines the common interface and functionality for interacting with
    various language model providers (OpenAI, Anthropic, etc.). It handles caching,
    response parsing, token usage tracking, and cost calculation, providing a
    consistent interface regardless of the underlying model.

    Subclasses must implement the async_execute_model_call method to handle the
    actual API call to the model provider. Other methods may also be overridden
    to customize behavior for specific models.

    The class uses several mixins to provide serialization, pretty printing, and
    hashing functionality, and a metaclass to automatically register model implementations.

    Attributes:
        _model_: The default model identifier (set by subclasses)
        key_sequence: Path to extract generated text from model responses
        DEFAULT_RPM: Default requests per minute rate limit
        DEFAULT_TPM: Default tokens per minute rate limit
    """

    _model_ = None
    key_sequence = (
        None  # This should be something like ["choices", 0, "message", "content"]
    )

    DEFAULT_RPM = 100
    DEFAULT_TPM = 1000

    @classproperty
    def response_handler(cls):
        """Get a handler for processing raw model responses.

        This property creates a RawResponseHandler configured for the specific
        model implementation, using the class's key_sequence and usage_sequence
        attributes to know how to extract information from the model's response format.

        Returns:
            RawResponseHandler: Handler configured for this model type
        """
        key_sequence = cls.key_sequence
        usage_sequence = cls.usage_sequence if hasattr(cls, "usage_sequence") else None
        return RawResponseHandler(key_sequence, usage_sequence)

    def __init__(
        self,
        tpm: Optional[float] = None,
        rpm: Optional[float] = None,
        omit_system_prompt_if_empty_string: bool = True,
        key_lookup: Optional["KeyLookup"] = None,
        **kwargs,
    ):
        """Initialize a new language model instance.

        Args:
            tpm: Optional tokens per minute rate limit override
            rpm: Optional requests per minute rate limit override
            omit_system_prompt_if_empty_string: Whether to omit the system prompt when empty
            key_lookup: Optional custom key lookup for API credentials
            **kwargs: Additional parameters to pass to the model provider

        The initialization process:
        1. Sets up the model identifier from the class attribute
        2. Configures model parameters by merging defaults with provided values
        3. Sets up API key lookup and rate limits
        4. Applies all parameters as instance attributes

        For subclasses that define _parameters_ class attribute, these will be
        used as default parameters that can be overridden by kwargs.
        """
        # Get the model identifier from the class attribute
        self.model = getattr(self, "_model_", None)

        # Set up model parameters by combining defaults with provided values
        default_parameters = getattr(self, "_parameters_", None)
        parameters = self._overide_default_parameters(kwargs, default_parameters)
        self.parameters = parameters

        # Initialize basic settings
        self.remote = False
        self.omit_system_prompt_if_empty = omit_system_prompt_if_empty_string

        # Set up API key lookup and fetch model information
        self.key_lookup = self._set_key_lookup(key_lookup)
        self.model_info = self.key_lookup.get(self._inference_service_)

        # Apply rate limit overrides if provided
        if rpm is not None:
            self._rpm = rpm

        if tpm is not None:
            self._tpm = tpm

        # Apply all parameters as instance attributes
        for key, value in parameters.items():
            setattr(self, key, value)

        # Apply any additional kwargs that aren't part of the standard parameters
        for key, value in kwargs.items():
            if key not in parameters:
                setattr(self, key, value)

        # Handle API key check skip for testing
        if kwargs.get("skip_api_key_check", False):
            # Skip the API key check. Sometimes this is useful for testing.
            self._api_token = None

        # Add canned response to parameters
        if "canned_response" in kwargs:
            self.parameters["canned_response"] = kwargs["canned_response"]

    def _set_key_lookup(self, key_lookup: "KeyLookup") -> "KeyLookup":
        """Set up the API key lookup mechanism.

        This method either uses the provided key lookup object or creates a default
        one that looks for API keys in config files and environment variables.

        Args:
            key_lookup: Optional custom key lookup object

        Returns:
            KeyLookup: The key lookup object to use for API credentials
        """
        if key_lookup is not None:
            return key_lookup
        else:
            klc = KeyLookupCollection()
            klc.add_key_lookup(fetch_order=("config", "env"))
            return klc.get(("config", "env"))

    def set_key_lookup(self, key_lookup: "KeyLookup") -> None:
        """Update the key lookup mechanism after initialization.

        This method allows changing the API key lookup after the model has been
        created, clearing any cached API tokens.

        Args:
            key_lookup: The new key lookup object to use
        """
        if hasattr(self, "_api_token"):
            del self._api_token
        self.key_lookup = key_lookup

    def ask_question(self, question: "QuestionBase") -> str:
        """Ask a question using this language model and return the response.

        This is a convenience method that extracts the necessary prompts from a
        question object and makes a model call.

        Args:
            question: The EDSL question object to ask

        Returns:
            str: The model's response to the question
        """
        user_prompt = question.get_instructions().render(question.data).text
        system_prompt = "You are a helpful agent pretending to be a human."
        return self.execute_model_call(user_prompt, system_prompt)

    @property
    def rpm(self):
        """Get the requests per minute rate limit for this model.

        This property provides the rate limit either from an explicitly set value,
        from the model info in the key lookup, or from the default value.

        Returns:
            float: The requests per minute rate limit
        """
        if not hasattr(self, "_rpm"):
            if self.model_info is None:
                self._rpm = self.DEFAULT_RPM
            else:
                self._rpm = self.model_info.rpm
        return self._rpm

    @property
    def tpm(self):
        """Get the tokens per minute rate limit for this model.

        This property provides the rate limit either from an explicitly set value,
        from the model info in the key lookup, or from the default value.

        Returns:
            float: The tokens per minute rate limit
        """
        if not hasattr(self, "_tpm"):
            if self.model_info is None:
                self._tpm = self.DEFAULT_TPM
            else:
                self._tpm = self.model_info.tpm
        return self._tpm

    # Setters for rate limits
    @tpm.setter
    def tpm(self, value):
        """Set the tokens per minute rate limit.

        Args:
            value: The new tokens per minute limit
        """
        self._tpm = value

    @rpm.setter
    def rpm(self, value):
        """Set the requests per minute rate limit.

        Args:
            value: The new requests per minute limit
        """
        self._rpm = value

    @property
    def api_token(self) -> str:
        """Get the API token for this model's service.

        This property lazily fetches the API token from the key lookup
        mechanism when first accessed, caching it for subsequent uses.

        Returns:
            str: The API token for authenticating with the model provider

        Raises:
            ValueError: If no API key is found for this model's service
        """
        if not hasattr(self, "_api_token"):
            info = self.key_lookup.get(self._inference_service_, None)
            if info is None:
                raise LanguageModelValueError(
                    f"No key found for service '{self._inference_service_}'"
                )
            self._api_token = info.api_token
        return self._api_token

    def copy(self) -> "LanguageModel":
        """Create a deep copy of this language model instance.

        This method creates a completely independent copy of the language model
        by creating a new instance with the same parameters and copying relevant attributes.

        Returns:
            LanguageModel: A new language model instance that is functionally identical to this one

        Examples:
            >>> m1 = LanguageModel.example()
            >>> m2 = m1.copy()
            >>> m1 == m2  # Functionally equivalent
            True
            >>> id(m1) == id(m2)  # But different objects
            False
        """
        # Create a new instance of the same class with the same parameters
        try:
            # For most models, we can instantiate with the saved parameters
            new_model = self.__class__(**self.parameters)

            # Copy all important instance attributes
            for key, value in self.__dict__.items():
                if key not in ("_api_token",) and not key.startswith("__"):
                    setattr(new_model, key, value)

            return new_model
        except Exception:
            # Fallback for dynamically created classes like TestServiceLanguageModel
            from ..inference_services import default

            # If this is a test model, create a new test model instance
            if getattr(self, "_inference_service_", "") == "test":
                service = default.get_service("test")
                model_class = service.create_model("test")
                new_model = model_class(**self.parameters)

                # Copy attributes
                for key, value in self.__dict__.items():
                    if key not in ("_api_token",) and not key.startswith("__"):
                        setattr(new_model, key, value)

                return new_model

            # If we can't create the model directly, just return a simple test model
            # This is a last resort fallback
            from ..inference_services import get_service

            service = get_service("test")
            model_class = service.create_model("test")
            return model_class()

    def __getitem__(self, key):
        """Allow dictionary-style access to model attributes.

        Args:
            key: The attribute name to access

        Returns:
            The value of the specified attribute
        """
        return getattr(self, key)

    def hello(self, verbose=False):
        """Run a simple test to verify the model connection is working.

        This method makes a basic model call to check if the API credentials
        are valid and the model is responsive.

        Args:
            verbose: If True, prints the masked API token

        Returns:
            str: The model's response to a simple greeting
        """
        token = self.api_token
        masked = token[: min(8, len(token))] + "..."
        if verbose:
            print(f"Current key is {masked}")
        return self.execute_model_call(
            user_prompt="Hello, model!", system_prompt="You are a helpful agent."
        )

    def has_valid_api_key(self) -> bool:
        """Check if the model has a valid API key available.

        This method verifies if the necessary API key is available in
        environment variables or configuration for this model's service.
        Test models always return True.

        Returns:
            bool: True if a valid API key is available, False otherwise

        Examples:
            >>> LanguageModel.example().has_valid_api_key() : # doctest: +SKIP
            True
        """
        from ..enums import service_to_api_keyname

        if self._model_ == "test":
            return True

        key_name = service_to_api_keyname.get(self._inference_service_, "NOT FOUND")
        key_value = os.getenv(key_name)
        return key_value is not None

    def __hash__(self) -> int:
        """Generate a hash value based on model identity and parameters.

        This method allows language model instances to be used as dictionary
        keys or in sets by providing a stable hash value based on the
        model's essential characteristics.

        Returns:
            int: A hash value for the model instance

        Examples:
            >>> m = LanguageModel.example()
            >>> hash(m)  # Actual value may vary
            325654563661254408
        """
        return dict_hash(self.to_dict(add_edsl_version=False))

    def __eq__(self, other) -> bool:
        """Check if two language model instances are functionally equivalent.

        Two models are considered equal if they have the same model identifier
        and the same parameter settings, meaning they would produce the same
        outputs given the same inputs.

        Args:
            other: Another model to compare with

        Returns:
            bool: True if the models are functionally equivalent

        Examples:
            >>> m1 = LanguageModel.example()
            >>> m2 = LanguageModel.example()
            >>> m1 == m2
            True
        """
        return self.model == other.model and self.parameters == other.parameters

    @staticmethod
    def _overide_default_parameters(passed_parameter_dict, default_parameter_dict):
        """Merge default parameters with user-specified parameters.

        This method creates a parameter dictionary where explicitly passed
        parameters take precedence over default values, while ensuring all
        required parameters have a value.

        Args:
            passed_parameter_dict: Dictionary of user-specified parameters
            default_parameter_dict: Dictionary of default parameter values

        Returns:
            dict: Combined parameter dictionary with defaults and overrides

        Examples:
            >>> LanguageModel._overide_default_parameters(
            ...     passed_parameter_dict={"temperature": 0.5},
            ...     default_parameter_dict={"temperature": 0.9})
            {'temperature': 0.5}

            >>> LanguageModel._overide_default_parameters(
            ...     passed_parameter_dict={"temperature": 0.5},
            ...     default_parameter_dict={"temperature": 0.9, "max_tokens": 1000})
            {'temperature': 0.5, 'max_tokens': 1000}
        """
        # Handle the case when data is loaded from a dict after serialization
        if "parameters" in passed_parameter_dict:
            passed_parameter_dict = passed_parameter_dict["parameters"]

        # Create new dict with defaults, overridden by passed parameters
        return {
            parameter_name: passed_parameter_dict.get(parameter_name, default_value)
            for parameter_name, default_value in default_parameter_dict.items()
        }

    def __call__(self, user_prompt: str, system_prompt: str):
        """Allow the model to be called directly as a function.

        This method provides a convenient way to use the model by calling
        it directly with prompts, like `response = model(user_prompt, system_prompt)`.

        Args:
            user_prompt: The user message or input prompt
            system_prompt: The system message or context

        Returns:
            The response from the model
        """
        return self.execute_model_call(user_prompt, system_prompt)

    @abstractmethod
    async def async_execute_model_call(
        self, user_prompt: str, system_prompt: str, question_name: Optional[str] = None
    ):
        """Execute the model call asynchronously.

        This abstract method must be implemented by all model subclasses
        to handle the actual API call to the language model provider.

        Args:
            user_prompt: The user message or input prompt
            system_prompt: The system message or context
            question_name: Optional name of the question being asked (primarily used for test models)

        Returns:
            Coroutine that resolves to the model response

        Note:
            Implementations should handle the actual API communication,
            including authentication, request formatting, and response parsing.
        """
        pass

    async def remote_async_execute_model_call(
        self, user_prompt: str, system_prompt: str, question_name: Optional[str] = None
    ):
        """Execute the model call remotely through the EDSL Coop service.

        This method allows offloading the model call to a remote server,
        which can be useful for models not available in the local environment
        or to avoid rate limits.

        Args:
            user_prompt: The user message or input prompt
            system_prompt: The system message or context
            question_name: Optional name of the question being asked (primarily used for test models)

        Returns:
            Coroutine that resolves to the model response from the remote service
        """
        from ..coop import Coop

        client = Coop()
        response_data = await client.remote_async_execute_model_call(
            self.to_dict(), user_prompt, system_prompt
        )
        return response_data

    @jupyter_nb_handler
    def execute_model_call(self, *args, **kwargs):
        """Execute a model call synchronously.

        This method is a synchronous wrapper around the asynchronous execution,
        making it easier to use the model in non-async contexts. It's decorated
        with jupyter_nb_handler to ensure proper handling in notebook environments.

        Args:
            *args: Positional arguments to pass to async_execute_model_call
            **kwargs: Keyword arguments to pass to async_execute_model_call
                     Can include question_name for test models

        Returns:
            The model response
        """

        async def main():
            results = await asyncio.gather(
                self.async_execute_model_call(*args, **kwargs)
            )
            return results[0]  # Since there's only one task, return its result

        return main()

    @classmethod
    def get_generated_token_string(cls, raw_response: dict[str, Any]) -> str:
        """Extract the generated text from a raw model response.

        This method navigates the response structure using the model's key_sequence
        to find and return just the generated text, without metadata.

        Args:
            raw_response: The complete response dictionary from the model API

        Returns:
            str: The generated text string

        Examples:
            >>> m = LanguageModel.example(test_model=True)
            >>> raw_response = m.execute_model_call("Hello, model!", "You are a helpful agent.")
            >>> m.get_generated_token_string(raw_response)
            'Hello world'
        """
        return cls.response_handler.get_generated_token_string(raw_response)

    @classmethod
    def get_usage_dict(cls, raw_response: dict[str, Any]) -> dict[str, Any]:
        """Extract token usage statistics from a raw model response.

        This method navigates the response structure to find and return
        information about token usage, which is used for cost calculation
        and monitoring.

        Args:
            raw_response: The complete response dictionary from the model API

        Returns:
            dict: Dictionary of token usage statistics (input tokens, output tokens, etc.)
        """
        return cls.response_handler.get_usage_dict(raw_response)

    @classmethod
    def parse_response(cls, raw_response: dict[str, Any]) -> EDSLOutput:
        """Parse the raw API response into a standardized EDSL output format.

        This method processes the model's response to extract the generated content
        and format it according to EDSL's expected structure, making it consistent
        across different model providers.

        Args:
            raw_response: The complete response dictionary from the model API

        Returns:
            EDSLOutput: Standardized output structure with answer and optional comment
        """
        return cls.response_handler.parse_response(raw_response)

    async def _async_get_intended_model_call_outcome(
        self,
        user_prompt: str,
        system_prompt: str,
        cache: "Cache",
        iteration: int = 0,
        files_list: Optional[List["FileStore"]] = None,
        invigilator=None,
    ) -> ModelResponse:
        """Handle model calls with caching for efficiency.

        This method implements the caching logic for model calls, checking if a
        response is already cached before making an actual API call. It handles
        the complete workflow of:
        1. Creating a cache key from the prompts and parameters
        2. Checking if the response is already in the cache
        3. Making the API call if needed
        4. Storing new responses in the cache
        5. Adding metadata like cost and cache status

        Args:
            user_prompt: The user's message or input prompt
            system_prompt: The system's message or context
            cache: The cache object to use for storing/retrieving responses
            iteration: The iteration number, used for the cache key
            files_list: Optional list of files to include in the prompt
            invigilator: Optional invigilator object, not used in caching

        Returns:
            ModelResponse: Response object with the model output and metadata

        Examples:
            >>> from edsl import Cache
            >>> m = LanguageModel.example(test_model=True)
            >>> m._get_intended_model_call_outcome(user_prompt="Hello", system_prompt="hello", cache=Cache())
            ModelResponse(...)
        """
        # Add file hashes to the prompt if files are provided
        if files_list:
            files_hash = "+".join([str(hash(file)) for file in files_list])
            user_prompt_with_hashes = user_prompt + f" {files_hash}"
        else:
            user_prompt_with_hashes = user_prompt

        # Prepare parameters for cache lookup
        cache_parameters = self.parameters.copy()
        if self.model == "test":
            cache_parameters.pop("canned_response", None)
        cache_call_params = {
            "model": str(self.model),
            "parameters": cache_parameters,
            "system_prompt": system_prompt,
            "user_prompt": user_prompt_with_hashes,
            "iteration": iteration,
        }

        # Try to fetch from cache
        cached_response, cache_key = cache.fetch(**cache_call_params)
        if cache_used := cached_response is not None:
            # Cache hit - use the cached response
            response = json.loads(cached_response)
        else:
            # Cache miss - make a new API call
            # Determine whether to use remote or local execution
            f = (
                self.remote_async_execute_model_call
                if hasattr(self, "remote") and self.remote
                else self.async_execute_model_call
            )

            # Prepare parameters for the model call
            params = {
                "user_prompt": user_prompt,
                "system_prompt": system_prompt,
                "files_list": files_list,
            }
            # Add question_name parameter for test models
            if self.model == "test" and invigilator:
                params["question_name"] = invigilator.question.question_name
            # Get timeout from configuration
            from ..config import CONFIG
            import logging

            logger = logging.getLogger(__name__)
            base_timeout = float(CONFIG.get("EDSL_API_TIMEOUT"))

            # Adjust timeout if files are present
            import time

            start = time.time()
            if files_list:
                # Calculate total size of attached files in MB
                file_sizes = []
                for file in files_list:
                    # Try different attributes that might contain the file content
                    if hasattr(file, "base64_string") and file.base64_string:
                        file_sizes.append(len(file.base64_string) / (1024 * 1024))
                    elif hasattr(file, "content") and file.content:
                        file_sizes.append(len(file.content) / (1024 * 1024))
                    elif hasattr(file, "data") and file.data:
                        file_sizes.append(len(file.data) / (1024 * 1024))
                    else:
                        # Default minimum size if we can't determine actual size
                        file_sizes.append(1)  # Assume at least 1MB
                total_size_mb = sum(file_sizes)

                # Increase timeout proportionally to file size
                # For each MB of file size, add 10 seconds to the timeout (adjust as needed)
                size_adjustment = total_size_mb * 10

                # Cap the maximum timeout adjustment at 5 minutes (300 seconds)
                size_adjustment = min(size_adjustment, 300)

                TIMEOUT = base_timeout + size_adjustment

                logger.info(
                    f"Adjusted timeout for API call with {len(files_list)} files (total size: {total_size_mb:.2f}MB). Base timeout: {base_timeout}s, New timeout: {TIMEOUT}s"
                )
            else:
                TIMEOUT = base_timeout

            # Execute the model call with timeout
            response = await asyncio.wait_for(f(**params), timeout=TIMEOUT)
            # Store the response in the cache
            new_cache_key = cache.store(
                **cache_call_params, response=response, service=self._inference_service_
            )
            assert new_cache_key == cache_key  # Verify cache key integrity

        # Calculate cost for the response
        cost = self.cost(response)
        # Return a structured response with metadata
        response = ModelResponse(
            response=response,
            cache_used=cache_used,
            cache_key=cache_key,
            cached_response=cached_response,
            input_tokens=cost.input_tokens,
            output_tokens=cost.output_tokens,
            input_price_per_million_tokens=cost.input_price_per_million_tokens,
            output_price_per_million_tokens=cost.output_price_per_million_tokens,
            total_cost=cost.total_cost,
        )
        return response

    _get_intended_model_call_outcome = sync_wrapper(
        _async_get_intended_model_call_outcome
    )

    def simple_ask(
        self,
        question: QuestionBase,
        system_prompt="You are a helpful agent pretending to be a human.",
        top_logprobs=2,
    ):
        """Ask a simple question with log probability tracking.

        This is a convenience method for getting responses with log probabilities,
        which can be useful for analyzing model confidence and alternatives.

        Args:
            question: The EDSL question object to ask
            system_prompt: System message to use (default is human agent instruction)
            top_logprobs: Number of top alternative tokens to return probabilities for

        Returns:
            The model response, including log probabilities if supported
        """
        self.logprobs = True
        self.top_logprobs = top_logprobs
        return self.execute_model_call(
            user_prompt=question.human_readable(), system_prompt=system_prompt
        )

    async def async_get_response(
        self,
        user_prompt: str,
        system_prompt: str,
        cache: Cache,
        iteration: int = 1,
        files_list: Optional[List[FileStore]] = None,
        **kwargs,
    ) -> AgentResponseDict:
        """Get a complete response with all metadata and parsed format.

        This method handles the complete pipeline for:
        1. Making a model call (with caching)
        2. Parsing the response
        3. Constructing a full response object with inputs, outputs, and parsed data

        It's the primary method used by higher-level components to interact with models.

        Args:
            user_prompt: The user's message or input prompt
            system_prompt: The system's message or context
            cache: The cache object to use for storing/retrieving responses
            iteration: The iteration number (default: 1)
            files_list: Optional list of files to include in the prompt
            **kwargs: Additional parameters (invigilator can be provided here)

        Returns:
            AgentResponseDict: Complete response object with inputs, raw outputs, and parsed data
        """
        # Prepare parameters for the cached model call
        params = {
            "user_prompt": user_prompt,
            "system_prompt": system_prompt,
            "iteration": iteration,
            "cache": cache,
            "files_list": files_list,
        }

        # Add invigilator if provided
        if "invigilator" in kwargs:
            params.update({"invigilator": kwargs["invigilator"]})

        # Create structured input record
        model_inputs = ModelInputs(user_prompt=user_prompt, system_prompt=system_prompt)
        # Get model response (using cache if available)
        model_outputs: ModelResponse = (
            await self._async_get_intended_model_call_outcome(**params)
        )

        # Parse the response into EDSL's standard format
        edsl_dict: EDSLOutput = self.parse_response(model_outputs.response)

        # Combine everything into a complete response object
        agent_response_dict = AgentResponseDict(
            model_inputs=model_inputs,
            model_outputs=model_outputs,
            edsl_dict=edsl_dict,
        )
        return agent_response_dict

    get_response = sync_wrapper(async_get_response)

    def cost(self, raw_response: dict[str, Any]) -> ResponseCost:
        """Calculate the monetary cost of a model API call.

        This method extracts token usage information from the response and
        uses the price manager to calculate the actual cost in dollars based
        on the model's pricing structure and token counts.

        Args:
            raw_response: The complete response dictionary from the model API

        Returns:
            ResponseCost: Object containing token counts and total cost
        """
        # Extract token usage data from the response
        usage = self.get_usage_dict(raw_response)

        # Use the price manager to calculate the actual cost
        from .price_manager import PriceManager

        price_manager = PriceManager.get_instance()

        return price_manager.calculate_cost(
            inference_service=self._inference_service_,
            model=self.model,
            usage=usage,
            input_token_name=self.input_token_name,
            output_token_name=self.output_token_name,
        )

    def to_dict(self, add_edsl_version: bool = True) -> dict[str, Any]:
        """Serialize the model instance to a dictionary representation.

        This method creates a dictionary containing all the information needed
        to recreate this model, including its identifier, parameters, and service.
        Optionally includes EDSL version information for compatibility checking.

        Args:
            add_edsl_version: Whether to include EDSL version and class name (default: True)

        Returns:
            dict: Dictionary representation of this model instance

        Examples:
            >>> m = LanguageModel.example()
            >>> m.to_dict()
            {'model': '...', 'parameters': {'temperature': ..., 'max_tokens': ..., 'top_p': ..., 'frequency_penalty': ..., 'presence_penalty': ..., 'logprobs': False, 'top_logprobs': ...}, 'inference_service': 'openai', 'edsl_version': '...', 'edsl_class_name': 'LanguageModel'}
        """
        # Build the base dictionary with essential model information
        parameters = self.parameters.copy()

        # For test models, ensure canned_response is included in serialization
        if self.model == "test" and hasattr(self, "canned_response"):
            parameters["canned_response"] = self.canned_response

        d = {
            "model": self.model,
            "parameters": parameters,
            "inference_service": self._inference_service_,
        }

        # Add EDSL version and class information if requested
        if add_edsl_version:
            from edsl import __version__

            d["edsl_version"] = __version__
            d["edsl_class_name"] = self.__class__.__name__

        return d

    @classmethod
    @remove_edsl_version
    def from_dict(cls, data: dict) -> "LanguageModel":
        """Create a language model instance from a dictionary representation.

        This class method deserializes a model from its dictionary representation,
        finding the correct model class based on the model identifier and service.

        Args:
            data: Dictionary containing the model configuration

        Returns:
            LanguageModel: A new model instance of the appropriate type

        Note:
            This method does not use the stored inference_service directly but
            fetches the model class based on the model name and service name.
        """
        from .model import get_model_class

        # Determine the appropriate model class
        model_class = get_model_class(
            data["model"], service_name=data.get("inference_service", None)
        )

        # Handle canned_response in parameters for test models
        if (
            data["model"] == "test"
            and "parameters" in data
            and "canned_response" in data["parameters"]
        ):
            # Extract canned_response from parameters to set as a direct attribute
            canned_response = data["parameters"]["canned_response"]
            params_copy = data.copy()

            # Direct attribute will be set during initialization
            # Add it as a top-level parameter for model initialization
            if isinstance(params_copy, dict) and "parameters" in params_copy:
                params_copy["canned_response"] = canned_response

            # Create the instance with canned_response as a direct parameter
            return model_class(**params_copy)

        # For non-test models or test models without canned_response
        return model_class(**data)

    def __repr__(self) -> str:
        """Generate a string representation of the model.

        This representation includes the model identifier and all parameters,
        providing a clear picture of how the model is configured.

        Returns:
            str: A string representation of the model
        """
        # Format the parameters as a string
        param_string = ", ".join(
            f'{key} = """{value}"""' if key == "canned_response" else f"{key} = {value}"
            for key, value in self.parameters.items()
        )

        # Combine model name and parameters
        return (
            f"Model(model_name = '{self.model}', service_name = '{self._inference_service_}'"
            + (f", {param_string}" if param_string else "")
            + ")"
        )

    def __add__(self, other_model: "LanguageModel") -> "LanguageModel":
        """Define behavior when models are combined with the + operator.

        This operator is used in survey builder contexts, but note that it
        replaces the left model with the right model rather than combining them.

        Args:
            other_model: Another model to combine with this one

        Returns:
            LanguageModel: The other model if provided, otherwise this model

        Warning:
            This doesn't truly combine models - it replaces one with the other.
            For running multiple models, use a single 'by' call with multiple models.
        """
        warnings.warn(
            """Warning: one model is replacing another. If you want to run both models, use a single `by` e.g., 
              by(m1, m2, m3) not by(m1).by(m2).by(m3)."""
        )
        return other_model or self

    @classmethod
    def example(
        cls,
        test_model: bool = False,
        canned_response: str = "Hello world",
        throw_exception: bool = False,
    ) -> "LanguageModel":
        """Create an example language model instance for testing and demonstration.

        This method provides a convenient way to create a model instance for
        examples, tests, and documentation. It can create either a real model
        (with API key checking disabled) or a test model that returns predefined
        responses.

        Args:
            test_model: If True, creates a test model that doesn't make real API calls
            canned_response: For test models, the predefined response to return
            throw_exception: For test models, whether to throw an exception instead of responding

        Returns:
            LanguageModel: An example model instance

        Examples:
            Create a test model with a custom response:

            >>> from edsl.language_models import LanguageModel
            >>> m = LanguageModel.example(test_model=True, canned_response="WOWZA!")
            >>> isinstance(m, LanguageModel)
            True

            Use the test model to answer a question:

            >>> from edsl import QuestionFreeText
            >>> q = QuestionFreeText(question_text="What is your name?", question_name='example')
            >>> q.by(m).run(cache=False, disable_remote_cache=True, disable_remote_inference=True).select('example').first()
            'WOWZA!'

            Create a test model that throws exceptions:

            >>> m = LanguageModel.example(test_model=True, canned_response="WOWZA!", throw_exception=True) # doctest: +SKIP
            >>> r = q.by(m).run(cache=False, disable_remote_cache=True, disable_remote_inference=True, print_exceptions=True) # doctest: +SKIP
            Exception report saved to ...
        """
        from ..language_models import Model

        if test_model:
            # Create a test model with predefined behavior
            m = Model(
                "test", canned_response=canned_response, throw_exception=throw_exception
            )
            return m
        else:
            # Create a regular model with API key checking disabled
            return Model(skip_api_key_check=True)

    def from_cache(self, cache: "Cache") -> "LanguageModel":
        """Create a new model that only returns responses from the cache.

        This method creates a modified copy of the model that will only use
        cached responses, never making new API calls. This is useful for
        offline operation or repeatable experiments.

        Args:
            cache: The cache object containing previously cached responses

        Returns:
            LanguageModel: A new model instance that only reads from cache
        """
        from copy import deepcopy
        from types import MethodType
        from ..caching import Cache

        # Create a deep copy of this model instance
        new_instance = deepcopy(self)
        print("Cache entries", len(cache))

        # Filter the cache to only include entries for this model
        new_instance.cache = Cache(
            data={k: v for k, v in cache.items() if v.model == self.model}
        )
        print("Cache entries with same model", len(new_instance.cache))

        # Store prompt lists for reference
        new_instance.user_prompts = [
            ce.user_prompt for ce in new_instance.cache.values()
        ]
        new_instance.system_prompts = [
            ce.system_prompt for ce in new_instance.cache.values()
        ]

        # Define a new async_execute_model_call that only reads from cache
        async def async_execute_model_call(
            self,
            user_prompt: str,
            system_prompt: str,
            question_name: Optional[str] = None,
        ):
            """Only use cached responses, never making new API calls."""
            cache_call_params = {
                "model": str(self.model),
                "parameters": self.parameters,
                "system_prompt": system_prompt,
                "user_prompt": user_prompt,
                "iteration": 1,
            }
            cached_response, cache_key = cache.fetch(**cache_call_params)
            response = json.loads(cached_response)

            try:
                usage = self.get_usage_dict(response)
                input_tokens = int(usage[self.input_token_name])
                output_tokens = int(usage[self.output_token_name])
            except Exception as e:
                print(f"Could not fetch tokens from model response: {e}")
                input_tokens = None
                output_tokens = None
            return ModelResponse(
                response=response,
                cache_used=True,
                cache_key=cache_key,
                cached_response=cached_response,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                input_price_per_million_tokens=0,
                output_price_per_million_tokens=0,
                total_cost=0,
            )

        # Bind the new method to the copied instance
        setattr(
            new_instance,
            "async_execute_model_call",
            MethodType(async_execute_model_call, new_instance),
        )

        return new_instance


if __name__ == "__main__":
    """Run the module's test suite."""
    import doctest

    doctest.testmod(optionflags=doctest.ELLIPSIS)
