"""
LLM-Driven Modular Strategy with intelligent constraint processing and early rejection.
"""

import asyncio
import json
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

from loguru import logger

from ...database.thread_local_session import thread_cleanup
from ...utilities.json_utils import extract_json

from ..candidate_exploration import AdaptiveExplorer
from ..constraint_checking import DualConfidenceChecker
from ..constraints import ConstraintAnalyzer
from ..questions import StandardQuestionGenerator
from .base_strategy import BaseSearchStrategy


@dataclass
class CandidateConfidence:
    """Track candidate confidence levels for early rejection"""

    candidate: object
    positive_confidence: float
    negative_confidence: float
    rejection_reason: Optional[str] = None
    should_continue: bool = True


class LLMConstraintProcessor:
    """LLM-driven intelligent constraint processing"""

    def __init__(self, model):
        self.model = model

    async def decompose_constraints_intelligently(self, constraints):
        """Let LLM intelligently break down constraints into searchable elements"""
        constraint_text = "\n".join([f"- {c.description}" for c in constraints])

        prompt = f"""
        I have these constraints from a search query:
        {constraint_text}

        Please intelligently decompose these constraints into atomic, searchable elements that can be combined in different ways.

        For each constraint, provide:
        1. **Atomic elements** - Break it into smallest meaningful parts
        2. **Variations** - Different ways to express the same concept
        3. **Granular specifics** - Specific values, years, numbers, etc.

        Example for "TV show aired between 1960s and 1980s":
        - Atomic elements: ["TV show", "television", "series", "program"]
        - Time variations: ["1960s", "1970s", "1980s", "60s", "70s", "80s"]
        - Granular years: ["1960", "1961", "1962", "1963", "1964", "1965", "1966", "1967", "1968", "1969", "1970", "1971", "1972", "1973", "1974", "1975", "1976", "1977", "1978", "1979", "1980", "1981", "1982", "1983", "1984", "1985", "1986", "1987", "1988", "1989"]

        Example for "fewer than 50 episodes":
        - Atomic elements: ["episodes", "installments", "parts"]
        - Quantity variations: ["under 50", "less than 50", "limited run", "short series"]
        - Granular numbers: ["13 episodes", "26 episodes", "39 episodes", "single season"]

        Return as valid JSON format:
        {{
            "constraint_1": {{
                "atomic_elements": [...],
                "variations": [...],
                "granular_specifics": [...]
            }},
            "constraint_2": {{
                "atomic_elements": [...],
                "variations": [...],
                "granular_specifics": [...]
            }}
        }}
        """

        response = await self.model.ainvoke(prompt)
        return self._parse_decomposition(response.content)

    async def generate_intelligent_combinations(self, decomposed_constraints):
        """LLM generates smart combinations of atomic elements"""

        # Flatten all elements for the LLM to see
        all_elements = dict(decomposed_constraints.items())

        prompt = f"""
        I have decomposed constraints into these atomic elements:
        {json.dumps(all_elements, indent=2)}

        Now intelligently combine these elements to create targeted search queries. Be creative and systematic:

        1. **Year-by-year combinations**: Take specific years and combine with other specifics
           Example: "1960 TV show 13 episodes", "1961 television 26 episodes", etc.

        2. **Cross-constraint combinations**: Mix elements from different constraints
           Example: "humor ascetic 1970s", "fourth wall short series vintage"

        3. **Granular progression**: Create systematic progressions
           Example: "1960 comedy", "1961 comedy", "1962 comedy"...

        4. **Semantic variations**: Same meaning, different words
           Example: "brief TV run 1970s" vs "short television series seventies"

        5. **Contextual combinations**: Add implied context
           Example: "monk-trained character 1978 television"

        Generate 60-80 diverse search combinations that would maximize finding the target.
        Focus on being comprehensive yet targeted.

        Return as a valid JSON list of search queries:
        ["query1", "query2", "query3"]
        """

        response = await self.model.ainvoke(prompt)
        return self._parse_combinations(response.content)

    async def generate_creative_search_angles(
        self, original_query, decomposed_constraints
    ):
        """LLM generates completely creative search approaches"""

        prompt = f"""
        Original query: "{original_query}"

        Now think like a detective - what are ALL the different ways someone might search for this character?
        Be extremely creative and think outside the box:

        1. **Character name guessing**: What names might this character have?
        2. **Show title guessing**: What might the TV show be called?
        3. **Cultural context**: What was happening in those decades?
        4. **Genre searches**: What genre/category would this fit?
        5. **Indirect searches**: What related topics might lead to this?
        6. **Reverse searches**: Start from known similar characters
        7. **Archetype searches**: What type of character is this?
        8. **Creator/studio searches**: Who might have made this?

        Generate 30-40 creative search angles that approach this from completely different directions.

        Examples of creative thinking:
        - "1970s cartoon characters who talk to camera"
        - "superhero trained by monks television"
        - "vintage comedy shows cancelled after one season"
        - "fourth wall breaking animation 70s"
        - "spiritual mentor origin story TV characters"
        - "Plastic Man TV show episodes"
        - "elastic superhero television series"

        Return as valid JSON list of creative searches:
        ["creative_query1", "creative_query2"]
        """

        response = await self.model.ainvoke(prompt)
        return self._parse_creative_searches(response.content)

    async def optimize_search_combinations(self, all_combinations):
        """LLM optimizes the search list for maximum effectiveness"""

        prompt = f"""
        I have generated {len(all_combinations)} search combinations. Here are the first 20:
        {json.dumps(all_combinations[:20], indent=2)}

        Please optimize this search strategy by organizing searches by priority and effectiveness:

        1. **Remove redundant searches** that are too similar
        2. **Prioritize high-value searches** likely to find results
        3. **Balance specificity vs breadth**
        4. **Add missing search angles** you notice
        5. **Organize by search strategy type**

        Return optimized searches organized by category as valid JSON:
        {{
            "high_priority": ["most likely to succeed - top 15 searches"],
            "systematic_granular": ["year-by-year, episode-by-episode combinations - 20 searches"],
            "creative_angles": ["outside-the-box approaches - 15 searches"],
            "contextual_searches": ["time period + cultural context - 15 searches"],
            "fallback_broad": ["broader searches if specifics fail - 10 searches"]
        }}
        """

        response = await self.model.ainvoke(prompt)
        return self._parse_optimized_searches(response.content)

    def _parse_decomposition(self, content):
        """Parse LLM decomposition response"""
        result = extract_json(content, expected_type=dict)
        if result is not None:
            return result

        # Fallback to simple structure
        return {
            "time_constraint": {
                "atomic_elements": ["TV show", "television", "series"],
                "variations": ["1960s", "1970s", "1980s"],
                "granular_specifics": [str(year) for year in range(1960, 1990)],
            }
        }

    def _parse_combinations(self, content):
        """Parse LLM combinations response"""
        result = extract_json(content, expected_type=list)
        if result is not None:
            return result

        # Fallback
        return [
            "fictional character humor",
            "TV show 1970s",
            "fourth wall breaking",
        ]

    def _parse_creative_searches(self, content):
        """Parse LLM creative searches response"""
        result = extract_json(content, expected_type=list)
        if result is not None:
            return result

        # Fallback
        return [
            "vintage cartoon character",
            "superhero TV show 1970s",
            "comedy series short run",
        ]

    def _parse_optimized_searches(self, content):
        """Parse LLM optimized searches response"""
        result = extract_json(content, expected_type=dict)
        if result is not None:
            return result

        # Fallback
        return {
            "high_priority": [
                "fictional character fourth wall humor",
                "1970s TV show limited episodes",
            ],
            "systematic_granular": [
                "1970 TV show",
                "1971 TV show",
                "1972 TV show",
            ],
            "creative_angles": [
                "superhero comedy television",
                "cartoon character talks to audience",
            ],
            "contextual_searches": [
                "vintage TV comedy",
                "classic television humor",
            ],
            "fallback_broad": ["fictional character", "TV show character"],
        }


