from __future__ import annotations
from typing import Union, Literal, Optional, List, Any

from jinja2 import Template
from pydantic import BaseModel, Field

from .question_base import QuestionBase
from .descriptors import QuestionOptionsDescriptor
from .decorators import inject_exception
from .response_validator_abc import ResponseValidatorABC
from .question_multiple_choice import (
    BaseMultipleChoiceResponse,
    MultipleChoiceResponseValidator,
)
from pydantic import BaseModel, Field


# Create a custom response model for MultipleChoiceWithOther
class BaseMultipleChoiceWithOtherResponse(BaseModel):
    """
    Base model for multiple choice with other responses.

    Attributes:
        answer: The selected choice
        other_text: Text for the "Other" option when selected
        comment: Optional comment field
        generated_tokens: Optional raw tokens generated by the model
    """

    answer: Any = Field(..., description="Selected choice")
    other_text: Optional[str] = Field(
        None, description="Specification when 'Other' is selected"
    )
    comment: Optional[str] = Field(None, description="Optional comment field")
    generated_tokens: Optional[Any] = Field(None, description="Generated tokens")


def create_response_model_with_other(choices: List[str], permissive: bool = False):
    """
    Create a ChoiceWithOtherResponse model class with a predefined list of choices.

    :param choices: A list of allowed values for the answer field.
    :param permissive: If True, any value will be accepted as an answer.
    :return: A new Pydantic model class.
    """
    choice_tuple = tuple(choices)

    if not permissive:

        class ChoiceWithOtherResponse(BaseMultipleChoiceWithOtherResponse):
            """
            A model for multiple choice with other responses with strict validation.
            """

            answer: Literal[choice_tuple] = Field(description="Selected choice")

            model_config = {
                "json_schema_extra": {"properties": {"answer": {"enum": choices}}}
            }

    else:

        class ChoiceWithOtherResponse(BaseMultipleChoiceWithOtherResponse):
            """
            A model for multiple choice with other responses with permissive validation.
            """

            answer: Any = Field(
                description=f"Selected choice (can be any value). Suggested choices are: {choices}"
            )

            model_config = {"title": "PermissiveChoiceWithOtherResponse"}

    return ChoiceWithOtherResponse


