"""
Base class for follow-up question generators.

This extends the standard question generator interface to handle
follow-up research that includes previous research context.
"""

from abc import abstractmethod
from typing import Dict, List
from ..base_question import BaseQuestionGenerator


class BaseFollowUpQuestionGenerator(BaseQuestionGenerator):
    """
    Abstract base class for follow-up question generators.

    These generators create contextualized queries that incorporate
    previous research findings and context.
    """

    def __init__(self, model):
        """
        Initialize the follow-up question generator.

        Args:
            model: The language model to use for question generation
        """
        super().__init__(model)
        self.follow_up_context = {}

    def set_follow_up_context(self, context: Dict):
        """
        Set the follow-up research context.

        Args:
            context: Dictionary containing:
                - past_findings: Previous research findings
                - original_query: The original research query
                - follow_up_query: The follow-up question from user
                - past_sources: Sources from previous research
                - key_entities: Key entities identified
        """
        self.follow_up_context = context

    @abstractmethod
    def generate_contextualized_query(
        self,
        follow_up_query: str,
        original_query: str,
        past_findings: str,
        **kwargs,
    ) -> str:
        """
        Generate a contextualized query for follow-up research.

        Args:
            follow_up_query: The user's follow-up question
            original_query: The original research query
            past_findings: The findings from previous research
            **kwargs: Additional context parameters

        Returns:
            str: A contextualized query that includes previous context
        """
        pass

    def generate_questions(
        self,
        current_knowledge: str,
        query: str,
        questions_per_iteration: int,
        questions_by_iteration: Dict[int, List[str]],
    ) -> List[str]:
        """
        Generate questions for follow-up research.

        For follow-up research, we typically return a single contextualized
        query rather than multiple questions, as the context is already rich.

        Args:
            current_knowledge: The accumulated knowledge so far
            query: The research query (already contextualized)
            questions_per_iteration: Number of questions to generate
            questions_by_iteration: Previous questions

        Returns:
            List[str]: List containing the contextualized query
        """
        # For follow-up research, the query is already contextualized
        # Just return it as a single-item list
        return [query]