class EarlyRejectionManager:
    """Manages early rejection and confidence tracking"""

    def __init__(self, model, positive_threshold=0.6, negative_threshold=0.3):
        self.model = model
        self.positive_threshold = positive_threshold
        self.negative_threshold = negative_threshold
        self.rejected_candidates = set()

    async def quick_confidence_check(self, candidate, constraints):
        """Quick confidence assessment for early rejection"""

        prompt = f"""
        Quickly assess if this candidate matches the search criteria:

        Candidate: {candidate.name}
        Available info: {getattr(candidate, "metadata", {})}

        Constraints to match:
        {[c.description for c in constraints]}

        Provide:
        1. **Positive confidence** (0.0-1.0): How likely this candidate matches
        2. **Negative confidence** (0.0-1.0): How likely this candidate does NOT match
        3. **Quick reasoning**: Brief explanation

        Return as JSON:
        {{
            "positive_confidence": 0.X,
            "negative_confidence": 0.X,
            "reasoning": "brief explanation"
        }}
        """

        try:
            response = await self.model.ainvoke(prompt)
            return self._parse_confidence(response.content)
        except Exception:
            logger.exception("Quick confidence check failed")
            return {
                "positive_confidence": 0.5,
                "negative_confidence": 0.3,
                "reasoning": "fallback",
            }

    def should_reject_early(self, confidence_result):
        """Determine if candidate should be rejected early"""
        positive = confidence_result.get("positive_confidence", 0.5)
        negative = confidence_result.get("negative_confidence", 0.3)

        # Reject if high negative confidence or very low positive confidence
        if negative > 0.7 or positive < 0.1:
            return (
                True,
                f"High negative confidence ({negative:.2f}) or low positive ({positive:.2f})",
            )

        return False, None

    def should_continue_search(self, all_candidates, high_confidence_count):
        """Determine if we should continue searching"""
        # Stop if we have enough high-confidence candidates
        if high_confidence_count >= 5:
            return False, "Found sufficient high-confidence candidates"

        # Stop if we have many candidates but low quality
        if len(all_candidates) > 50 and high_confidence_count == 0:
            return False, "Too many low-quality candidates"

        return True, None

    def _parse_confidence(self, content):
        """Parse confidence assessment"""
        result = extract_json(content, expected_type=dict)
        if result is not None:
            return result

        return {
            "positive_confidence": 0.5,
            "negative_confidence": 0.3,
            "reasoning": "parse_error",
        }


