from typing import List

from langchain_core.language_models import BaseLLM
from loguru import logger

from .base_question import BaseQuestionGenerator


class DecompositionQuestionGenerator(BaseQuestionGenerator):
    """Question generator for decomposing complex queries into sub-queries."""

    def __init__(self, model: BaseLLM, max_subqueries: int = 5):
        """
        Initialize the question generator.

        Args:
            model: The language model to use for question generation
            max_subqueries: Maximum number of sub-queries to generate
        """
        super().__init__(model)
        self.max_subqueries = max_subqueries

    def generate_questions(
        self,
        query: str,
        context: str,
        **kwargs,
    ) -> List[str]:
        """
        Generate sub-queries by decomposing the original query.

        Args:
            query: The main research query
            context: Additional context for question generation
            **kwargs: Additional keyword arguments

        Returns:
            List of generated sub-queries
        """
        # Extract subject if the query is in question format
        subject = query
        lower_query = query.lower()

        if lower_query.endswith("?"):
            # Handle question-format queries by extracting the subject
            question_prefixes = [
                "what is",
                "what are",
                "how does",
                "how do",
                "how can",
                "why is",
                "why are",
                "when did",
                "where is",
                "which",
                "who is",
                "can",
                "will",
            ]

            # Remove the question mark
            subject_candidate = query[:-1].strip()

            # Check for common question beginnings and extract the subject
            for prefix in question_prefixes:
                if lower_query.startswith(prefix):
                    # Extract everything after the question prefix
                    subject_candidate = query[len(prefix) :].strip()
                    # Remove trailing ? if present
                    if subject_candidate.endswith("?"):
                        subject_candidate = subject_candidate[:-1].strip()
                    subject = subject_candidate
                    break

            # For compound questions, extract just the primary subject
            conjunctions = [
                " and ",
                " or ",
                " but ",
                " as ",
                " that ",
                " which ",
                " when ",
                " where ",
                " how ",
            ]
            for conjunction in conjunctions:
                if conjunction in subject.lower():
                    # Take only the part before the conjunction
                    subject = subject.split(conjunction)[0].strip()
                    logger.info(
                        f"Split compound question at '{conjunction}', extracted: '{subject}'"
                    )
                    break

            # Clean up the subject if it starts with articles
            for article in ["a ", "an ", "the "]:
                if subject.lower().startswith(article):
                    subject = subject[len(article) :].strip()

        logger.info(
            f"Original query: '{query}', Extracted subject: '{subject}'"
        )

        # Create a prompt to decompose the query into sub-questions
        prompt = f"""Decompose the main research topic into 3-5 specific sub-queries that can be answered independently.
Focus on breaking down complex concepts and identifying key aspects requiring separate investigation.
Ensure sub-queries are clear, targeted, and help build a comprehensive understanding.

Main Research Topic: {subject}
Original Query: {query}

Context Information:
{context[:2000]}  # Limit context length to prevent token limit issues

Your task is to create 3-5 specific questions that will help thoroughly research this topic.
If the original query is already a question, extract the core subject and formulate questions around that subject.

Return ONLY the sub-queries, one per line, without numbering or bullet points.
Example format:
What is X technology?
How does X compare to Y?
What are the security implications of X?
"""

        logger.info(
            f"Generating sub-questions for query: '{query}', subject: '{subject}'"
        )

        try:
            # Get response from LLM
            response = self.model.invoke(prompt)

            # Handle different response formats (string or object with content attribute)
            sub_queries_text = ""
            if hasattr(response, "content"):
                sub_queries_text = response.content.strip()
            else:
                # Handle string responses
                sub_queries_text = str(response).strip()

            # Check for the common "No language models available" error
            if (
                "No language models are available" in sub_queries_text
                or "Please install Ollama" in sub_queries_text
            ):
                logger.warning(
                    "LLM returned error about language models not being available, using default questions"
                )
                # Create topic-specific default questions based on the query
                return self._generate_default_questions(query)

            # Extract sub-queries (one per line)
            sub_queries = []
            for line in sub_queries_text.split("\n"):
                line = line.strip()
                # Skip empty lines and lines that are just formatting (bullets, numbers)
                if (
                    not line
                    or line in ["*", "-", "•"]
                    or line.startswith(("- ", "* ", "• ", "1. ", "2. ", "3. "))
                ):
                    continue

                # Remove any leading bullets or numbers if they exist
                clean_line = line
                for prefix in [
                    "- ",
                    "* ",
                    "• ",
                    "1. ",
                    "2. ",
                    "3. ",
                    "4. ",
                    "5. ",
                    "- ",
                    "#",
                ]:
                    if clean_line.startswith(prefix):
                        clean_line = clean_line[len(prefix) :]

                if (
                    clean_line and len(clean_line) > 10
                ):  # Ensure it's a meaningful question
                    sub_queries.append(clean_line)

            # If no sub-queries were extracted, try again with a simpler prompt
            if not sub_queries:
                logger.warning(
                    "No sub-queries extracted from first attempt, trying simplified approach"
                )

                # Determine if the query is already a question and extract the subject
                topic_text = query
                if query.lower().endswith("?"):
                    # Try to extract subject from question
                    for prefix in [
                        "what is",
                        "what are",
                        "how does",
                        "how can",
                        "why is",
                    ]:
                        if query.lower().startswith(prefix):
                            topic_text = query[len(prefix) :].strip()
                            if topic_text.endswith("?"):
                                topic_text = topic_text[:-1].strip()
                            break

                    # For compound topics, extract just the primary subject
                    conjunctions = [
                        " and ",
                        " or ",
                        " but ",
                        " as ",
                        " that ",
                        " which ",
                        " when ",
                        " where ",
                        " how ",
                    ]
                    for conjunction in conjunctions:
                        if conjunction in topic_text.lower():
                            # Take only the part before the conjunction
                            topic_text = topic_text.split(conjunction)[
                                0
                            ].strip()
                            logger.info(
                                f"Simplified prompt: Split compound query at '{conjunction}', extracted: '{topic_text}'"
                            )
                            break

                    # Clean up the topic if it starts with articles
                    for article in ["a ", "an ", "the "]:
                        if topic_text.lower().startswith(article):
                            topic_text = topic_text[len(article) :].strip()

                # Simpler prompt
                simple_prompt = f"""Break down this research topic into 3 simpler sub-questions:

Research Topic: {topic_text}
Original Query: {query}

Your task is to create 3 specific questions that will help thoroughly research this topic.
If the original query is already a question, use the core subject of that question.

Sub-questions:
1.
2.
3. """

                simple_response = self.model.invoke(simple_prompt)

                # Handle different response formats
                simple_text = ""
                if hasattr(simple_response, "content"):
                    simple_text = simple_response.content.strip()
                else:
                    simple_text = str(simple_response).strip()

                # Check again for language model errors
                if (
                    "No language models are available" in simple_text
                    or "Please install Ollama" in simple_text
                ):
                    logger.warning(
                        "LLM returned error in simplified prompt, using default questions"
                    )
                    return self._generate_default_questions(query)

                # Extract sub-queries from the simpler response
                for line in simple_text.split("\n"):
                    line = line.strip()
                    if (
                        line
                        and not line.startswith("Sub-questions:")
                        and len(line) > 10
                    ):
                        # Clean up numbering
                        for prefix in ["1. ", "2. ", "3. ", "- ", "* "]:
                            if line.startswith(prefix):
                                line = line[len(prefix) :]
                        sub_queries.append(line.strip())

            # If still no sub-queries, create default ones based on the original query
            if not sub_queries:
                logger.warning(
                    "Failed to generate meaningful sub-queries, using default decomposition"
                )
                return self._generate_default_questions(query)

            logger.info(
                f"Generated {len(sub_queries)} sub-questions: {sub_queries}"
            )
            return sub_queries[: self.max_subqueries]  # Limit to max_subqueries

        except Exception:
            logger.exception("Error generating sub-questions")
            # Fallback to basic questions in case of error
            return self._generate_default_questions(query)

    def _generate_default_questions(self, query: str) -> List[str]:
        """
        Generate default questions for a given query when LLM fails.

        Args:
            query: The main research query

        Returns:
            List of default questions
        """
        # Adjust questions based on the type of query
        query = query.strip()

        # Check if the query is already in question format
        question_prefixes = [
            "what is",
            "what are",
            "how does",
            "how do",
            "how can",
            "why is",
            "why are",
            "when did",
            "where is",
            "which",
            "who is",
            "can",
            "will",
        ]

        # Extract the subject from a question-format query
        subject = query
        lower_query = query.lower()

        # Check for common question formats and extract the subject
        if lower_query.endswith("?"):
            # Remove the question mark
            subject = query[:-1].strip()

            # Check for common question beginnings and extract the subject
            for prefix in question_prefixes:
                if lower_query.startswith(prefix):
                    # Extract everything after the question prefix
                    subject = query[len(prefix) :].strip()
                    # Remove trailing ? if present
                    if subject.endswith("?"):
                        subject = subject[:-1].strip()
                    break

        # For compound questions, extract just the primary subject
        # Look for conjunctions and prepositions that typically separate the subject from the rest
        conjunctions = [
            " and ",
            " or ",
            " but ",
            " as ",
            " that ",
            " which ",
            " when ",
            " where ",
            " how ",
        ]
        for conjunction in conjunctions:
            if conjunction in subject.lower():
                # Take only the part before the conjunction
                subject = subject.split(conjunction)[0].strip()
                logger.info(
                    f"Split compound question at '{conjunction}', extracted: '{subject}'"
                )
                break

        # Clean up the subject if it starts with articles
        for article in ["a ", "an ", "the "]:
            if subject.lower().startswith(article):
                subject = subject[len(article) :].strip()

        # For single word or very short subjects, adapt the question format
        is_short_subject = len(subject.split()) <= 2

        logger.info(
            f"Query: '{query}', Identified subject: '{subject}', Short subject: {is_short_subject}"
        )

        # Special case for CSRF - if we've extracted just "csrf" from a longer query
        if (
            subject.lower() == "csrf"
            or subject.lower() == "cross-site request forgery"
        ):
            # CSRF-specific questions
            default_questions = [
                "What is Cross-Site Request Forgery (CSRF)?",
                "How do CSRF attacks work and what are common attack vectors?",
                "What are effective CSRF prevention methods and best practices?",
                "How do CSRF tokens work to prevent attacks?",
                "What are real-world examples of CSRF vulnerabilities and their impact?",
            ]
        elif not subject:
            # Empty query case
            default_questions = [
                "What is the definition of this topic?",
                "What are the key aspects of this topic?",
                "What are practical applications of this concept?",
            ]
        elif any(
            term in subject.lower()
            for term in ["secure", "security", "vulnerability", "attack"]
        ):
            # Security-related questions
            default_questions = [
                f"What is {subject} and how does it work?",
                f"What are common {subject} vulnerabilities or attack vectors?",
                f"What are best practices for preventing {subject} issues?",
                f"How can {subject} be detected and mitigated?",
                f"What are real-world examples of {subject} incidents?",
            ]
        elif any(
            term in subject.lower()
            for term in ["programming", "language", "code", "software"]
        ):
            # Programming-related questions
            default_questions = [
                f"What is {subject} and how does it work?",
                f"What are the main features and advantages of {subject}?",
                f"What are common use cases and applications for {subject}?",
                f"How does {subject} compare to similar technologies?",
                f"What are best practices when working with {subject}?",
            ]
        elif is_short_subject:
            # For short subjects (1-2 words), use a dedicated format
            default_questions = [
                f"What is {subject}?",
                f"What are the main characteristics of {subject}?",
                f"How is {subject} used in practice?",
                f"What are the advantages and disadvantages of {subject}?",
                f"How has {subject} evolved over time?",
            ]
        else:
            # Generic questions for any topic
            default_questions = [
                f"What is the definition of {subject}?",
                f"What are the key components or features of {subject}?",
                f"What are common applications or use cases for {subject}?",
                f"What are the advantages and limitations of {subject}?",
                f"How does {subject} compare to alternatives?",
            ]

        logger.info(
            f"Using {len(default_questions)} default questions: {default_questions}"
        )
        return default_questions[: self.max_subqueries]