class MultipleChoiceWithOtherResponseValidator(MultipleChoiceResponseValidator):
    """
    Validator for multiple choice with "other" responses.

    This validator extends the MultipleChoiceResponseValidator to handle
    the case where a user selects "Other" and provides a custom response.

    Examples:
        >>> from edsl.questions import QuestionMultipleChoiceWithOther
        >>> q = QuestionMultipleChoiceWithOther(
        ...     question_name="feeling",
        ...     question_text="How are you feeling?",
        ...     question_options=["Good", "Bad", "Neutral"],
        ...     other_option_text="Other"
        ... )
        >>> validator = q.response_validator
        >>> result = validator.validate({"answer": "Good"})
        >>> sorted(result.keys())
        ['answer', 'comment', 'generated_tokens', 'other_text']
        >>> result["answer"]
        'Good'
        >>> result = validator.validate({"answer": "Other", "other_text": "Excited"})
        >>> sorted(result.keys())
        ['answer', 'comment', 'generated_tokens', 'other_text']
        >>> result["answer"]
        'Other'
        >>> result["other_text"]
        'Excited'
        >>> # Direct "Other: X" format
        >>> result = validator.validate({"answer": "Other: Paris"})
        >>> result["answer"]
        'Other'
        >>> result["other_text"]
        'Paris'
    """

    required_params = ["question_options", "use_code", "other_option_text"]

    def __init__(self, **kwargs):
        """
        Initialize the validator.

        Ensures that "Other" is added to the question_options.
        """
        super().__init__(**kwargs)

        # Make sure "Other" is always in the list of valid options
        if "Other" not in self.question_options:
            # Create a new list but don't modify the original reference
            self.question_options = list(self.question_options) + ["Other"]

    def validate(self, response_dict, verbose=False):
        """
        Validate the response according to the schema.

        This overrides the parent validate method to handle the "Other" option specially.
        If the answer is in the format "Other: X", it splits it into "Other" and the other_text.

        Parameters:
            response_dict: The response to validate
            verbose: Whether to print debug information

        Returns:
            A validated response dict with "Other" as the answer and other_text field set
        """
        # Keep original for reference
        orig_response_dict = response_dict.copy()

        # Create a copy to avoid modifying the original that may be needed elsewhere
        response_dict = response_dict.copy()
        answer = str(response_dict.get("answer", ""))
        other_text = None

        # Check for "Other: X" format directly in the answer field
        if ":" in answer:
            parts = answer.split(":", 1)
            if len(parts) == 2 and parts[0].strip().lower() == "other":
                # Extract the "Other" part and the custom text
                other_text = parts[1].strip()
                if verbose:
                    print(f"Detected 'Other: X' format: {answer}")
                    print(f"Extracted other_text: {other_text}")

                # Change answer to just "Other" for validation
                response_dict["answer"] = "Other"

                # Add the other_text to the response
                if "other_text" not in response_dict or not response_dict["other_text"]:
                    response_dict["other_text"] = other_text

        # Make sure "Other" is in the list of valid options for validation
        question_options = list(self.question_options)
        if "Other" not in question_options:
            question_options.append("Other")

        # Validate with the parent validator
        validated_response = super().validate(response_dict, verbose)

        # If we extracted other_text but it wasn't set in the validated response, set it now
        if other_text and (
            "other_text" not in validated_response
            or not validated_response["other_text"]
        ):
            validated_response["other_text"] = other_text

        return validated_response

    def fix(self, response, verbose=False):
        """
        Attempt to fix an invalid multiple choice with other response.

        Extends the MultipleChoiceResponseValidator fix method to also
        handle "Other" responses with custom text.

        Parameters:
            response: The invalid response to fix
            verbose: Whether to print debug information

        Returns:
            A fixed response dict if possible, otherwise the original response
        """
        # Check if this is an "Other" response with additional text
        response_text = str(response.get("answer", ""))

        # Handle "Other: X" format - extract "Other" and other_text separately
        if ":" in response_text:
            parts = response_text.split(":", 1)
            if len(parts) == 2 and parts[0].strip().lower() == "other":
                if verbose:
                    print(f"Identified as 'Other: X' format: {response_text}")

                # Extract the custom text after "Other:"
                other_text = parts[1].strip()

                # Create response with "Other" and the other_text separately
                proposed_data = {
                    "answer": "Other",
                    "other_text": other_text,
                    "comment": response.get("comment"),
                    "generated_tokens": response.get("generated_tokens"),
                }

                try:
                    # Validate the fixed answer
                    self.response_model.model_validate(proposed_data)
                    if verbose:
                        print(f"Fixed answer: Other with other_text={other_text}")
                    return proposed_data
                except Exception as e:
                    if verbose:
                        print(f"Validation failed for 'Other: X' format: {e}")

        # If the response is not in question_options but contains "Other",
        # it might be an "Other" response with custom text in a different format
        elif (
            response_text not in self.question_options
            and "other" in response_text.lower()
        ):

            # Try to extract the custom text after "Other"
            # This is a fallback for other formats
            after_other = response_text.lower().split("other", 1)
            if len(after_other) > 1:
                other_text = after_other[1].strip()
                if other_text.startswith(":"):
                    other_text = other_text[1:].strip()

                if verbose:
                    print(f"Extracted text after 'Other': {other_text}")

                # Create response with "Other" and the other_text separately
                proposed_data = {
                    "answer": "Other",
                    "other_text": other_text,
                    "comment": response.get("comment"),
                    "generated_tokens": response.get("generated_tokens"),
                }

                try:
                    # Validate the fixed answer
                    self.response_model.model_validate(proposed_data)
                    if verbose:
                        print(f"Fixed answer: Other with other_text={other_text}")
                    return proposed_data
                except Exception as e:
                    if verbose:
                        print(f"Validation failed for extracted 'Other' text: {e}")

        # If not an "Other" response or validation failed, try the parent class fix method
        return super().fix(response, verbose)

    valid_examples = [
        (
            {"answer": "Good"},
            {
                "question_options": ["Good", "Great", "OK", "Bad"],
                "other_option_text": "Other",
            },
        ),
        (
            {"answer": "Other", "other_text": "Fantastic"},
            {
                "question_options": ["Good", "Great", "OK", "Bad"],
                "other_option_text": "Other",
            },
        ),
    ]

    invalid_examples = [
        (
            {"answer": "Terrible"},
            {
                "question_options": ["Good", "Great", "OK", "Bad"],
                "other_option_text": "Other",
            },
            "Value error, Permitted values are 'Good', 'Great', 'OK', 'Bad', 'Other'",
        ),
        (
            {"answer": "Other", "other_text": ""},
            {
                "question_options": ["Good", "Great", "OK", "Bad"],
                "other_option_text": "Other",
            },
            "When selecting 'Other', you must provide text in the other_text field",
        ),
    ]


