# SPDX-License-Identifier: Apache-2.0.
# Copyright (c) 2024 - 2025 Waldiez and contributors.
"""Generate the whole folw content."""

from typing import Any

from ..core import (
    ContentGenerator,
    ExporterContentError,
    ExporterContext,
    ExportPosition,
    ExportResult,
    PositionedContent,
    get_comment,
)
from .execution_generator import ExecutionGenerator
from .utils.common import generate_header


class FileGenerator(ContentGenerator):
    """Generate the complete flow notebook content."""

    def __init__(
        self,
        context: ExporterContext,
    ) -> None:
        """Initialize the notebook generator.

        Parameters
        ----------
        context : ExporterContext
            The exporter context containing dependencies.
        """
        self.context = context
        self.config = context.get_config()

    # pylint: disable=too-many-locals
    def generate(
        self,
        merged_result: ExportResult,
        is_async: bool,
        after_run: str,
        **kwargs: Any,
    ) -> str:
        """Generate the complete flow notebook content.

        Parameters
        ----------
        merged_result : ExportResult
            The merged export result containing all content.
        is_async : bool
            Whether to generate async conten
        after_run : str
            Additional content to add after the main flow execution.
        **kwargs : Any
            Additional keyword arguments for the generator.

        Returns
        -------
        str
            The complete flow notebook content.

        Raises
        ------
        ExporterContentError
            If there is no content to export.
        """
        # 1. Generate header
        header = self.get_header(merged_result)

        # 2. Generate imports
        imports_section = merged_result.get_content_by_position(
            ExportPosition.IMPORTS
        )

        # 3. Generate content sections
        tools_section = merged_result.get_content_by_position(
            ExportPosition.TOOLS
        )
        models_section = merged_result.get_content_by_position(
            ExportPosition.MODELS
        )
        agents_section = merged_result.get_content_by_position(
            ExportPosition.AGENTS,
            # Skip agent arguments (should already be there)
            skip_agent_arguments=True,
        )
        chats_content = merged_result.get_content_by_position(
            ExportPosition.CHATS
        )
        if not chats_content:
            raise ExporterContentError(
                "No content to export. Please ensure that the flow has chats."
            )
        after_chats = merged_result.get_content_by_position(
            ExportPosition.BOTTOM
        )

        main, call_main, execution_block = self._get_execution_content(
            chats_content=chats_content,
            is_async=is_async,
            after_run=after_run,
            for_notebook=self.config.for_notebook,
        )

        # 5. Combine everything
        everything: list[str] = [header]
        if imports_section:
            comment = get_comment(
                "Imports",
                for_notebook=self.config.for_notebook,
            )
            everything.append(comment)
            everything.append(
                "\n".join([entry.content for entry in imports_section])
            )
        if tools_section:
            comment = get_comment(
                "Tools",
                for_notebook=self.config.for_notebook,
            )
            everything.append(comment)
            everything.append(
                "\n".join([entry.content for entry in tools_section]) + "\n"
            )
        if models_section:
            comment = get_comment(
                "Models",
                for_notebook=self.config.for_notebook,
            )
            everything.append(comment)
            everything.append(
                "\n\n".join([entry.content for entry in models_section]) + "\n"
            )
        if agents_section:
            comment = get_comment(
                "Agents",
                for_notebook=self.config.for_notebook,
            )
            if self.config.for_notebook:
                comment += "# pyright: reportUnnecessaryIsInstance=false\n"
            everything.append(comment)
            everything.append(
                "\n\n".join([entry.content for entry in agents_section]) + "\n"
            )
        everything.append(main)
        if after_chats:
            everything.append(
                "\n".join([entry.content for entry in after_chats])
            )
        everything.append(call_main)
        if execution_block:
            everything.append(execution_block)

        return "\n".join(everything)

    def _get_execution_content(
        self,
        chats_content: list[PositionedContent],
        is_async: bool,
        for_notebook: bool,
        after_run: str,
    ) -> tuple[str, str, str]:
        cache_seed = (
            self.context.config.cache_seed if self.context.config else None
        )
        execution_gen = ExecutionGenerator()
        chat_contents = "\n".join(chat.content for chat in chats_content)
        main = execution_gen.generate_main_function(
            content=chat_contents,
            is_async=is_async,
            for_notebook=for_notebook,
            cache_seed=cache_seed,
            after_run=after_run,
        )
        call_main = execution_gen.generate_call_main_function(
            is_async=is_async,
            for_notebook=for_notebook,
        )
        execution_block = (
            execution_gen.generate_execution_block(
                is_async=is_async,
            )
            if not for_notebook
            else ""
        )
        return main, call_main, execution_block

    def get_header(self, merged_result: ExportResult) -> str:
        """Get or generate the header for the script.

        Parameters
        ----------
        merged_result : ExportResult
            The merged export result containing all content.

        Returns
        -------
        str
            The header content.
        """
        from_result = merged_result.get_content_by_position(ExportPosition.TOP)
        if not from_result:
            return generate_header(
                name=self.config.name,
                description=self.config.description,
                requirements=self.config.requirements,
                tags=self.config.tags,
                for_notebook=self.config.for_notebook,
            )
        header_string = "\n".join(content.content for content in from_result)
        while not header_string.endswith("\n\n"):
            header_string += "\n"
        return header_string
