"""
BrowseComp-specific question generation that creates progressive, entity-focused searches.
"""

import re
from typing import Dict, List, Optional

from loguru import logger

from .base_question import BaseQuestionGenerator


class BrowseCompQuestionGenerator(BaseQuestionGenerator):
    """
    Question generator optimized for BrowseComp-style queries.

    Key features:
    1. Extract concrete entities (dates, numbers, names, places)
    2. Generate progressive search combinations
    3. Start broad, then narrow systematically
    4. Focus on verifiable facts
    """

    def __init__(
        self,
        model,
        knowledge_truncate_length: int = 1500,
        previous_searches_limit: int = 10,
    ):
        """Initialize the question generator.

        Args:
            model: The LLM model to use for generation
            knowledge_truncate_length: Max chars for knowledge in prompts (None=unlimited)
            previous_searches_limit: Max previous searches to show (None=unlimited)
        """
        super().__init__(model)
        self.extracted_entities: Dict[str, List[str]] = {}
        self.search_progression: List[str] = []
        self.knowledge_truncate_length = knowledge_truncate_length
        self.previous_searches_limit = previous_searches_limit

    def generate_questions(
        self,
        current_knowledge: str,
        query: str,
        questions_per_iteration: int = 5,
        questions_by_iteration: Optional[dict] = None,
        results_by_iteration: Optional[dict] = None,
        iteration: int = 1,
    ) -> List[str]:
        """Generate progressive search queries for BrowseComp problems."""
        questions_by_iteration = questions_by_iteration or {}

        # First iteration: Extract entities and create initial searches
        if iteration == 1 or not self.extracted_entities:
            self.extracted_entities = self._extract_entities(query)
            return self._generate_initial_searches(
                query, self.extracted_entities, questions_per_iteration
            )

        # Subsequent iterations: Progressive refinement
        return self._generate_progressive_searches(
            query,
            current_knowledge,
            self.extracted_entities,
            questions_by_iteration,
            results_by_iteration or {},
            questions_per_iteration,
            iteration,
        )

    def _extract_entities(self, query: str) -> Dict[str, List[str]]:
        """Extract concrete entities from the query."""
        prompt = f"""Extract ALL concrete, searchable entities from this query:

Query: {query}

Extract:
1. TEMPORAL: All years, dates, time periods (e.g., "2018", "between 1995 and 2006", "2023")
2. NUMERICAL: All numbers, statistics, counts (e.g., "300", "more than 3", "4-3", "84.5%")
3. NAMES: Partial names, name hints, proper nouns (e.g., "Dartmouth", "EMNLP", "Plastic Man")
4. LOCATIONS: Places, institutions, geographic features (e.g., "Pennsylvania", "Grand Canyon")
5. DESCRIPTORS: Key descriptive terms (e.g., "fourth wall", "ascetics", "decider game")

For TEMPORAL entities, if there's a range (e.g., "between 2018-2023"), list EACH individual year.

Format your response as:
TEMPORAL: [entity1], [entity2], ...
NUMERICAL: [entity1], [entity2], ...
NAMES: [entity1], [entity2], ...
LOCATIONS: [entity1], [entity2], ...
DESCRIPTORS: [entity1], [entity2], ...
"""

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

        entities: Dict[str, List[str]] = {
            "temporal": [],
            "numerical": [],
            "names": [],
            "locations": [],
            "descriptors": [],
        }

        for line in content.strip().split("\n"):
            line = line.strip()
            if ":" in line:
                category, values = line.split(":", 1)
                category = category.strip().lower()
                if category in entities:
                    # Parse comma-separated values
                    values = [v.strip() for v in values.split(",") if v.strip()]
                    entities[category].extend(values)

        # Expand temporal ranges
        entities["temporal"] = self._expand_temporal_ranges(
            entities["temporal"]
        )

        logger.info(f"Extracted entities: {entities}")
        return entities

    def _expand_temporal_ranges(
        self, temporal_entities: List[str]
    ) -> List[str]:
        """Expand year ranges into individual years."""
        expanded = []
        for entity in temporal_entities:
            # Check for range patterns like "2018-2023" or "between 1995 and 2006"
            range_match = re.search(
                r"(\d{4})[-\s]+(?:to|and)?\s*(\d{4})", entity
            )
            if range_match:
                start_year = int(range_match.group(1))
                end_year = int(range_match.group(2))
                for year in range(start_year, end_year + 1):
                    expanded.append(str(year))
            else:
                # Single year or other temporal entity
                year_match = re.search(r"\d{4}", entity)
                if year_match:
                    expanded.append(year_match.group())
                else:
                    expanded.append(entity)

        return list(set(expanded))  # Remove duplicates

    def _generate_initial_searches(
        self, query: str, entities: Dict[str, List[str]], num_questions: int
    ) -> List[str]:
        """Generate initial broad searches."""
        searches = []

        # 1. Original query (always include)
        searches.append(query)

        # If only 1 question requested, return just the original query
        if num_questions <= 1:
            return searches[:1]

        # 2. Domain exploration searches (combine key entities)
        if entities["names"] and len(searches) < num_questions:
            for name in entities["names"][:2]:  # Top 2 names
                if len(searches) >= num_questions:
                    break
                searches.append(f"{name}")
                if entities["descriptors"] and len(searches) < num_questions:
                    searches.append(f"{name} {entities['descriptors'][0]}")

        # 3. Temporal searches if years are important
        if (
            entities["temporal"]
            and len(entities["temporal"]) <= 10
            and len(searches) < num_questions
        ):
            # For small year ranges, search each year with a key term
            key_term = (
                entities["names"][0]
                if entities["names"]
                else entities["descriptors"][0]
                if entities["descriptors"]
                else ""
            )
            for year in entities["temporal"][:5]:  # Limit to 5 years initially
                if len(searches) >= num_questions:
                    break
                if key_term:
                    searches.append(f"{key_term} {year}")

        # 4. Location-based searches
        if entities["locations"] and len(searches) < num_questions:
            for location in entities["locations"][:2]:
                if len(searches) >= num_questions:
                    break
                searches.append(f"{location}")
                if entities["descriptors"] and len(searches) < num_questions:
                    searches.append(f"{location} {entities['descriptors'][0]}")

        # Remove duplicates and limit to requested number
        seen = set()
        unique_searches = []
        for s in searches:
            if s.lower() not in seen:
                seen.add(s.lower())
                unique_searches.append(s)
                if len(unique_searches) >= num_questions:
                    break

        return unique_searches[:num_questions]

    def _generate_progressive_searches(
        self,
        query: str,
        current_knowledge: str,
        entities: Dict[str, List[str]],
        questions_by_iteration: dict,
        results_by_iteration: dict,
        num_questions: int,
        iteration: int,
    ) -> List[str]:
        """Generate progressively more specific searches based on findings."""

        # Only add strategy instructions if we have actual result data (adaptive mode)
        strategy_instruction = ""
        if results_by_iteration:
            # Check if recent searches are failing (returning 0 results)
            recent_iterations = list(range(max(1, iteration - 5), iteration))
            zero_count = sum(
                1
                for i in recent_iterations
                if results_by_iteration.get(i, 1) == 0
            )
            searches_failing = zero_count >= 3

            # Adjust strategy based on success/failure
            if searches_failing:
                strategy_instruction = """
IMPORTANT: Your recent searches are returning 0 results - they are TOO NARROW!
- Use FEWER constraints (1-2 terms instead of 4-5)
- Try BROADER, more general searches
- Remove overly specific combinations
- Focus on key concepts, not detailed entity combinations
"""
            else:
                strategy_instruction = """
Focus on finding the specific answer by combining entities systematically.
"""

        # Analyze what we've found so far
        prompt = f"""Based on our search progress, generate targeted follow-up searches.
{strategy_instruction}

Original Query: {query}

Entities Found:
- Names/Terms: {", ".join(entities["names"][:5])}
- Years: {", ".join(entities["temporal"][:5])}
- Locations: {", ".join(entities["locations"][:3])}
- Key Features: {", ".join(entities["descriptors"][:3])}

Current Knowledge Summary:
{current_knowledge[: self.knowledge_truncate_length] if self.knowledge_truncate_length else current_knowledge}

Previous Searches:
{self._format_previous_searches(questions_by_iteration, results_by_iteration)}

Generate {num_questions} NEW search queries that:
1. Combine 2-3 entities we haven't tried together
2. If we found candidate names, search for them with other constraints
3. For year ranges, systematically cover years we haven't searched
4. Use quotes for exact phrases when beneficial

Focus on finding the specific answer, not general information.

Format: One search per line
"""

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

        # Extract searches from response
        searches = []
        for line in content.strip().split("\n"):
            line = line.strip()
            if line and not line.endswith(":") and len(line) > 5:
                # Clean up common prefixes
                for prefix in ["Q:", "Search:", "-", "*", "•"]:
                    if line.startswith(prefix):
                        line = line[len(prefix) :].strip()
                if line:
                    searches.append(line)

        # Ensure we have enough searches, but respect the limit
        while len(searches) < num_questions:
            # Generate combinations programmatically
            if iteration <= 5 and entities["temporal"]:
                # Continue with year-based searches
                added_any = False
                for year in entities["temporal"]:
                    if not self._was_searched(year, questions_by_iteration):
                        base_term = (
                            entities["names"][0] if entities["names"] else ""
                        )
                        searches.append(f"{base_term} {year}".strip())
                        added_any = True
                        if len(searches) >= num_questions:
                            break
                if not added_any:
                    break  # No more year searches to add
            else:
                # Combine multiple constraints
                added_any = False
                if entities["names"] and entities["descriptors"]:
                    for name in entities["names"]:
                        for desc in entities["descriptors"]:
                            combo = f"{name} {desc}"
                            if not self._was_searched(
                                combo, questions_by_iteration
                            ):
                                searches.append(combo)
                                added_any = True
                                if len(searches) >= num_questions:
                                    break
                        if len(searches) >= num_questions:
                            break
                if not added_any:
                    break  # No more combinations to add

        return searches[:num_questions]

    def _format_previous_searches(
        self, questions_by_iteration: dict, results_by_iteration: dict
    ) -> str:
        """Format previous searches for context with result counts."""
        formatted = []
        for iteration, questions in questions_by_iteration.items():
            if isinstance(questions, list):
                result_count = results_by_iteration.get(iteration, "?")
                # Limit questions per iteration (main uses 3)
                questions_to_show = (
                    questions[:3] if self.previous_searches_limit else questions
                )
                for q in questions_to_show:
                    # Only show result counts if we have actual data (not "?")
                    if result_count != "?":
                        formatted.append(
                            f"Iteration {iteration}: {q} ({result_count} results)"
                        )
                    else:
                        formatted.append(f"Iteration {iteration}: {q}")
        # Apply limit if configured (main uses last 10)
        if self.previous_searches_limit:
            return "\n".join(formatted[-self.previous_searches_limit :])
        return "\n".join(formatted)

    def _was_searched(self, term: str, questions_by_iteration: dict) -> bool:
        """Check if a term was already searched."""
        term_lower = term.lower()
        for questions in questions_by_iteration.values():
            if isinstance(questions, list):
                for q in questions:
                    if term_lower in q.lower():
                        return True
        return False
