"""
Focused Iteration Strategy - **PROVEN HIGH-PERFORMANCE STRATEGY FOR SIMPLEQA**

**PERFORMANCE RECORD:**
- SimpleQA Accuracy: 96.51% (CONFIRMED HIGH PERFORMER)
- Optimal Configuration: 8 iterations, 5 questions/iteration, GPT-4.1 Mini
- Status: PRESERVE THIS STRATEGY - Core SimpleQA implementation

This strategy achieves excellent SimpleQA performance by:
1. Using simple, direct search execution (like source-based)
2. Progressive entity-focused exploration
3. No early filtering or complex constraint checking
4. Trusting the LLM for final synthesis

IMPORTANT: This strategy works exceptionally well for SimpleQA. Any modifications
should preserve the core approach that achieves 96.51% accuracy.

**BrowseComp Enhancement:** Also includes BrowseComp-specific optimizations
when use_browsecomp_optimization=True, but SimpleQA performance is the priority.
"""

import concurrent.futures
from typing import Dict, List

from loguru import logger

from ...citation_handler import CitationHandler

# Model and search should be provided by AdvancedSearchSystem
from ...utilities.thread_context import (
    preserve_research_context,
    get_search_context,
)
from ..candidate_exploration import ProgressiveExplorer
from ..findings.repository import FindingsRepository
from ..questions import BrowseCompQuestionGenerator
from .base_strategy import BaseSearchStrategy


