import importlib
from typing import Dict, List
from datetime import datetime, UTC

from langchain_core.language_models import BaseChatModel
from loguru import logger

# Fix circular import by importing directly from source modules
from .config.llm_config import get_llm
from .search_system import AdvancedSearchSystem
from .utilities import search_utilities


def get_report_generator(search_system=None):
    """Return an instance of the report generator with default settings.

    Args:
        search_system: Optional existing AdvancedSearchSystem to use
    """
    return IntegratedReportGenerator(search_system=search_system)


class IntegratedReportGenerator:
    def __init__(
        self,
        searches_per_section: int = 2,
        search_system=None,
        llm: BaseChatModel | None = None,
    ):
        """
        Args:
            searches_per_section: Number of searches to perform for each
                section in the report.
            search_system: Custom search system to use, otherwise just uses
                the default.
            llm: Custom LLM to use. Required if search_system is not provided.

        """
        # If search_system is provided, use its LLM; otherwise use the provided LLM
        if search_system:
            self.search_system = search_system
            self.model = llm or search_system.model
        elif llm:
            self.model = llm
            self.search_system = AdvancedSearchSystem(llm=self.model)
        else:
            # Fallback for backwards compatibility - will only work with auth
            self.model = get_llm()
            self.search_system = AdvancedSearchSystem(llm=self.model)

        self.searches_per_section = (
            searches_per_section  # Control search depth per section
        )

    def generate_report(self, initial_findings: Dict, query: str) -> Dict:
        """Generate a complete research report with section-specific research."""

        # Step 1: Determine structure
        structure = self._determine_report_structure(initial_findings, query)

        # Step 2: Research and generate content for each section in one step
        sections = self._research_and_generate_sections(
            initial_findings, structure, query
        )

        # Step 3: Format final report
        report = self._format_final_report(sections, structure, query)

        return report

    def _determine_report_structure(
        self, findings: Dict, query: str
    ) -> List[Dict]:
        """Analyze content and determine optimal report structure."""
        combined_content = findings["current_knowledge"]
        prompt = f"""
        Analyze this research content about: {query}

        Content Summary:
        {combined_content[:1000]}... [truncated]

        Determine the most appropriate report structure by:
        1. Analyzing the type of content (technical, business, academic, etc.)
        2. Identifying main themes and logical groupings
        3. Considering the depth and breadth of the research

        Return a table of contents structure in this exact format:
        STRUCTURE
        1. [Section Name]
           - [Subsection] | [purpose]
        2. [Section Name]
           - [Subsection] | [purpose]
        ...
        END_STRUCTURE

        Make the structure specific to the content, not generic.
        Each subsection must include its purpose after the | symbol.
        DO NOT include sections about sources, citations, references, or methodology.
        """

        response = search_utilities.remove_think_tags(
            self.model.invoke(prompt).content
        )

        # Parse the structure
        structure = []
        current_section = None

        for line in response.split("\n"):
            if line.strip() in ["STRUCTURE", "END_STRUCTURE"]:
                continue

            if line.strip().startswith(tuple("123456789")):
                # Main section
                section_name = line.split(".")[1].strip()
                current_section = {"name": section_name, "subsections": []}
                structure.append(current_section)
            elif line.strip().startswith("-") and current_section:
                # Subsection with or without purpose
                parts = line.strip("- ").split("|")
                if len(parts) == 2:
                    current_section["subsections"].append(
                        {"name": parts[0].strip(), "purpose": parts[1].strip()}
                    )
                elif len(parts) == 1 and parts[0].strip():
                    # Subsection without purpose - add default
                    current_section["subsections"].append(
                        {
                            "name": parts[0].strip(),
                            "purpose": f"Provide detailed information about {parts[0].strip()}",
                        }
                    )

        # Check if the last section is source-related and remove it
        if structure:
            last_section = structure[-1]
            section_name_lower = last_section["name"].lower()
            source_keywords = [
                "source",
                "citation",
                "reference",
                "bibliography",
            ]

            # Only check the last section for source-related content
            if any(
                keyword in section_name_lower for keyword in source_keywords
            ):
                logger.info(
                    f"Removed source-related last section: {last_section['name']}"
                )
                structure = structure[:-1]

        return structure

    def _research_and_generate_sections(
        self,
        initial_findings: Dict,
        structure: List[Dict],
        query: str,
    ) -> Dict[str, str]:
        """Research and generate content for each section in one step."""
        sections = {}

        # Preserve questions from initial research to avoid repetition
        # This follows the same pattern as citation tracking (all_links_of_system)
        existing_questions = initial_findings.get("questions_by_iteration", {})
        if existing_questions:
            # Set questions on both search system and its strategy
            if hasattr(self.search_system, "questions_by_iteration"):
                self.search_system.questions_by_iteration = (
                    existing_questions.copy()
                )

            # More importantly, set it on the strategy which actually uses it
            if hasattr(self.search_system, "strategy") and hasattr(
                self.search_system.strategy, "questions_by_iteration"
            ):
                self.search_system.strategy.questions_by_iteration = (
                    existing_questions.copy()
                )
                logger.info(
                    f"Initialized strategy with {len(existing_questions)} iterations of previous questions"
                )

        for section in structure:
            logger.info(f"Processing section: {section['name']}")
            section_content = []

            section_content.append(f"# {section['name']}\n")

            # If section has no subsections, create one from the section itself
            if not section["subsections"]:
                # Parse section name for purpose
                if "|" in section["name"]:
                    parts = section["name"].split("|", 1)
                    section["subsections"] = [
                        {"name": parts[0].strip(), "purpose": parts[1].strip()}
                    ]
                else:
                    # No purpose provided - use section name as subsection
                    section["subsections"] = [
                        {
                            "name": section["name"],
                            "purpose": f"Provide comprehensive content for {section['name']}",
                        }
                    ]

            # Process each subsection by directly researching it
            for subsection in section["subsections"]:
                # Only add subsection header if there are multiple subsections
                if len(section["subsections"]) > 1:
                    section_content.append(f"## {subsection['name']}\n")
                    section_content.append(f"_{subsection['purpose']}_\n\n")

                # Get other subsections in this section for context
                other_subsections = [
                    f"- {s['name']}: {s['purpose']}"
                    for s in section["subsections"]
                    if s["name"] != subsection["name"]
                ]
                other_subsections_text = (
                    "\n".join(other_subsections)
                    if other_subsections
                    else "None"
                )

                # Get all other sections for broader context
                other_sections = [
                    f"- {s['name']}"
                    for s in structure
                    if s["name"] != section["name"]
                ]
                other_sections_text = (
                    "\n".join(other_sections) if other_sections else "None"
                )

                # Check if this is actually a section-level content (only one subsection, likely auto-created)
                is_section_level = len(section["subsections"]) == 1

                # Generate appropriate search query
                if is_section_level:
                    # Section-level prompt - more comprehensive
                    subsection_query = (
                        f"Research task: Create comprehensive content for the '{subsection['name']}' section in a report about '{query}'. "
                        f"Section purpose: {subsection['purpose']} "
                        f"\n"
                        f"Other sections in the report:\n{other_sections_text}\n"
                        f"\n"
                        f"This is a standalone section requiring comprehensive coverage of its topic. "
                        f"Provide a thorough exploration that may include synthesis of information from previous sections where relevant. "
                        f"Include unique insights, specific examples, and concrete data. "
                        f"Use tables to organize information where applicable. "
                        f"For conclusion sections: synthesize key findings and provide forward-looking insights. "
                        f"Build upon the research findings from earlier sections to create a cohesive narrative."
                    )
                else:
                    # Subsection-level prompt - more focused
                    subsection_query = (
                        f"Research task: Create content for subsection '{subsection['name']}' in a report about '{query}'. "
                        f"This subsection's purpose: {subsection['purpose']} "
                        f"Part of section: '{section['name']}' "
                        f"\n"
                        f"Other sections in the report:\n{other_sections_text}\n"
                        f"\n"
                        f"Other subsections in this section will cover:\n{other_subsections_text}\n"
                        f"\n"
                        f"Focus ONLY on information specific to your subsection's purpose. "
                        f"Include unique details, specific examples, and concrete data. "
                        f"Use tables to organize information where applicable. "
                        f"IMPORTANT: Avoid repeating information that would logically be covered in other sections - focus on what makes this subsection unique. "
                        f"Previous research exists - find specific angles for this subsection."
                    )

                logger.info(
                    f"Researching subsection: {subsection['name']} with query: {subsection_query}"
                )

                # Configure search system for focused search
                original_max_iterations = self.search_system.max_iterations
                self.search_system.max_iterations = 1  # Keep search focused

                # Perform search for this subsection
                subsection_results = self.search_system.analyze_topic(
                    subsection_query
                )

                # Restore original iterations setting
                self.search_system.max_iterations = original_max_iterations

                # Add the researched content for this subsection
                if subsection_results.get("current_knowledge"):
                    section_content.append(
                        subsection_results["current_knowledge"]
                    )
                else:
                    section_content.append(
                        "*Limited information was found for this subsection.*\n"
                    )

                section_content.append("\n\n")

            # Combine all content for this section
            sections[section["name"]] = "\n".join(section_content)

        return sections

    def _generate_sections(
        self,
        initial_findings: Dict,
        section_research: Dict[str, List[Dict]],
        structure: List[Dict],
        query: str,
    ) -> Dict[str, str]:
        """
        This method is kept for compatibility but no longer used.
        The functionality has been moved to _research_and_generate_sections.
        """
        return {}

    def _format_final_report(
        self,
        sections: Dict[str, str],
        structure: List[Dict],
        query: str,
    ) -> Dict:
        """Format the final report with table of contents and sections."""
        # Generate TOC
        toc = ["# Table of Contents\n"]
        for i, section in enumerate(structure, 1):
            toc.append(f"{i}. **{section['name']}**")
            for j, subsection in enumerate(section["subsections"], 1):
                toc.append(
                    f"   {i}.{j} {subsection['name']} | _{subsection['purpose']}_"
                )

        # Combine TOC and sections
        report_parts = ["\n".join(toc), ""]

        # Add a summary of the research
        report_parts.append("# Research Summary")
        report_parts.append(
            "This report was researched using an advanced search system."
        )
        report_parts.append(
            "Research included targeted searches for each section and subsection."
        )
        report_parts.append("\n---\n")

        # Add each section's content
        for section in structure:
            if section["name"] in sections:
                report_parts.append(sections[section["name"]])
                report_parts.append("")

        # Format links from search system
        # Get utilities module dynamically to avoid circular imports
        utilities = importlib.import_module("local_deep_research.utilities")
        formatted_all_links = (
            utilities.search_utilities.format_links_to_markdown(
                all_links=self.search_system.all_links_of_system
            )
        )

        # Create final report with all parts
        final_report_content = "\n\n".join(report_parts)
        final_report_content = (
            final_report_content + "\n\n## Sources\n\n" + formatted_all_links
        )

        # Create metadata dictionary
        metadata = {
            "generated_at": datetime.now(UTC).isoformat(),
            "initial_sources": len(self.search_system.all_links_of_system),
            "sections_researched": len(structure),
            "searches_per_section": self.searches_per_section,
            "query": query,
        }

        # Return both content and metadata
        return {"content": final_report_content, "metadata": metadata}

    def _generate_error_report(self, query: str, error_msg: str) -> str:
        error_report = (
            f"=== ERROR REPORT ===\nQuery: {query}\nError: {error_msg}"
        )
        return error_report
