"""
Parallel candidate explorer implementation.

This explorer runs multiple search queries in parallel to quickly discover
a wide range of candidates.
"""

import concurrent.futures
import time
from typing import List, Optional

from loguru import logger

from ...database.thread_local_session import thread_cleanup
from ..candidates.base_candidate import Candidate
from ..constraints.base_constraint import Constraint
from .base_explorer import (
    BaseCandidateExplorer,
    ExplorationResult,
    ExplorationStrategy,
)


class ParallelExplorer(BaseCandidateExplorer):
    """
    Parallel candidate explorer that runs multiple searches concurrently.

    This explorer:
    1. Generates multiple search queries from the initial query
    2. Runs searches in parallel for speed
    3. Collects and deduplicates candidates
    4. Focuses on breadth-first exploration
    """

    def __init__(
        self,
        *args,
        max_workers: int = 5,
        queries_per_round: int = 8,
        max_rounds: int = 3,
        **kwargs,
    ):
        """
        Initialize parallel explorer.

        Args:
            max_workers: Maximum number of parallel search threads
            queries_per_round: Number of queries to generate per round
            max_rounds: Maximum exploration rounds
        """
        super().__init__(*args, **kwargs)
        self.max_workers = max_workers
        self.queries_per_round = queries_per_round
        self.max_rounds = max_rounds

    def explore(
        self,
        initial_query: str,
        constraints: Optional[List[Constraint]] = None,
        entity_type: Optional[str] = None,
    ) -> ExplorationResult:
        """Explore candidates using parallel search strategy."""
        start_time = time.time()
        logger.info(f"Starting parallel exploration for: {initial_query}")

        all_candidates = []
        exploration_paths = []
        total_searched = 0

        # Initial search
        current_queries = [initial_query]

        with concurrent.futures.ThreadPoolExecutor(
            max_workers=self.max_workers
        ) as executor:
            for round_num in range(self.max_rounds):
                if not self._should_continue_exploration(
                    start_time, len(all_candidates)
                ):
                    break

                logger.info(
                    f"Exploration round {round_num + 1}: {len(current_queries)} queries"
                )

                # Submit all queries for parallel execution
                future_to_query = {
                    executor.submit(
                        thread_cleanup(self._execute_search), query
                    ): query
                    for query in current_queries
                }

                round_candidates = []

                # Collect results as they complete
                for future in concurrent.futures.as_completed(future_to_query):
                    query = future_to_query[future]
                    total_searched += 1

                    try:
                        results = future.result()
                        candidates = self._extract_candidates_from_results(
                            results, entity_type
                        )
                        round_candidates.extend(candidates)
                        exploration_paths.append(
                            f"Round {round_num + 1}: {query} -> {len(candidates)} candidates"
                        )

                    except Exception:
                        logger.exception(f"Error processing query '{query}'")

                # Add new candidates
                all_candidates.extend(round_candidates)

                # Generate queries for next round
                if round_num < self.max_rounds - 1:
                    current_queries = self.generate_exploration_queries(
                        initial_query, all_candidates, constraints
                    )[: self.queries_per_round]

                    if not current_queries:
                        logger.info("No more queries to explore")
                        break

        # Deduplicate and rank
        unique_candidates = self._deduplicate_candidates(all_candidates)
        ranked_candidates = self._rank_candidates_by_relevance(
            unique_candidates, initial_query
        )

        # Limit to max candidates
        final_candidates = ranked_candidates[: self.max_candidates]

        elapsed_time = time.time() - start_time
        logger.info(
            f"Parallel exploration completed: {len(final_candidates)} unique candidates in {elapsed_time:.1f}s"
        )

        return ExplorationResult(
            candidates=final_candidates,
            total_searched=total_searched,
            unique_candidates=len(unique_candidates),
            exploration_paths=exploration_paths,
            metadata={
                "strategy": "parallel",
                "rounds": min(round_num + 1, self.max_rounds),
                "max_workers": self.max_workers,
                "entity_type": entity_type,
            },
            elapsed_time=elapsed_time,
            strategy_used=ExplorationStrategy.BREADTH_FIRST,
        )

    def generate_exploration_queries(
        self,
        base_query: str,
        found_candidates: List[Candidate],
        constraints: Optional[List[Constraint]] = None,
    ) -> List[str]:
        """Generate queries for parallel exploration."""
        queries = []

        # Query variations based on base query
        base_variations = self._generate_query_variations(base_query)
        queries.extend(base_variations)

        # Queries based on found candidates
        if found_candidates:
            candidate_queries = self._generate_candidate_based_queries(
                found_candidates, base_query
            )
            queries.extend(candidate_queries)

        # Constraint-based queries
        if constraints:
            constraint_queries = self._generate_constraint_queries(
                constraints, base_query
            )
            queries.extend(constraint_queries)

        # Remove already explored queries
        new_queries = [
            q for q in queries if q.lower() not in self.explored_queries
        ]

        return new_queries[: self.queries_per_round]

    def _generate_query_variations(self, base_query: str) -> List[str]:
        """Generate variations of the base query."""
        try:
            prompt = f"""
Generate 4 search query variations for: "{base_query}"

Each variation should:
1. Use different keywords but same intent
2. Be specific and searchable
3. Focus on finding concrete examples or instances

Format as numbered list:
1. [query]
2. [query]
3. [query]
4. [query]
"""

            response = self.model.invoke(prompt).content.strip()

            # Parse numbered list
            queries = []
            for line in response.split("\n"):
                line = line.strip()
                if line and any(line.startswith(f"{i}.") for i in range(1, 10)):
                    # Remove number prefix
                    query = line.split(".", 1)[1].strip()
                    if query:
                        queries.append(query)

            return queries[:4]

        except Exception:
            logger.exception("Error generating query variations")
            return []

    def _generate_candidate_based_queries(
        self, candidates: List[Candidate], base_query: str
    ) -> List[str]:
        """Generate queries based on found candidates."""
        queries = []

        # Sample a few candidates to avoid too many queries
        sample_candidates = candidates[:3]

        for candidate in sample_candidates:
            # Query for similar entities
            queries.append(f'similar to "{candidate.name}"')
            queries.append(f'like "{candidate.name}" examples')

        return queries

    def _generate_constraint_queries(
        self, constraints: List[Constraint], base_query: str
    ) -> List[str]:
        """Generate queries focusing on specific constraints."""
        queries = []

        # Sample constraints to avoid too many queries
        for constraint in constraints[:2]:
            queries.append(f"{constraint.value} examples")
            queries.append(f'"{constraint.value}" instances')

        return queries
