from typing import List, Dict, Any, Union, Tuple, Optional, Callable
import json
import re
from opto.utils.llm import LLM, AbstractModel
from opto.trainer.suggest import Suggest

def exact_match_metric(question, student_answer, info):
    """ Exact match metric """
    return float(student_answer == info)

class AutoGuide:
    """
    Base class for all guides that provide feedback on content.

    Guide evaluates generated content and provide feedback to help improve it.
    Different implementations may use different strategies for evaluation,
    such as LLM-based comparison, keyword matching, or custom verification.
    """

    def __call__(self, task: str, response: str, info: Any, **kwargs) -> Tuple[float, str]:
        """
        Generate feedback for the provided response.

        Args:
            task: The task to analyze (e.g., user query, task, etc.)
            response: The response generated by LLM (e.g., student answer, generated code)
            info: additional information (used by some suggest but not others)
            **kwargs: Optional reference information (e.g., expected answer, execution logs),
                     Additional context or parameters for specialized guide implementations

        Returns:
            score: score from the teacher
            feedback: feedback from the teacher
        """
        return self.forward(task, response, info, **kwargs)

    def forward(self, task: str, response: str, info: Any, **kwargs) -> Tuple[float, str]:
        return self.get_feedback(task, response, info, **kwargs)
    
    def get_feedback(self, query: str, response: str, reference: Optional[str] = None, **kwargs) -> Tuple[float, str]:
        raise NotImplementedError

    def metric(self, query: str, response: str, reference: Optional[str] = None, **kwargs) -> float:
        """ Exact match metric """
        return self.get_feedback(query, response, reference)[0]


class VerbalJudgeGuide(AutoGuide):
    """
    This is a combined metric + feedback guide that asks LLM to provide a binary judgment (True/False)
    and then if False, provide feedback.

    This is an implementation of LLM-as-a-judge.
    """

    DEFAULT_PROMPT_TEMPLATE = (
        "The query is: {query}. The student answered: {response}. The correct answer is: {reference}. "
        "If the student answer is correct, please say 'Correct [TERMINATE]'. "
        "Otherwise, if the student answer is incorrect, please provide feedback to the student. "
        "The feedback should be specific and actionable."
    )

    DEFAULT_SYSTEM_PROMPT = "You're a helpful teacher who provides clear and constructive feedback."
    DEFAULT_CORRECTNESS_TEMPLATE = "Correct [TERMINATE]"

    def __init__(self,
                 model: Optional[str] = None,
                 llm: Optional[AbstractModel] = None,
                 prompt_template: Optional[str] = None,
                 system_prompt: Optional[str] = None,
                 correctness_template: Optional[str] = None):
        """
        Initialize the VerbalGuide with an LLM and prompt templates.

        Args:
            model: The name of the LLM model to use (if llm is not provided)
            llm: An instance of AbstractModel to use for generating feedback
            prompt_template: Custom prompt template with {response} and {reference} placeholders
            system_prompt: Custom system prompt for the LLM
            correctness_template: Template to use when response is deemed correct by metric
        """
        self.model = model
        self.llm = llm or LLM(model=model)
        self.prompt_template = prompt_template or self.DEFAULT_PROMPT_TEMPLATE
        self.system_prompt = system_prompt or self.DEFAULT_SYSTEM_PROMPT
        self.correctness_template = correctness_template or self.DEFAULT_CORRECTNESS_TEMPLATE

    def get_feedback(self, query: str, response: str, reference: Optional[str] = None, **kwargs) -> Tuple[float, str]:
        """
        Get LLM-generated feedback by comparing response with reference information.

        Args:
            query: The query to analyze (e.g., user query, task, etc.)
            response: The response generated by LLM (e.g., student answer, code, etc.)
            reference: The expected information or correct answer
            **kwargs: Additional parameters (unused in this implementation)

        Returns:
            score: a float number provided by the metric function
            feedback: A string containing the LLM-generated feedback
        """
        if reference is None:
            raise ValueError("ReferenceGuide requires reference information to generate feedback")

        # Check if metric function indicates perfect match
        user_prompt = self.prompt_template.format(query=query, response=response, reference=reference)

        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": user_prompt}
        ]

        llm_response = self.llm(messages=messages)

        # Extract the content from the response based on the LLM's response format
        if hasattr(llm_response, 'choices') and hasattr(llm_response.choices[0], 'message'):
            # For OpenAI-like response format
            llm_response = llm_response.choices[0].message.content
            # else: response is already the content

        # Format the output
        deliminator = "\n-------------------------------------\n"
        formatted_response = (
                f"The query is: {query}. The student answered: {response}. The correct answer is: {reference}. " + deliminator +
                "Expert feedback for generating the net content (if exists, please pay most attention to it):\n" +
                llm_response
        )

        score = 1 if 'Correct [TERMINATE]' in llm_response else 0

        return score, formatted_response

    def forward(self, task: str, response: str, info: Any, **kwargs) -> Tuple[float, str]:
        score, feedback = self.get_feedback(task, response, info, **kwargs)
        return score, feedback
