#!/usr/bin/env python3
"""
Universal TreeShell Tool Generator

Generates HEAVEN tools for any domain-specific TreeShell application.
Creates {domain}_agent_tree_shell_tool.py files in heaven_data_dir/tools.
"""

import os
from pathlib import Path
from typing import Dict, Any, Optional, List
from datetime import datetime


def generate_treeshell_tool(
    domain: str,
    tree_shell_class_name: str,
    tree_shell_import_path: str,
    description: str,
    commands: List[str] = None,
    example_workflow: List[str] = None,
    heaven_data_dir: str = None
) -> str:
    """
    Generate a HEAVEN tool for a domain-specific TreeShell.
    
    Args:
        domain: Domain name (e.g., "conversation", "code_analysis", "data_processing")
        tree_shell_class_name: Name of the TreeShell class (e.g., "ConversationTreeShellBase")
        tree_shell_import_path: Import path for the TreeShell class (e.g., "heaven_tree_repl.tools")
        description: Description of what this TreeShell does
        commands: List of common commands for documentation
        example_workflow: List of example workflow steps
        heaven_data_dir: HEAVEN data directory path (defaults to env var)
    
    Returns:
        Path to the generated tool file
    """
    
    # Get heaven data directory
    if heaven_data_dir is None:
        heaven_data_dir = os.environ.get('HEAVEN_DATA_DIR')
        if not heaven_data_dir:
            raise ValueError("HEAVEN_DATA_DIR environment variable not set and heaven_data_dir not provided")
    
    # Create tools directory
    tools_dir = Path(heaven_data_dir) / "tools"
    tools_dir.mkdir(parents=True, exist_ok=True)
    
    # Generate file names
    domain_snake = domain.lower().replace(" ", "_").replace("-", "_")
    tool_filename = f"{domain_snake}_agent_tree_shell_tool.py"
    tool_path = tools_dir / tool_filename
    
    # Generate class names
    domain_pascal = "".join(word.capitalize() for word in domain_snake.split("_"))
    tool_class_name = f"{domain_pascal}AgentTreeShellTool"
    args_schema_name = f"{domain_pascal}AgentTreeShellToolArgsSchema"
    session_manager_name = f"{domain_pascal}TreeShellSessionManager"
    func_name = f"{domain_snake}_treeshell_func"
    
    # Default commands if not provided
    if commands is None:
        commands = [
            "'' (empty) - Show current menu/options",
            "'jump X.Y' - Navigate to specific coordinate",
            "'1 {...}' - Execute menu option with arguments",
            "'back' - Go back to previous position",
            "'menu' - Show current node menu"
        ]
    
    # Default example workflow if not provided
    if example_workflow is None:
        example_workflow = [
            "'' - Show main menu",
            "'jump 0.1' - Navigate to primary function",
            "'1 {...}' - Execute with parameters",
            "'back' - Return to menu"
        ]
    
    # Generate the tool code
    tool_code = f'''#!/usr/bin/env python3
"""
{domain_pascal} TreeShell Tool - Generated by HEAVEN TreeShell Tool Generator

This module contains a HEAVEN tool that wraps a {domain} TreeShell application
to make it available as a tool for HEAVEN agents.

Generated on: {datetime.now().isoformat()}
Domain: {domain}
TreeShell Class: {tree_shell_class_name}
"""

from typing import Dict, Any
from heaven_base.baseheaventool import BaseHeavenTool, ToolArgsSchema, ToolResult
from {tree_shell_import_path} import {tree_shell_class_name}


class {args_schema_name}(ToolArgsSchema):
    arguments: Dict[str, Dict[str, Any]] = {{
        "command": {{
            "name": "command",
            "type": "str", 
            "description": "TreeShell command to execute (e.g., 'jump 0.1', '1 {{...}}', or '' for menu)",
            "required": True
        }}
    }}


class {session_manager_name}:
    """Manages a single persistent TreeShell session for {domain}."""
    
    def __init__(self):
        self.shell = None
    
    def get_shell(self) -> {tree_shell_class_name}:
        """Get the persistent shell instance."""
        if self.shell is None:
            self.shell = {tree_shell_class_name}()
        return self.shell


async def {func_name}(command: str) -> str:
    """Execute a command in the {domain} TreeShell."""
    
    # Get the persistent shell instance
    shell = {tool_class_name}.get_session_manager().get_shell()
    
    # Execute command
    result = await shell.handle_command(command)
    
    # Format the result
    if isinstance(result, dict):
        if result.get("error"):
            return f"❌ Error: {{result['error']}}"
        else:
            # Format successful result
            return f"✅ TreeShell Result:\\n{{result}}"
    else:
        return str(result)


class {tool_class_name}(BaseHeavenTool):
    name = "{tool_class_name}"
    description = """Navigate and execute commands in a {domain} TreeShell.
    
{description}

Commands:
{chr(10).join(f"- {cmd}" for cmd in commands)}

Example workflow:
{chr(10).join(f"{i+1}. {step}" for i, step in enumerate(example_workflow))}
"""
    func = {func_name}
    args_schema = {args_schema_name}
    is_async = True
    
    _session_manager = None  # Class-level singleton
    
    @classmethod
    def get_session_manager(cls):
        """Get or create the persistent session manager."""
        if cls._session_manager is None:
            cls._session_manager = {session_manager_name}()
        return cls._session_manager
'''
    
    # Write the file
    with open(tool_path, 'w') as f:
        f.write(tool_code)
    
    print(f"Generated TreeShell tool: {tool_path}")
    print(f"Tool class: {tool_class_name}")
    print(f"Domain: {domain}")
    
    return str(tool_path)


