"""
Parallel search strategy implementation for maximum search speed.
"""

import concurrent.futures
import logging
from typing import Dict

from ...citation_handler import CitationHandler
from ...config.llm_config import get_llm
from ...config.search_config import get_search
from ...utilities.db_utils import get_db_setting
from ...utilities.search_utilities import extract_links_from_search_results
from ..filters.cross_engine_filter import CrossEngineFilter
from ..findings.repository import FindingsRepository
from ..questions.standard_question import StandardQuestionGenerator
from .base_strategy import BaseSearchStrategy

logger = logging.getLogger(__name__)


class ParallelSearchStrategy(BaseSearchStrategy):
    """
    Parallel search strategy that generates questions and runs all searches
    simultaneously for maximum speed.
    """

    def __init__(
        self,
        search=None,
        model=None,
        citation_handler=None,
        include_text_content: bool = True,
        use_cross_engine_filter: bool = True,
        filter_reorder: bool = True,
        filter_reindex: bool = True,
        cross_engine_max_results: int = None,
        all_links_of_system=None,
    ):
        """Initialize with optional dependency injection for testing.

        Args:
            search: Optional search engine instance
            model: Optional LLM model instance
            citation_handler: Optional citation handler instance
            include_text_content: If False, only includes metadata and links in search results
            use_cross_engine_filter: If True, filter search results across engines
            filter_reorder: Whether to reorder results by relevance
            filter_reindex: Whether to update result indices after filtering
            cross_engine_max_results: Maximum number of results to keep after cross-engine filtering
            all_links_of_system: Optional list of links to initialize with
        """
        super().__init__(all_links_of_system=all_links_of_system)
        self.search = search or get_search()
        self.model = model or get_llm()
        self.progress_callback = None
        self.questions_by_iteration = {}
        self.include_text_content = include_text_content
        self.use_cross_engine_filter = use_cross_engine_filter
        self.filter_reorder = filter_reorder
        self.filter_reindex = filter_reindex

        # Get max_filtered_results from database if not provided
        if cross_engine_max_results is None:
            cross_engine_max_results = get_db_setting(
                "search.cross_engine_max_results", 100
            )

        # Initialize the cross-engine filter
        self.cross_engine_filter = CrossEngineFilter(
            model=self.model,
            max_results=cross_engine_max_results,
            default_reorder=filter_reorder,
            default_reindex=filter_reindex,
        )

        # Set include_full_content on the search engine if it supports it
        if hasattr(self.search, "include_full_content"):
            self.search.include_full_content = include_text_content

        # Use provided citation_handler or create one
        self.citation_handler = citation_handler or CitationHandler(self.model)

        # Initialize components
        self.question_generator = StandardQuestionGenerator(self.model)
        self.findings_repository = FindingsRepository(self.model)

    def analyze_topic(self, query: str) -> Dict:
        """
        Analyze a topic using parallel search, supporting multiple iterations.

        Args:
            query: The research query to analyze
        """
        logger.info(f"Starting parallel research on topic: {query}")

        findings = []
        all_search_results = []
        current_knowledge = ""

        # Track all search results across iterations
        self.all_links_of_system = list()
        self.questions_by_iteration = {}

        self._update_progress(
            "Initializing parallel research",
            5,
            {
                "phase": "init",
                "strategy": "parallel",
                "include_text_content": self.include_text_content,
            },
        )

        # Check search engine
        if not self._validate_search_engine():
            return {
                "findings": [],
                "iterations": 0,
                "questions_by_iteration": {},
                "formatted_findings": "Error: Unable to conduct research without a search engine.",
                "current_knowledge": "",
                "error": "No search engine available",
            }

        # Determine number of iterations to run
        iterations_to_run = get_db_setting("search.iterations")
        logger.debug("Selected amount of iterations: " + str(iterations_to_run))
        iterations_to_run = int(iterations_to_run)
        try:
            # Run each iteration
            for iteration in range(1, iterations_to_run + 1):
                iteration_progress_base = 5 + (iteration - 1) * (70 / iterations_to_run)

                self._update_progress(
                    f"Starting iteration {iteration}/{iterations_to_run}",
                    iteration_progress_base,
                    {"phase": f"iteration_{iteration}", "iteration": iteration},
                )

                # Step 1: Generate questions
                self._update_progress(
                    f"Generating search questions for iteration {iteration}",
                    iteration_progress_base + 5,
                    {"phase": "question_generation", "iteration": iteration},
                )

                # For first iteration, generate initial questions
                # For subsequent iterations, generate follow-up questions
                logger.info("Starting to generate questions")
                if iteration == 1:
                    # Generate additional questions (plus the main query)
                    if iterations_to_run > 1:
                        context = f"""Iteration: {1} of {iterations_to_run}"""
                    else:
                        context = ""
                    questions = self.question_generator.generate_questions(
                        current_knowledge=context,
                        query=query,
                        questions_per_iteration=int(
                            get_db_setting("search.questions_per_iteration")
                        ),
                        questions_by_iteration=self.questions_by_iteration,
                    )

                    # Add the original query as the first question
                    all_questions = [query] + questions

                    # Store in questions_by_iteration
                    self.questions_by_iteration[iteration] = questions
                    logger.info(
                        f"Generated questions for iteration {iteration}: {questions}"
                    )
                else:
                    # Get past questions from all previous iterations
                    past_questions = []
                    for prev_iter in range(1, iteration):
                        if prev_iter in self.questions_by_iteration:
                            past_questions.extend(
                                self.questions_by_iteration[prev_iter]
                            )

                    # Generate follow-up questions based on accumulated knowledge if iterations > 2
                    use_knowledge = iterations_to_run > 2
                    knowledge_for_questions = current_knowledge if use_knowledge else ""
                    context = f"""Current Knowledge: {knowledge_for_questions}
                    Iteration: {iteration} of {iterations_to_run}"""

                    # Generate questions
                    questions = self.question_generator.generate_questions(
                        current_knowledge=context,
                        query=query,
                        questions_per_iteration=int(
                            get_db_setting("search.questions_per_iteration")
                        ),
                        questions_by_iteration=self.questions_by_iteration,
                    )

                    # Use only the new questions for this iteration's searches
                    all_questions = questions

                    # Store in questions_by_iteration
                    self.questions_by_iteration[iteration] = questions
                    logger.info(
                        f"Generated questions for iteration {iteration}: {questions}"
                    )

                # Step 2: Run all searches in parallel for this iteration
                self._update_progress(
                    f"Running parallel searches for iteration {iteration}",
                    iteration_progress_base + 10,
                    {"phase": "parallel_search", "iteration": iteration},
                )

                # Function for thread pool
                def search_question(q):
                    try:
                        result = self.search.run(q)
                        return {"question": q, "results": result or []}
                    except Exception as e:
                        logger.error(f"Error searching for '{q}': {str(e)}")
                        return {"question": q, "results": [], "error": str(e)}

                # Run searches in parallel
                with concurrent.futures.ThreadPoolExecutor(
                    max_workers=len(all_questions)
                ) as executor:
                    futures = [
                        executor.submit(search_question, q) for q in all_questions
                    ]
                    iteration_search_dict = {}
                    iteration_search_results = []

                    # Process results as they complete
                    for i, future in enumerate(
                        concurrent.futures.as_completed(futures)
                    ):
                        result_dict = future.result()
                        question = result_dict["question"]
                        search_results = result_dict["results"]
                        iteration_search_dict[question] = search_results

                        self._update_progress(
                            f"Completed search {i + 1} of {len(all_questions)}: {question[:30]}...",
                            iteration_progress_base
                            + 10
                            + ((i + 1) / len(all_questions) * 30),
                            {
                                "phase": "search_complete",
                                "iteration": iteration,
                                "result_count": len(search_results),
                                "question": question,
                            },
                        )

                        # Collect all search results for this iteration
                        iteration_search_results.extend(search_results)

                # Step 3: Filter and analyze results for this iteration
                self._update_progress(
                    f"Analyzing results for iteration {iteration}",
                    iteration_progress_base + 45,
                    {"phase": "iteration_analysis", "iteration": iteration},
                )

                # Apply cross-engine filtering if enabled
                if self.use_cross_engine_filter:
                    self._update_progress(
                        f"Filtering search results for iteration {iteration}",
                        iteration_progress_base + 45,
                        {"phase": "cross_engine_filtering", "iteration": iteration},
                    )

                    # Get the current link count (for indexing)
                    existing_link_count = len(self.all_links_of_system)

                    # Filter the search results
                    filtered_search_results = self.cross_engine_filter.filter_results(
                        iteration_search_results,
                        query,
                        reorder=self.filter_reorder,
                        reindex=self.filter_reindex,
                        start_index=existing_link_count,  # Start indexing after existing links
                    )

                    links = extract_links_from_search_results(filtered_search_results)
                    self.all_links_of_system.extend(links)

                    self._update_progress(
                        f"Filtered from {len(iteration_search_results)} to {len(filtered_search_results)} results",
                        iteration_progress_base + 50,
                        {
                            "phase": "filtering_complete",
                            "iteration": iteration,
                            "links_count": len(self.all_links_of_system),
                        },
                    )

                    # Use filtered results for analysis
                    iteration_search_results = filtered_search_results
                else:
                    # Just extract links without filtering
                    links = extract_links_from_search_results(iteration_search_results)
                    self.all_links_of_system.extend(links)

                # Add to all search results
                all_search_results.extend(iteration_search_results)

                # Create a finding for this iteration's results
                if self.include_text_content and iteration_search_results:
                    # For iteration > 1 with knowledge accumulation, use follow-up analysis
                    if iteration > 1 and iterations_to_run > 2:
                        citation_result = self.citation_handler.analyze_followup(
                            query,
                            iteration_search_results,
                            current_knowledge,
                            len(self.all_links_of_system) - len(links),
                        )
                    else:
                        # For first iteration or without knowledge accumulation, use initial analysis
                        citation_result = self.citation_handler.analyze_initial(
                            query, iteration_search_results
                        )

                    if citation_result:
                        # Create a finding for this iteration
                        iteration_content = citation_result["content"]

                        # Update current knowledge if iterations > 2
                        if iterations_to_run > 2:
                            if current_knowledge:
                                current_knowledge = f"{current_knowledge}\n\n## FINDINGS FROM ITERATION {iteration}:\n\n{iteration_content}"
                            else:
                                current_knowledge = iteration_content

                        finding = {
                            "phase": f"Iteration {iteration}",
                            "content": iteration_content,
                            "question": query,
                            "search_results": iteration_search_results,
                            "documents": citation_result.get("documents", []),
                        }
                        findings.append(finding)

                        # Add documents to repository
                        if "documents" in citation_result:
                            self.findings_repository.add_documents(
                                citation_result["documents"]
                            )

                # Mark iteration as complete
                iteration_progress = 5 + iteration * (70 / iterations_to_run)
                self._update_progress(
                    f"Completed iteration {iteration}/{iterations_to_run}",
                    iteration_progress,
                    {"phase": "iteration_complete", "iteration": iteration},
                )

            # Final synthesis after all iterations
            self._update_progress(
                "Generating final synthesis", 80, {"phase": "synthesis"}
            )

            # Handle final synthesis based on include_text_content flag
            if self.include_text_content:
                # Generate a final synthesis from all search results
                if iterations_to_run > 1:
                    final_citation_result = self.citation_handler.analyze_initial(
                        query, all_search_results
                    )
                    # Add null check for final_citation_result
                    if final_citation_result:
                        synthesized_content = final_citation_result["content"]
                    else:
                        synthesized_content = (
                            "No relevant results found in final synthesis."
                        )
                else:
                    # For single iteration, use the content from findings
                    synthesized_content = (
                        findings[0]["content"]
                        if findings
                        else "No relevant results found."
                    )
                # Add a final synthesis finding
                final_finding = {
                    "phase": "Final synthesis",
                    "content": synthesized_content,
                    "question": query,
                    "search_results": all_search_results,
                    "documents": [],
                }
                findings.append(final_finding)
            else:
                # Skip LLM analysis, just format the raw search results
                synthesized_content = "LLM analysis skipped"
                final_finding = {
                    "phase": "Raw search results",
                    "content": "LLM analysis was skipped. Displaying raw search results with links.",
                    "question": query,
                    "search_results": all_search_results,
                    "documents": [],
                }
                findings.append(final_finding)

            # Transfer questions to repository
            self.findings_repository.set_questions_by_iteration(
                self.questions_by_iteration
            )

            # Format findings
            formatted_findings = self.findings_repository.format_findings_to_text(
                findings, synthesized_content
            )

        except Exception as e:
            import traceback

            error_msg = f"Error in research process: {str(e)}"
            logger.error(error_msg)
            logger.error(traceback.format_exc())
            synthesized_content = f"Error: {str(e)}"
            formatted_findings = f"Error: {str(e)}"
            finding = {
                "phase": "Error",
                "content": synthesized_content,
                "question": query,
                "search_results": [],
                "documents": [],
            }
            findings.append(finding)

        self._update_progress("Research complete", 100, {"phase": "complete"})

        return {
            "findings": findings,
            "iterations": iterations_to_run,
            "questions_by_iteration": self.questions_by_iteration,
            "formatted_findings": formatted_findings,
            "current_knowledge": synthesized_content,
        }