class LLMDrivenModularStrategy(BaseSearchStrategy):
    """
    LLM-driven modular strategy with intelligent constraint processing and early rejection.
    """

    def __init__(
        self,
        model,
        search,
        all_links_of_system=None,
        constraint_checker_type: str = "dual_confidence",
        exploration_strategy: str = "adaptive",
        early_rejection: bool = True,
        **kwargs,
    ):
        super().__init__(all_links_of_system=all_links_of_system)

        self.model = model
        self.search_engine = search
        self.search_engines = getattr(search, "search_engines", [])

        # Initialize components
        self.constraint_analyzer = ConstraintAnalyzer(self.model)
        self.llm_processor = LLMConstraintProcessor(self.model)
        self.early_rejection_manager = (
            EarlyRejectionManager(self.model) if early_rejection else None
        )

        # Initialize constraint checker
        self.constraint_checker = DualConfidenceChecker(
            model=self.model,
            evidence_gatherer=self._gather_evidence_for_constraint,
            negative_threshold=0.25,
            positive_threshold=0.4,
            uncertainty_penalty=0.2,
            negative_weight=2.0,
        )

        # Initialize candidate explorer
        self.candidate_explorer = AdaptiveExplorer(
            search_engine=self.search_engine,
            model=self.model,
            learning_rate=0.1,
            max_search_time=45.0,  # Reduced since we have more searches
            max_candidates=30,  # Increased since we filter early
        )

        # Initialize question generator
        self.question_generator = StandardQuestionGenerator(model=self.model)

        # Strategy configuration
        self.constraint_checker_type = constraint_checker_type
        self.exploration_strategy = exploration_strategy
        self.early_rejection = early_rejection

        logger.info(
            f"Initialized LLMDrivenModularStrategy with {constraint_checker_type} checker, "
            f"{exploration_strategy} explorer, early_rejection={early_rejection}"
        )

    def analyze_topic(self, query: str) -> Dict:
        """Main entry point - sync wrapper for async search"""
        try:
            import asyncio

            # Create a new event loop if none exists or if the current loop is running
            try:
                loop = asyncio.get_event_loop()
                if loop.is_running():
                    # If we're already in an async context, run in a new thread
                    import concurrent.futures

                    with concurrent.futures.ThreadPoolExecutor() as executor:
                        future = executor.submit(
                            thread_cleanup(
                                lambda: asyncio.run(self.search(query))
                            )
                        )
                        answer, metadata = future.result()
                else:
                    # If not in async context, run directly
                    answer, metadata = loop.run_until_complete(
                        self.search(query)
                    )
            except RuntimeError:
                # No event loop, create one
                answer, metadata = asyncio.run(self.search(query))

            return {
                "findings": [{"content": answer}],
                "iterations": 1,
                "final_answer": answer,
                "metadata": metadata,
                "links": getattr(self, "all_links_of_system", []),
                "questions_by_iteration": getattr(
                    self, "questions_by_iteration", []
                ),
            }

        except Exception as e:
            logger.exception("Error in analyze_topic")
            return {
                "findings": [],
                "iterations": 0,
                "final_answer": f"Analysis failed: {e!s}",
                "metadata": {"error": str(e)},
                "links": [],
                "questions_by_iteration": [],
            }

    async def search(
        self,
        query: str,
        search_engines: List[str] = None,
        progress_callback=None,
        **kwargs,
    ) -> Tuple[str, Dict]:
        """Execute the LLM-driven modular search strategy"""
        try:
            logger.info(f"Starting LLM-driven modular search for: {query}")

            # Phase 1: Extract base constraints
            if progress_callback:
                progress_callback(
                    {
                        "phase": "constraint_analysis",
                        "progress": 5,
                        "message": "Analyzing query constraints",
                    }
                )

            base_constraints = self.constraint_analyzer.extract_constraints(
                query
            )
            logger.info(f"Extracted {len(base_constraints)} base constraints")

            # Phase 2: LLM intelligent decomposition
            if progress_callback:
                progress_callback(
                    {
                        "phase": "llm_decomposition",
                        "progress": 15,
                        "message": "LLM decomposing constraints intelligently",
                    }
                )

            decomposed = (
                await self.llm_processor.decompose_constraints_intelligently(
                    base_constraints
                )
            )
            logger.info(
                f"LLM decomposed constraints into {len(decomposed)} groups"
            )

            # Phase 3: LLM intelligent combinations
            if progress_callback:
                progress_callback(
                    {
                        "phase": "llm_combinations",
                        "progress": 25,
                        "message": "LLM generating intelligent search combinations",
                    }
                )

            intelligent_combinations = (
                await self.llm_processor.generate_intelligent_combinations(
                    decomposed
                )
            )
            logger.info(
                f"LLM generated {len(intelligent_combinations)} intelligent combinations"
            )

            # Phase 4: LLM creative search angles
            if progress_callback:
                progress_callback(
                    {
                        "phase": "llm_creative",
                        "progress": 35,
                        "message": "LLM generating creative search angles",
                    }
                )

            creative_searches = (
                await self.llm_processor.generate_creative_search_angles(
                    query, decomposed
                )
            )
            logger.info(
                f"LLM generated {len(creative_searches)} creative searches"
            )

            # Phase 5: LLM optimization
            if progress_callback:
                progress_callback(
                    {
                        "phase": "llm_optimization",
                        "progress": 45,
                        "message": "LLM optimizing search strategy",
                    }
                )

            all_searches = intelligent_combinations + creative_searches
            optimized_searches = (
                await self.llm_processor.optimize_search_combinations(
                    all_searches
                )
            )
            total_searches = sum(
                len(searches) for searches in optimized_searches.values()
            )
            logger.info(
                f"LLM optimized to {total_searches} total searches across categories"
            )

            # Phase 6: Execute searches by priority with early rejection
            all_candidates = []
            high_confidence_count = 0
            search_progress = 50

            for category, searches in optimized_searches.items():
                if not searches:
                    continue

                logger.info(
                    f"Executing {category} searches: {len(searches)} queries"
                )

                if progress_callback:
                    progress_callback(
                        {
                            "phase": f"search_{category}",
                            "progress": search_progress,
                            "message": f"Searching with {category} strategy",
                        }
                    )

                # Execute in parallel batches
                batch_size = 3 if category == "high_priority" else 5
                category_candidates = []

                for i in range(0, len(searches), batch_size):
                    batch = searches[i : i + batch_size]

                    # Execute batch searches in parallel
                    batch_tasks = []
                    for search_query in batch:
                        task = self.candidate_explorer._execute_search(
                            search_query
                        )
                        batch_tasks.append(task)

                    # Wait for batch completion
                    batch_results = await asyncio.gather(
                        *batch_tasks, return_exceptions=True
                    )

                    # Process batch results
                    for j, result in enumerate(batch_results):
                        if isinstance(result, Exception):
                            logger.error(
                                f"Search failed: {batch[j]} - {result}"
                            )
                            continue

                        candidates = self.candidate_explorer._extract_candidates_from_results(
                            result, entity_type="fictional character"
                        )

                        # Early rejection if enabled
                        if self.early_rejection_manager:
                            for candidate in candidates:
                                confidence = await self.early_rejection_manager.quick_confidence_check(
                                    candidate, base_constraints
                                )

                                should_reject, reason = (
                                    self.early_rejection_manager.should_reject_early(
                                        confidence
                                    )
                                )
                                if should_reject:
                                    logger.debug(
                                        f"Early rejected {candidate.name}: {reason}"
                                    )
                                    continue

                                if (
                                    confidence.get("positive_confidence", 0)
                                    > 0.6
                                ):
                                    high_confidence_count += 1

                                category_candidates.append(candidate)
                        else:
                            category_candidates.extend(candidates)

                    logger.info(
                        f"{category} batch {i // batch_size + 1}: found {len(category_candidates)} candidates"
                    )

                    # Early stopping check
                    if self.early_rejection_manager:
                        should_continue, stop_reason = (
                            self.early_rejection_manager.should_continue_search(
                                all_candidates + category_candidates,
                                high_confidence_count,
                            )
                        )
                        if not should_continue:
                            logger.info(f"Early stopping: {stop_reason}")
                            break

                all_candidates.extend(category_candidates)
                search_progress += 8  # Distribute remaining progress

                # Stop if we have enough high-confidence candidates
                if high_confidence_count >= 5:
                    logger.info(
                        "Found sufficient high-confidence candidates, stopping search"
                    )
                    break

            logger.info(
                f"Search completed: {len(all_candidates)} total candidates, {high_confidence_count} high-confidence"
            )

            # Phase 7: Constraint checking on remaining candidates
            if progress_callback:
                progress_callback(
                    {
                        "phase": "constraint_evaluation",
                        "progress": 85,
                        "message": f"Evaluating {len(all_candidates)} candidates",
                    }
                )

            if not all_candidates:
                return "No valid candidates found", {
                    "strategy": "llm_driven_modular",
                    "total_searches": total_searches,
                    "candidates_found": 0,
                    "high_confidence_count": 0,
                }

            # Evaluate top candidates (limit to avoid long processing)
            candidates_to_evaluate = all_candidates[:20]  # Top 20 candidates
            evaluated_candidates = []

            for i, candidate in enumerate(candidates_to_evaluate):
                try:
                    result = self.constraint_checker.check_candidate(
                        candidate, base_constraints
                    )

                    candidate.evaluation_results = result.detailed_results
                    candidate.score = result.total_score
                    candidate.should_reject = result.should_reject

                    if not result.should_reject:
                        evaluated_candidates.append(candidate)

                except Exception:
                    logger.exception(
                        f"Error evaluating candidate {candidate.name}"
                    )
                    continue

            if not evaluated_candidates:
                return "No valid candidates passed constraint evaluation", {
                    "strategy": "llm_driven_modular",
                    "total_searches": total_searches,
                    "candidates_found": len(all_candidates),
                    "candidates_evaluated": len(candidates_to_evaluate),
                    "high_confidence_count": high_confidence_count,
                }

            # Select best candidate
            evaluated_candidates.sort(key=lambda x: x.score, reverse=True)
            best_candidate = evaluated_candidates[0]

            logger.info(
                f"Best candidate: {best_candidate.name} with score {best_candidate.score:.2%}"
            )

            # Generate final answer
            if progress_callback:
                progress_callback(
                    {
                        "phase": "final_answer",
                        "progress": 95,
                        "message": "Generating final answer",
                    }
                )

            answer = await self._generate_final_answer(
                query, best_candidate, base_constraints
            )

            metadata = {
                "strategy": "llm_driven_modular",
                "constraint_checker": self.constraint_checker_type,
                "exploration_strategy": self.exploration_strategy,
                "early_rejection_enabled": self.early_rejection,
                "total_searches_generated": total_searches,
                "candidates_found": len(all_candidates),
                "candidates_evaluated": len(candidates_to_evaluate),
                "candidates_valid": len(evaluated_candidates),
                "high_confidence_count": high_confidence_count,
                "best_candidate": best_candidate.name,
                "best_score": best_candidate.score,
            }

            return answer, metadata

        except Exception:
            logger.exception("Error in LLM-driven search")
            return "Search failed", {"error": "Search failed"}

    async def _generate_final_answer(self, query, best_candidate, constraints):
        """Generate comprehensive final answer"""
        constraint_info = "\n".join([f"- {c.description}" for c in constraints])

        evaluation_info = ""
        if hasattr(best_candidate, "evaluation_results"):
            evaluation_info = "\n".join(
                [
                    f"- {result.get('constraint', 'Unknown')}: {result.get('score', 0):.0%}"
                    for result in best_candidate.evaluation_results
                ]
            )

        prompt = f"""Based on the search results, provide a comprehensive answer to: {query}

Best candidate found: {best_candidate.name}
Score: {best_candidate.score:.0%}

Constraints analyzed:
{constraint_info}

Constraint evaluation results:
{evaluation_info}

Evidence summary: {getattr(best_candidate, "summary", "No summary available")}

Provide a clear, factual answer that addresses the original question and explains how the candidate satisfies the constraints."""

        response = await self.model.ainvoke(prompt)
        return response.content

    def _gather_evidence_for_constraint(self, candidate, constraint):
        """Gather evidence for a constraint using actual search"""
        try:
            # Create a focused search query
            query = f"{candidate.name} {constraint.description}"

            # Use the search engine properly
            if hasattr(self.search_engine, "run"):
                results = self.search_engine.run(query)
            else:
                logger.warning("Search engine doesn't have run method")
                return []

            # Handle different result formats
            if isinstance(results, list):
                result_list = results[:3]  # Top 3 results
            elif isinstance(results, dict):
                result_list = results.get("results", [])[:3]  # Top 3 results
            else:
                logger.warning(f"Unknown search result format: {type(results)}")
                return []

            # Extract evidence from search results
            evidence = []
            for result in result_list:
                evidence.append(
                    {
                        "text": result.get("snippet", "")
                        or result.get("content", ""),
                        "source": result.get("url", "search_result"),
                        "confidence": 0.7,
                        "title": result.get("title", ""),
                    }
                )

            return evidence

        except Exception:
            logger.exception("Error gathering evidence")
            return []
