import importlib
from typing import Dict, List

from langchain_core.language_models import BaseChatModel

# 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, otherwise just uses the default.

        """
        self.model = llm or get_llm()
        # Use provided search_system or create a new one
        self.search_system = search_system or 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.
        """

        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 purpose
                parts = line.strip("- ").split("|")
                if len(parts) == 2:
                    current_section["subsections"].append(
                        {"name": parts[0].strip(), "purpose": parts[1].strip()}
                    )

        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 = {}

        for section in structure:
            print(f"Processing section: {section['name']}")
            section_content = []
            section_content.append(f"# {section['name']}\n")

            # Process each subsection by directly researching it
            for subsection in section["subsections"]:
                # Add subsection header
                section_content.append(f"## {subsection['name']}\n")
                section_content.append(f"_{subsection['purpose']}_\n\n")

                # Generate a specific search query for this subsection
                subsection_query = f"{query} {section['name']} {subsection['name']} {subsection['purpose']}"

                print(
                    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 (
                    "current_knowledge" in subsection_results
                    and subsection_results["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
        from datetime import datetime

        metadata = {
            "generated_at": datetime.utcnow().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
