"""Python code generator for LangGraph workflows.

This module generates executable Python code from LangGraph workflow definitions,
allowing visual graphs to be exported as standalone scripts.
"""

from typing import Dict, List, Any
from pathlib import Path
import textwrap

from airunner.settings import AIRUNNER_LOG_LEVEL
from airunner.utils.application import get_logger


logger = get_logger(__name__, AIRUNNER_LOG_LEVEL)


class LangGraphCodeGenerator:
    """Generate Python code from LangGraph workflow definitions.

    This class takes a workflow structure and generates clean,
    executable Python code that can be saved and run independently.
    """

    def __init__(
        self,
        workflow_name: str = "agent_workflow",
        state_class_name: str = "AgentState",
    ):
        """Initialize the code generator.

        Args:
            workflow_name: Name for the workflow
            state_class_name: Name for the state class
        """
        self.workflow_name = workflow_name
        self.state_class_name = state_class_name

    def generate(
        self,
        nodes: Dict[str, Dict[str, Any]],
        edges: List[tuple],
        conditional_edges: List[Dict[str, Any]],
        state_fields: Dict[str, str],
        entry_point: str,
    ) -> str:
        """Generate complete Python code.

        Args:
            nodes: Dict of node_name: node_config
            edges: List of (source, target) tuples
            conditional_edges: List of conditional edge configs
            state_fields: Dict of field_name: type_string
            entry_point: Entry point node name

        Returns:
            Generated Python code
        """
        code_parts = [
            self._generate_header(),
            self._generate_imports(),
            self._generate_state_class(state_fields),
            self._generate_node_functions(nodes),
            self._generate_conditional_functions(conditional_edges),
            self._generate_graph_construction(
                nodes, edges, conditional_edges, entry_point
            ),
            self._generate_main_execution(),
        ]

        code = "\n\n\n".join(code_parts)
        return code

    def _generate_header(self) -> str:
        """Generate file header with docstring."""
        return f'''"""LangGraph workflow: {self.workflow_name}.

Auto-generated by AI Runner LangGraph integration.
This file can be executed standalone or imported as a module.
"""'''

    def _generate_imports(self) -> str:
        """Generate necessary imports."""
        return textwrap.dedent(
            """
        from typing import TypedDict, List, Dict, Any, Optional
        from langgraph.graph import StateGraph, END
        import logging

        logging.basicConfig(level=logging.INFO)
        logger = logging.getLogger(__name__)
        """
        ).strip()

    def _generate_state_class(self, state_fields: Dict[str, str]) -> str:
        """Generate state class definition.

        Args:
            state_fields: Dict of field_name: type_string

        Returns:
            State class code
        """
        lines = [f"class {self.state_class_name}(TypedDict):"]
        lines.append(f'    """State for {self.workflow_name} workflow."""')
        lines.append("")

        for field_name, field_type in state_fields.items():
            lines.append(f"    {field_name}: {field_type}")

        return "\n".join(lines)

    def _generate_node_functions(
        self, nodes: Dict[str, Dict[str, Any]]
    ) -> str:
        """Generate node function definitions.

        Args:
            nodes: Dict of node_name: node_config

        Returns:
            Node functions code
        """
        functions = []

        for node_name, node_config in nodes.items():
            func_code = self._generate_node_function(node_name, node_config)
            functions.append(func_code)

        return "\n\n".join(functions)

    def _generate_node_function(
        self, node_name: str, node_config: Dict[str, Any]
    ) -> str:
        """Generate a single node function.

        Args:
            node_name: Name of the node
            node_config: Node configuration

        Returns:
            Function code
        """
        func_name = self._sanitize_name(node_name)
        node_type = node_config.get("type", "custom")
        description = node_config.get("description", f"{node_name} node")

        lines = [
            f"def {func_name}(state: {self.state_class_name}) -> {self.state_class_name}:",
            f'    """{description}"""',
            f'    logger.info("Executing node: {node_name}")',
            "",
        ]

        # Add node-specific logic
        if node_type == "llm":
            lines.extend(self._generate_llm_logic(node_config))
        elif node_type == "rag":
            lines.extend(self._generate_rag_logic(node_config))
        elif node_type == "tool":
            lines.extend(self._generate_tool_logic(node_config))
        elif node_type == "transform":
            lines.extend(self._generate_transform_logic(node_config))
        else:
            lines.extend(self._generate_custom_logic(node_config))

        lines.append("    return state")

        return "\n".join(lines)

    def _generate_llm_logic(self, config: Dict[str, Any]) -> List[str]:
        """Generate LLM node logic."""
        return [
            "    # LLM call logic",
            "    # TODO: Implement LLM call",
            '    state["messages"].append("LLM response")',
        ]

    def _generate_rag_logic(self, config: Dict[str, Any]) -> List[str]:
        """Generate RAG node logic."""
        return [
            "    # RAG search logic",
            "    # TODO: Implement RAG search",
            '    state["rag_context"] = "Retrieved context"',
        ]

    def _generate_tool_logic(self, config: Dict[str, Any]) -> List[str]:
        """Generate tool node logic."""
        tool_name = config.get("tool_name", "unknown")
        return [
            f"    # Execute tool: {tool_name}",
            "    # TODO: Implement tool execution",
            '    state["tool_results"] = []',
        ]

    def _generate_transform_logic(self, config: Dict[str, Any]) -> List[str]:
        """Generate transform node logic."""
        return [
            "    # Transform state",
            "    # TODO: Implement state transformation",
            "    pass",
        ]

    def _generate_custom_logic(self, config: Dict[str, Any]) -> List[str]:
        """Generate custom node logic."""
        custom_code = config.get("code", "")
        if custom_code:
            return ["    " + line for line in custom_code.split("\n")]
        return ["    # Custom logic", "    pass"]

    def _generate_conditional_functions(
        self, conditional_edges: List[Dict[str, Any]]
    ) -> str:
        """Generate conditional routing functions.

        Args:
            conditional_edges: List of conditional edge configs

        Returns:
            Conditional functions code
        """
        if not conditional_edges:
            return ""

        functions = []
        for i, edge in enumerate(conditional_edges):
            func_code = self._generate_conditional_function(i, edge)
            functions.append(func_code)

        return "\n\n".join(functions)

    def _generate_conditional_function(
        self, index: int, edge_config: Dict[str, Any]
    ) -> str:
        """Generate a single conditional function."""
        source = edge_config.get("source", "unknown")
        condition_code = edge_config.get("condition_code", "")

        lines = [
            f"def condition_{index}(state: {self.state_class_name}) -> str:",
            f'    """Routing condition for {source}."""',
        ]

        if condition_code:
            lines.extend(
                ["    " + line for line in condition_code.split("\n")]
            )
        else:
            lines.extend(
                [
                    "    # Default routing based on next_action",
                    '    return state.get("next_action", "END")',
                ]
            )

        return "\n".join(lines)

    def _generate_graph_construction(
        self,
        nodes: Dict[str, Dict[str, Any]],
        edges: List[tuple],
        conditional_edges: List[Dict[str, Any]],
        entry_point: str,
    ) -> str:
        """Generate graph construction code.

        Args:
            nodes: Dict of node_name: node_config
            edges: List of (source, target) tuples
            conditional_edges: List of conditional edge configs
            entry_point: Entry point node name

        Returns:
            Graph construction code
        """
        lines = [
            "# Build the workflow graph",
            f"workflow = StateGraph({self.state_class_name})",
            "",
            "# Add nodes",
        ]

        for node_name in nodes.keys():
            func_name = self._sanitize_name(node_name)
            lines.append(f'workflow.add_node("{node_name}", {func_name})')

        lines.append("")
        lines.append("# Add edges")

        for source, target in edges:
            if target == "END":
                lines.append(f'workflow.add_edge("{source}", END)')
            else:
                lines.append(f'workflow.add_edge("{source}", "{target}")')

        if conditional_edges:
            lines.append("")
            lines.append("# Add conditional edges")

            for i, edge in enumerate(conditional_edges):
                source = edge.get("source")
                mapping = edge.get("mapping", {})

                # Format mapping
                mapping_str = "{"
                for key, value in mapping.items():
                    if value == "END":
                        mapping_str += f'"{key}": END, '
                    else:
                        mapping_str += f'"{key}": "{value}", '
                mapping_str += "}"

                lines.append(
                    f'workflow.add_conditional_edges("{source}", condition_{i}, {mapping_str})'
                )

        lines.append("")
        lines.append("# Set entry point and compile")
        lines.append(f'workflow.set_entry_point("{entry_point}")')
        lines.append("app = workflow.compile()")

        return "\n".join(lines)

    def _generate_main_execution(self) -> str:
        """Generate main execution block."""
        return textwrap.dedent(
            f'''
        def run_{self.workflow_name}(initial_state: {self.state_class_name}) -> {self.state_class_name}:
            """Execute the workflow.
            
            Args:
                initial_state: Initial state for the workflow
                
            Returns:
                Final state after workflow execution
            """
            logger.info("Starting workflow execution")
            result = app.invoke(initial_state)
            logger.info("Workflow execution completed")
            return result


        if __name__ == "__main__":
            # Example execution
            initial_state = {{
                "messages": [],
                "next_action": "",
                "error": None,
                "metadata": {{}},
            }}
            
            result = run_{self.workflow_name}(initial_state)
            print("Workflow result:", result)
        '''
        ).strip()

    @staticmethod
    def _sanitize_name(name: str) -> str:
        """Sanitize node name to valid Python identifier.

        Args:
            name: Node name

        Returns:
            Valid Python identifier
        """
        # Replace invalid characters
        sanitized = "".join(
            c if c.isalnum() or c == "_" else "_" for c in name
        )

        # Ensure it doesn't start with a number
        if sanitized and sanitized[0].isdigit():
            sanitized = f"node_{sanitized}"

        return sanitized or "unnamed_node"

    def save(self, code: str, output_path: Path) -> None:
        """Save generated code to file.

        Args:
            code: Generated Python code
            output_path: Path to save file
        """
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(code)
        logger.info(f"Code saved to: {output_path}")

    def format_code(self, code: str) -> str:
        """Format code with Black (if available).

        Args:
            code: Python code to format

        Returns:
            Formatted code
        """
        try:
            import black

            mode = black.FileMode(line_length=79)
            formatted = black.format_str(code, mode=mode)
            return formatted
        except ImportError:
            logger.warning("Black not available, returning unformatted code")
            return code
        except Exception as e:
            logger.error(f"Code formatting error: {e}")
            return code
