"""
Entity-aware question generation for improved entity identification.
"""

from datetime import datetime, UTC
from typing import List

from loguru import logger

from .base_question import BaseQuestionGenerator


class EntityAwareQuestionGenerator(BaseQuestionGenerator):
    """Question generator that creates more targeted searches for entity identification."""

    def generate_questions(
        self,
        current_knowledge: str,
        query: str,
        questions_per_iteration: int = 2,
        questions_by_iteration: dict = None,
    ) -> List[str]:
        """Generate questions with entity-aware search patterns."""
        now = datetime.now(UTC)
        current_time = now.strftime("%Y-%m-%d")
        questions_by_iteration = questions_by_iteration or {}

        logger.info("Generating entity-aware follow-up questions...")

        # Detect if this is likely an entity identification query
        entity_keywords = [
            "who",
            "what",
            "which",
            "identify",
            "name",
            "character",
            "person",
            "place",
            "organization",
            "company",
            "author",
            "scientist",
            "inventor",
            "city",
            "country",
            "book",
            "movie",
        ]

        is_entity_query = any(
            keyword in query.lower() for keyword in entity_keywords
        )

        if is_entity_query:
            # Use more direct entity-focused prompt
            if questions_by_iteration:
                prompt = f"""Generate {questions_per_iteration} targeted search queries to identify the specific entity in the query.

Query: {query}
Today: {current_time}
Past questions: {questions_by_iteration!s}
Current knowledge: {current_knowledge}

Create direct search queries that combine the key identifying features to find the specific name/entity.
Focus on:
1. Combining multiple constraints in a single search
2. Using quotation marks for exact phrases
3. Including specific details that narrow down results

Format: One question per line, e.g.
Q: "fictional character" "breaks fourth wall" "TV show" 1960s 1980s
Q: character name ascetics humor television fewer than 50 episodes
"""
            else:
                prompt = f"""Generate {questions_per_iteration} direct search queries to identify the specific entity in: {query}

Today: {current_time}

Create search queries that:
1. Combine multiple identifying features
2. Target the specific entity name/identification
3. Use variations of key terms

Format: One question per line, e.g.
Q: question1
Q: question2
"""
        else:
            # Fall back to standard question generation for non-entity queries
            return super().generate_questions(
                current_knowledge,
                query,
                questions_per_iteration,
                questions_by_iteration,
            )

        response = self.model.invoke(prompt)

        # Handle both string responses and responses with .content attribute
        response_text = ""
        if hasattr(response, "content"):
            response_text = response.content
        else:
            response_text = str(response)

        questions = [
            q.replace("Q:", "").strip()
            for q in response_text.split("\n")
            if q.strip().startswith("Q:")
        ][:questions_per_iteration]

        logger.info(f"Generated {len(questions)} entity-aware questions")

        return questions

    def generate_sub_questions(
        self, query: str, context: str = ""
    ) -> List[str]:
        """Generate sub-questions with entity focus when appropriate."""
        # Check if this is an entity identification query
        entity_keywords = [
            "who",
            "what",
            "which",
            "identify",
            "name",
            "character",
            "person",
            "place",
            "organization",
            "company",
        ]

        is_entity_query = any(
            keyword in query.lower() for keyword in entity_keywords
        )

        if is_entity_query:
            prompt = f"""Break down this entity identification query into targeted sub-questions.

Original Question: {query}
{context}

Generate 2-5 sub-questions that will help identify the specific entity.
Focus on:
1. Combining constraints to narrow down results
2. Finding the actual name/identity
3. Verifying the entity matches all criteria

Format your response as:
1. First sub-question
2. Second sub-question
...

Only provide the numbered sub-questions."""
        else:
            return super().generate_sub_questions(query, context)

        try:
            response = self.model.invoke(prompt)
            content = ""
            if hasattr(response, "content"):
                content = response.content
            else:
                content = str(response)

            # Extract numbered questions
            questions = []
            for line in content.strip().split("\n"):
                line = line.strip()
                if line and (line[0].isdigit() or line.startswith("-")):
                    # Remove the number/bullet and clean up
                    question = line.split(".", 1)[-1].strip()
                    question = question.lstrip("- ").strip()
                    if question:
                        questions.append(question)

            return questions

        except Exception as e:
            logger.exception(f"Error generating sub-questions: {e!s}")
            return []