class FocusedIterationStrategy(BaseSearchStrategy):
    """
    A hybrid strategy that combines the simplicity of source-based search
    with BrowseComp-optimized progressive exploration.

    Key principles:
    1. Start broad, then narrow progressively
    2. Extract and systematically search entities
    3. Keep all results without filtering
    4. Trust LLM for final constraint matching
    5. Use more iterations for thorough exploration
    """

    def __init__(
        self,
        model,
        search,
        citation_handler=None,
        all_links_of_system=None,
        max_iterations: int = 8,  # OPTIMAL FOR SIMPLEQA: 90%+ accuracy achieved
        questions_per_iteration: int = 5,  # OPTIMAL FOR SIMPLEQA: proven config
        use_browsecomp_optimization: bool = True,  # True for 90%+ accuracy with forced_answer handler
        settings_snapshot=None,
        # Options to match main branch behavior for testing:
        enable_adaptive_questions: bool = False,  # Pass results_by_iteration to question generator (False=main behavior)
        enable_early_termination: bool = False,  # Allow early stopping when high confidence found
        knowledge_summary_limit: int = 10,  # Limit results in knowledge summary (None=unlimited, 10=main behavior)
        knowledge_snippet_truncate: int = 200,  # Truncate snippets to N chars (None=no truncation, 200=main behavior)
        prompt_knowledge_truncate: int = 1500,  # Truncate knowledge in LLM prompt (None=unlimited, 1500=main behavior)
        previous_searches_limit: int = 10,  # Limit previous searches shown to LLM (None=unlimited, 10=main behavior)
    ):
        """Initialize with components optimized for focused iteration."""
        super().__init__(all_links_of_system, settings_snapshot)
        self.search = search
        self.model = model
        self.progress_callback = None

        # Configuration - ensure these are integers with defaults
        self.max_iterations = (
            int(max_iterations) if max_iterations is not None else 3
        )
        self.questions_per_iteration = (
            int(questions_per_iteration)
            if questions_per_iteration is not None
            else 3
        )
        self.use_browsecomp_optimization = use_browsecomp_optimization

        # Options to control behavior (for A/B testing main vs dev behavior)
        self.enable_adaptive_questions = enable_adaptive_questions
        self.enable_early_termination = enable_early_termination
        self.knowledge_summary_limit = knowledge_summary_limit
        self.knowledge_snippet_truncate = knowledge_snippet_truncate
        self.prompt_knowledge_truncate = prompt_knowledge_truncate
        self.previous_searches_limit = previous_searches_limit

        logger.info(
            f"FocusedIterationStrategy configuration - max_iterations: {self.max_iterations}, questions_per_iteration: {self.questions_per_iteration}"
        )
        logger.debug(
            f"FocusedIterationStrategy - use_browsecomp_optimization: {self.use_browsecomp_optimization}"
        )

        # Initialize specialized components
        if use_browsecomp_optimization:
            # Pass truncation settings to question generator
            self.question_generator = BrowseCompQuestionGenerator(
                self.model,
                knowledge_truncate_length=prompt_knowledge_truncate,
                previous_searches_limit=previous_searches_limit,
            )
            self.explorer = ProgressiveExplorer(self.search, self.model)
        else:
            # Fall back to standard components
            from ..questions import StandardQuestionGenerator

            self.question_generator = StandardQuestionGenerator(self.model)
            self.explorer = None

        # Use forced answer handler for BrowseComp optimization
        handler_type = (
            "forced_answer" if use_browsecomp_optimization else "standard"
        )
        self.citation_handler = citation_handler or CitationHandler(
            self.model, handler_type=handler_type
        )
        self.findings_repository = FindingsRepository(self.model)

        # Track all search results
        self.all_search_results = []
        # Note: questions_by_iteration is already initialized by parent class
        # Track result counts per iteration for question generator feedback
        self.results_by_iteration = {}

    def analyze_topic(self, query: str) -> Dict:
        """
        Analyze topic using focused iteration approach.

        Combines simplicity of source-based with progressive BrowseComp optimization.
        """
        logger.info(f"Starting focused iteration search: {query}")

        self._update_progress(
            "Initializing focused iteration search",
            5,
            {
                "phase": "init",
                "strategy": "focused_iteration",
                "max_iterations": self.max_iterations,
                "browsecomp_optimized": self.use_browsecomp_optimization,
            },
        )

        # Validate search engine
        if not self._validate_search_engine():
            return self._create_error_response("No search engine available")

        findings = []
        extracted_entities = {}

        try:
            # Main iteration loop
            for iteration in range(1, self.max_iterations + 1):
                iteration_progress = 10 + (iteration - 1) * (
                    80 / self.max_iterations
                )

                # Show context-aware progress message for question generation
                self._emit_question_generation_progress(
                    iteration=iteration,
                    progress_percent=iteration_progress + 2,
                    source_count=len(self.all_search_results)
                    if iteration > 1
                    else 0,
                    query=query,
                )

                # Generate questions for this iteration
                if self.use_browsecomp_optimization:
                    # Use BrowseComp-aware question generation
                    # Only pass results_by_iteration if adaptive questions enabled
                    # (when disabled, LLM won't see result counts or get "too narrow" warnings)
                    questions = self.question_generator.generate_questions(
                        current_knowledge=self._get_current_knowledge_summary(),
                        query=query,
                        questions_per_iteration=self.questions_per_iteration,
                        questions_by_iteration=self.questions_by_iteration,
                        results_by_iteration=self.results_by_iteration
                        if self.enable_adaptive_questions
                        else None,
                        iteration=iteration,
                    )

                    # Extract entities on first iteration
                    if iteration == 1 and hasattr(
                        self.question_generator, "extracted_entities"
                    ):
                        extracted_entities = (
                            self.question_generator.extracted_entities
                        )
                else:
                    # Standard question generation
                    questions = self.question_generator.generate_questions(
                        current_knowledge=self._get_current_knowledge_summary(),
                        query=query,
                        questions_per_iteration=self.questions_per_iteration,
                        questions_by_iteration=self.questions_by_iteration,
                    )

                # Always include original query in first iteration, but respect question limit
                if iteration == 1 and query not in questions:
                    questions = [query] + questions
                    # Trim to respect questions_per_iteration limit
                    questions = questions[: self.questions_per_iteration]

                self.questions_by_iteration[iteration] = questions
                logger.info(f"Iteration {iteration} questions: {questions}")

                # Skip search phase if no questions were generated
                if not questions:
                    logger.warning(
                        f"No questions generated for iteration {iteration}, skipping search phase"
                    )
                    continue

                # Execute searches
                if self.explorer and self.use_browsecomp_optimization:
                    # Use progressive explorer for better tracking
                    iteration_results, search_progress = self.explorer.explore(
                        queries=questions,
                        max_workers=len(questions),
                        extracted_entities=extracted_entities,
                    )

                    # Log results but don't send progress update to avoid jumps
                    logger.info(
                        f"Found {len(search_progress.found_candidates)} candidates, "
                        f"covered {sum(len(v) for v in search_progress.entity_coverage.values())} entities"
                    )

                    # Check if we should generate verification searches
                    if iteration > 3 and search_progress.found_candidates:
                        verification_searches = (
                            self.explorer.suggest_next_searches(
                                extracted_entities, max_suggestions=2
                            )
                        )
                        if verification_searches:
                            logger.info(
                                f"Adding verification searches: {verification_searches}"
                            )
                            questions.extend(verification_searches)
                            # Re-run with verification searches
                            verification_results, _ = self.explorer.explore(
                                queries=verification_searches,
                                max_workers=len(verification_searches),
                            )
                            iteration_results.extend(verification_results)
                else:
                    # Simple parallel search (like source-based)
                    iteration_results = self._execute_parallel_searches(
                        questions
                    )

                # Accumulate all results (no filtering!)
                self.all_search_results.extend(iteration_results)
                self.all_links_of_system.extend(iteration_results)

                # Track result count for this iteration
                self.results_by_iteration[iteration] = len(iteration_results)

                # Add iteration finding
                finding = {
                    "phase": f"Iteration {iteration}",
                    "content": f"Searched with {len(questions)} questions, found {len(iteration_results)} results.",
                    "question": questions[0]
                    if len(questions) == 1
                    else "\n".join(questions)
                    if questions
                    else query,
                    "documents": [],
                }
                findings.append(finding)

                # Early termination check (controlled by enable_early_termination parameter)
                if (
                    self.enable_early_termination
                    and self._should_terminate_early(iteration)
                ):
                    logger.info(f"Early termination at iteration {iteration}")
                    break

            # Final synthesis (like source-based - trust the LLM!)
            self._update_progress(
                f"Synthesizing {len(self.all_search_results)} sources from {len(self.questions_by_iteration)} iterations...",
                90,
                {"phase": "synthesis", "type": "milestone"},
            )

            # Use citation handler for final synthesis
            # Note: nr_of_links=0 ensures citations start at [1]
            final_result = self.citation_handler.analyze_followup(
                query,
                self.all_search_results,
                previous_knowledge="",
                nr_of_links=0,
            )

            synthesized_content = final_result.get(
                "content", "No relevant results found."
            )
            documents = final_result.get("documents", [])

            # Add final synthesis finding
            final_finding = {
                "phase": "Final synthesis",
                "content": synthesized_content,
                "question": query,
                "search_results": self.all_search_results,
                "documents": documents,
            }
            findings.append(final_finding)

            # Add documents to repository
            self.findings_repository.add_documents(documents)
            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
                )
            )

            # Note: "Search complete" progress is handled by research_service after strategy returns

            # Return results
            result = {
                "findings": findings,
                "iterations": len(self.questions_by_iteration),
                "questions_by_iteration": self.questions_by_iteration,
                "formatted_findings": formatted_findings,
                "current_knowledge": synthesized_content,
                "all_links_of_system": self.all_links_of_system,
                "sources": self.all_links_of_system,
            }

            # Add BrowseComp-specific data if available
            if self.explorer and hasattr(self.explorer, "progress"):
                result["candidates"] = dict(
                    self.explorer.progress.found_candidates
                )
                result["entity_coverage"] = {
                    k: list(v)
                    for k, v in self.explorer.progress.entity_coverage.items()
                }

            return result

        except Exception as e:
            logger.exception(f"Error in focused iteration search: {e!s}")
            import traceback

            logger.exception(traceback.format_exc())
            return self._create_error_response("Search iteration failed")

    def _execute_parallel_searches(self, queries: List[str]) -> List[Dict]:
        """Execute searches in parallel (like source-based strategy)."""
        if not queries:
            logger.warning("No queries provided for parallel search")
            return []

        all_results = []

        def search_question(q):
            try:
                # Get the current research context to pass explicitly

                current_context = get_search_context()
                result = self.search.run(q, research_context=current_context)
                return {"question": q, "results": result or []}
            except Exception as e:
                logger.exception(f"Error searching '{q}': {e!s}")
                return {"question": q, "results": [], "error": "Search failed"}

        # Create context-preserving wrapper for the search function
        context_aware_search = preserve_research_context(search_question)

        # Run searches in parallel
        with concurrent.futures.ThreadPoolExecutor(
            max_workers=len(queries)
        ) as executor:
            futures = [
                executor.submit(context_aware_search, q) for q in queries
            ]

            for future in concurrent.futures.as_completed(futures):
                result_dict = future.result()
                all_results.extend(result_dict.get("results", []))

        return all_results

    def _execute_parallel_searches_with_progress(
        self, queries: List[str], iteration: int
    ) -> List[Dict]:
        """Execute searches in parallel with detailed progress reporting."""
        if not queries:
            logger.warning(
                f"No queries provided for parallel search in iteration {iteration}"
            )
            return []

        all_results = []
        completed_searches = 0
        total_searches = len(queries)

        def search_question_with_progress(q):
            nonlocal completed_searches
            try:
                # Report starting this search
                self._update_progress(
                    f"Searching: {q[:50]}{'...' if len(q) > 50 else ''}",
                    None,  # Don't update overall progress for individual searches
                    {
                        "phase": f"iteration_{iteration}_individual_search",
                        "current_query": q,
                        "search_progress": f"{completed_searches + 1}/{total_searches}",
                    },
                )

                # Get the current research context to pass explicitly

                current_context = get_search_context()
                result = self.search.run(q, research_context=current_context)
                completed_searches += 1

                # Report completion of this search
                result_count = len(result) if result else 0
                self._update_progress(
                    f"Completed search for '{q[:30]}{'...' if len(q) > 30 else ''}' - {result_count} results",
                    None,
                    {
                        "phase": f"iteration_{iteration}_search_complete",
                        "completed_query": q,
                        "results_found": result_count,
                        "search_progress": f"{completed_searches}/{total_searches}",
                    },
                )

                return {
                    "question": q,
                    "results": result or [],
                    "result_count": result_count,
                }
            except Exception as e:
                completed_searches += 1
                logger.exception(f"Error searching '{q}': {e!s}")
                self._update_progress(
                    f"Search failed for '{q[:30]}{'...' if len(q) > 30 else ''}'",
                    None,
                    {
                        "phase": f"iteration_{iteration}_search_error",
                        "failed_query": q,
                        "error": "Search error",
                        "search_progress": f"{completed_searches}/{total_searches}",
                    },
                )
                return {
                    "question": q,
                    "results": [],
                    "error": "Search failed",
                    "result_count": 0,
                }

        # Create context-preserving wrapper for the search function
        context_aware_search_with_progress = preserve_research_context(
            search_question_with_progress
        )

        # Run searches in parallel
        with concurrent.futures.ThreadPoolExecutor(
            max_workers=min(len(queries), 5)
        ) as executor:
            futures = [
                executor.submit(context_aware_search_with_progress, q)
                for q in queries
            ]

            total_results_found = 0
            for future in concurrent.futures.as_completed(futures):
                result_dict = future.result()
                results = result_dict.get("results", [])
                all_results.extend(results)
                total_results_found += result_dict.get("result_count", 0)

        # Report final iteration summary
        self._update_progress(
            f"Iteration {iteration} complete: {total_results_found} total results from {total_searches} searches",
            None,
            {
                "phase": f"iteration_{iteration}_summary",
                "total_searches": total_searches,
                "total_results": total_results_found,
                "average_results": (
                    round(total_results_found / total_searches, 1)
                    if total_searches > 0
                    else 0
                ),
            },
        )

        return all_results

    def _get_current_knowledge_summary(self) -> str:
        """Get summary of current knowledge for question generation."""
        if not self.all_search_results:
            return ""

        # Apply limit if configured (knowledge_summary_limit=10 matches main branch behavior)
        results_to_summarize = self.all_search_results
        if self.knowledge_summary_limit is not None:
            results_to_summarize = self.all_search_results[
                : self.knowledge_summary_limit
            ]

        summary_parts = []
        for i, result in enumerate(results_to_summarize):
            title = result.get("title", "")
            snippet = result.get("snippet", "")
            if title or snippet:
                # Truncate snippet if configured (200=main behavior, None=no truncation)
                if self.knowledge_snippet_truncate is not None:
                    snippet = snippet[: self.knowledge_snippet_truncate] + "..."
                summary_parts.append(f"{i + 1}. {title}: {snippet}")

        return "\n".join(summary_parts)

    def _should_terminate_early(self, iteration: int) -> bool:
        """Check if we should terminate early based on findings."""
        # For BrowseComp, continue if we're making progress
        if self.explorer and hasattr(self.explorer, "progress"):
            progress = self.explorer.progress

            # Continue if we're still finding new candidates
            if iteration > 3 and len(progress.found_candidates) > 0:
                # Check if top candidate has very high confidence
                if progress.found_candidates:
                    top_confidence = max(progress.found_candidates.values())
                    if top_confidence > 0.9:
                        return True

            # Continue if we haven't covered all entities
            if extracted_entities := getattr(
                self.question_generator, "extracted_entities", {}
            ):
                total_entities = sum(
                    len(v) for v in extracted_entities.values()
                )
                covered_entities = sum(
                    len(v) for v in progress.entity_coverage.values()
                )
                coverage_ratio = (
                    covered_entities / total_entities
                    if total_entities > 0
                    else 0
                )

                # Continue if coverage is low
                if coverage_ratio < 0.8 and iteration < 6:
                    return False

        # Default: continue to max iterations for thoroughness
        return False

    def _create_error_response(self, error_msg: str) -> Dict:
        """Create standardized error response."""
        return {
            "findings": [],
            "iterations": 0,
            "questions_by_iteration": {},
            "formatted_findings": f"Error: {error_msg}",
            "current_knowledge": "",
            "error": error_msg,
        }