class QuestionMultipleChoiceWithOther(QuestionBase):
    """
    A question that prompts the agent to select one option from a list of choices or specify "Other".

    QuestionMultipleChoiceWithOther extends QuestionMultipleChoice to include an "Other" option
    that allows the agent to provide a custom response when none of the predefined options
    are suitable. This is especially useful for surveys and open-ended questions where
    you want to capture responses that don't fit into predefined categories.

    Key Features:
    - All features of QuestionMultipleChoice
    - Additional "Other" option with free-text field
    - Customizable text for the "Other" option
    - Validation ensures that when "Other" is selected, a text explanation is provided
    - Supports "Other: X" format where X is the custom response

    Technical Details:
    - Uses extended Pydantic models for validation
    - Preserves all functionality of QuestionMultipleChoice
    - Adds 'other_text' field for custom responses
    - Post-processes responses to handle "Other: X" format

    Examples:
        Basic usage:

        ```python
        q = QuestionMultipleChoiceWithOther(
            question_name="preference",
            question_text="Which color do you prefer?",
            question_options=["Red", "Green", "Blue", "Yellow"]
        )
        ```

        Custom "Other" option text:

        ```python
        q = QuestionMultipleChoiceWithOther(
            question_name="preference",
            question_text="Which color do you prefer?",
            question_options=["Red", "Green", "Blue", "Yellow"],
            other_option_text="Something else (please specify)"
        )
        ```

        Handling "Other: X" format:

        ```python
        # If the model responds with "Other: Paris"
        q = QuestionMultipleChoiceWithOther(
            question_name="capital",
            question_text="What is the capital of France?",
            question_options=["London", "Berlin", "Madrid"]
        )
        result = q.by(model).run()
        # result will have:
        # answer.capital = "Other"
        # other_text.capital_other_text = "Paris"
        ```
    """

    question_type = "multiple_choice_with_other"
    purpose = "When options are known but you want to allow for custom responses"
    question_options: Union[list[str], list[list], list[float], list[int]] = (
        QuestionOptionsDescriptor()
    )
    _response_model = None
    response_validator_class = MultipleChoiceWithOtherResponseValidator

    def post_process_result(self, result):
        """
        Post-process the result to handle "Other: X" format.

        This method is called after the result is generated by the model
        and before it's returned to the user. It checks if the answer
        has the "Other: X" format and stores the full answer.

        Parameters:
            result: The result object to process

        Returns:
            The processed result object
        """
        # Process each result in the results list
        for r in result:
            if "answer" in r:
                question_name = self.question_name
                answer = r["answer"].get(question_name)

                # Check if the answer has the "Other: X" format or custom "Something else: X" format
                if isinstance(answer, str) and ":" in answer:
                    # Split into two parts
                    parts = answer.split(":", 1)
                    if len(parts) == 2:
                        # Check if the first part is either "Other" or matches other_option_text
                        prefix = parts[0].strip()
                        if prefix == "Other" or prefix == self.other_option_text:
                            # Extract the custom text
                            other_text = parts[1].strip()

                            # Set the answer to just "Other" (not the full "Other: X" format)
                            r["answer"][question_name] = "Other"

                            # Store the other_text in the result object as well as an attribute
                            other_text_key = f"{question_name}_other_text"
                            if "other_text" not in r:
                                r["other_text"] = {}
                            r["other_text"][other_text_key] = other_text

                            # Store the other_text as an attribute of the question instance
                            self._other_text = other_text

        return result

    @property
    def other_text(self):
        """
        Get the text entered for the 'Other' option.

        Returns:
            The text entered for the 'Other' option, or None if not applicable
        """
        return getattr(self, "_other_text", None)

    def __init__(
        self,
        question_name: str,
        question_text: str,
        question_options: Union[list[str], list[list], list[float], list[int]],
        include_comment: bool = True,
        use_code: bool = False,
        answering_instructions: Optional[str] = None,
        question_presentation: Optional[str] = None,
        permissive: bool = False,
        other_option_text: str = "Other",
        other_instructions: Optional[str] = None,
    ):
        """
        Initialize a new multiple choice with "Other" question.

        Parameters
        ----------
        question_name : str
            The name of the question, used as an identifier. Must be a valid Python variable name.

        question_text : str
            The actual text of the question to be asked.

        question_options : Union[list[str], list[list], list[float], list[int]]
            The list of options the agent can select from. The "Other" option will be
            automatically added to this list.

        include_comment : bool, default=True
            Whether to include a comment field in the response.

        use_code : bool, default=False
            If True, the answer will be the index of the selected option (0-based) instead of
            the option text itself.

        answering_instructions : Optional[str], default=None
            Custom instructions for how the model should answer the question.

        question_presentation : Optional[str], default=None
            Custom template for how the question is presented to the model.

        permissive : bool, default=False
            If True, the validator will accept answers that are not in the provided options list.

        other_option_text : str, default="Other"
            The text to use for the "Other" option. This will be added to the list of options.

        other_instructions : Optional[str], default=None
            Custom instructions for how to provide the "Other" response. If None,
            default instructions will be used.

        Examples
        --------
        >>> q = QuestionMultipleChoiceWithOther(
        ...     question_name="color_preference",
        ...     question_text="What is your favorite color?",
        ...     question_options=["Red", "Blue", "Green", "Yellow"],
        ...     other_option_text="Something else (please specify)"
        ... )
        """
        # Initialize base attributes from QuestionBase
        self.question_name = question_name
        self.question_text = question_text
        self.question_options = question_options

        # Add other specific attributes
        self._include_comment = include_comment
        self.use_code = use_code
        self.answering_instructions = answering_instructions
        self.question_presentation = question_presentation
        self.permissive = permissive
        self.other_option_text = other_option_text
        self.other_instructions = other_instructions

    def create_response_model(self, replacement_dict: dict = None):
        """Create a response model that allows for the 'Other' option."""
        if replacement_dict is None:
            replacement_dict = {}

        # Create options list with "Other" option added
        # Always use "Other" (not self.other_option_text) to ensure consistency
        options = list(self.question_options)
        if "Other" not in options:
            options.append("Other")

        if self.use_code:
            return create_response_model_with_other(
                list(range(len(options))), self.permissive
            )
        else:
            return create_response_model_with_other(options, self.permissive)

    @property
    def question_html_content(self) -> str:
        """Return the HTML version of the question with the Other option."""
        if hasattr(self, "option_labels"):
            option_labels = self.option_labels
        else:
            option_labels = {}

        # Create a list of all options including the "Other" option
        all_options = list(self.question_options) + [self.other_option_text]

        question_html_content = Template(
            """
        {% for option in question_options %} 
        <div>
        <input type="radio" id="{{ option }}" name="{{ question_name }}" value="{{ option }}">
        <label for="{{ option }}">
        {{ option }}
        {% if option in option_labels %}
        : {{ option_labels[option] }}
        {% endif %}
        </label>
        </div>
        {% endfor %}
        
        <div>
        <input type="radio" id="{{ other_option }}" name="{{ question_name }}" value="{{ other_option }}">
        <label for="{{ other_option }}">{{ other_option }}</label>
        <input type="text" id="{{ question_name }}_other_text" name="{{ question_name }}_other_text" 
               placeholder="Please specify" style="display:none;">
        </div>
        
        <script>
        document.getElementById('{{ other_option }}').addEventListener('change', function() {
            document.getElementById('{{ question_name }}_other_text').style.display = 'inline-block';
        });
        
        // Hide the text input when any other option is selected
        {% for option in question_options %}
        document.getElementById('{{ option }}').addEventListener('change', function() {
            document.getElementById('{{ question_name }}_other_text').style.display = 'none';
        });
        {% endfor %}
        </script>
        """
        ).render(
            question_name=self.question_name,
            question_options=self.question_options,
            option_labels=option_labels,
            other_option=self.other_option_text,
        )
        return question_html_content

    def by(self, *models, **kwargs):
        """
        Chain this question with one or more models.

        This method overrides the parent class's by method to add post-processing
        for the 'Other: X' format in results.

        Parameters:
            *models: One or more models to chain with this question
            **kwargs: Additional kwargs to pass to the parent class

        Returns:
            The chained object with post-processing added
        """
        # Call the parent class's by method first
        chained = super().by(*models, **kwargs)

        # Add a hook to post-process the results
        original_run = chained.run

        def run_with_post_processing(*args, **kwargs):
            # Call the original run method
            results = original_run(*args, **kwargs)

            # Post-process the results to handle 'Other: X' format
            return self.post_process_result(results)

        # Replace the run method with our wrapped version
        chained.run = run_with_post_processing

        return chained

    @classmethod
    @inject_exception
    def example(
        cls, include_comment=False, use_code=False
    ) -> QuestionMultipleChoiceWithOther:
        """Return an example instance."""
        return cls(
            question_text="How are you?",
            question_options=["Good", "Great", "OK", "Bad"],
            question_name="how_feeling_with_other",
            include_comment=include_comment,
            use_code=use_code,
            other_option_text="Other (please specify)",
        )


if __name__ == "__main__":
    import doctest

    doctest.testmod(optionflags=doctest.ELLIPSIS)