def generate_tool_from_config(config: Dict[str, Any]) -> str:
    """
    Generate a TreeShell tool from a configuration dictionary.
    
    Args:
        config: Configuration dictionary with required fields:
            - domain: Domain name
            - tree_shell_class_name: TreeShell class name
            - tree_shell_import_path: Import path
            - description: Tool description
            Optional fields:
            - commands: List of commands
            - example_workflow: List of workflow steps
            - heaven_data_dir: HEAVEN data directory
    """
    return generate_treeshell_tool(**config)


def _generate_treeshell_tool(args: dict) -> tuple[dict, bool]:
    """
    TreeShell node function to generate HEAVEN agent tools from TreeShell classes.
    
    Args:
        import_string: Python import path (e.g., "my_project.shells")
        class_name: TreeShell class name (e.g., "DataAnalysisTreeShell") 
        domain: Optional domain name (auto-extracted if not provided)
        description: Optional description (auto-generated if not provided)
    
    Returns:
        (result_dict, success_bool)
    """
    import importlib
    
    # Extract arguments
    import_string = args.get("import_string", "").strip()
    class_name = args.get("class_name", "").strip()
    domain = args.get("domain", "").strip()
    description = args.get("description", "").strip()
    
    # Validate required arguments
    if not import_string:
        return {"error": "import_string is required (e.g., 'my_project.shells')"}, False
        
    if not class_name:
        return {"error": "class_name is required (e.g., 'DataAnalysisTreeShell')"}, False
    
    try:
        # Try to import the module to validate
        module = importlib.import_module(import_string)
        if not hasattr(module, class_name):
            return {"error": f"Class '{class_name}' not found in module '{import_string}'"}, False
            
    except ImportError as e:
        return {"error": f"Cannot import module '{import_string}': {e}"}, False
    
    # Auto-generate domain from class name if not provided
    if not domain:
        # Remove "TreeShell" suffix and convert to snake_case
        domain_raw = class_name.replace("TreeShell", "").replace("Base", "")
        # Convert CamelCase to snake_case
        domain = ""
        for i, char in enumerate(domain_raw):
            if char.isupper() and i > 0:
                domain += "_"
            domain += char.lower()
        
        if not domain:
            domain = "custom"
    
    # Auto-generate description if not provided
    if not description:
        domain_title = domain.replace("_", " ").title()
        description = f"This tool provides access to a TreeShell interface for {domain_title} operations."
    
    try:
        # Generate the tool using existing function
        tool_path = generate_treeshell_tool(
            domain=domain,
            tree_shell_class_name=class_name,
            tree_shell_import_path=import_string,
            description=description
        )
        
        return {
            "action": "generate_treeshell_tool",
            "domain": domain,
            "class_name": class_name,
            "import_string": import_string,
            "tool_path": tool_path,
            "result": f"Successfully generated {domain} agent TreeShell tool at {tool_path}"
        }, True
        
    except Exception as e:
        return {"error": f"Failed to generate tool: {e}"}, False


# Example usage
if __name__ == "__main__":
    # Example: Generate a conversation TreeShell tool
    conversation_config = {
        "domain": "conversation",
        "tree_shell_class_name": "ConversationTreeShellBase", 
        "tree_shell_import_path": "heaven_tree_repl.tools",
        "description": "This tool provides access to a TreeShell interface for managing conversations with:\\n- Starting new conversations\\n- Continuing existing conversations\\n- Listing and searching conversations\\n- Loading specific conversations",
        "commands": [
            "'' (empty) - Show current menu/options",
            "'jump 0.1' - Navigate to start_chat",
            "'1 {\"title\": \"My Chat\", \"message\": \"Hello\"}' - Execute with arguments",
            "'back' - Go back to previous position", 
            "'menu' - Show current node menu"
        ],
        "example_workflow": [
            "'' - Show main menu",
            "'jump 0.1' - Go to start_chat", 
            "'1 {\"title\": \"Test\", \"message\": \"Hello\"}' - Start conversation",
            "'jump 0.2' - Go to continue_chat",
            "'1 {\"message\": \"How are you?\"}' - Continue conversation"
        ]
    }
    
    generate_tool_from_config(conversation_config)