from fastapi import APIRouter, HTTPException, Body
from pytrace.core.registry import NodeRegistry, get_type_string
import importlib.util
import sys
from pathlib import Path
import os

router = APIRouter()

# Create a directory to store custom nodes uploaded by clients
CUSTOM_NODES_DIR = Path("custom_nodes")
CUSTOM_NODES_DIR.mkdir(exist_ok=True)
if str(CUSTOM_NODES_DIR.resolve()) not in sys.path:
    sys.path.insert(0, str(CUSTOM_NODES_DIR.resolve()))


@router.post("/register")
async def register_custom_node(code: str = Body(..., embed=True)):
    """
    Receives Python code for a custom node, saves it, and loads it.
    """
    # This is a simplified and dangerous implementation for demonstration.
    # A real implementation would need heavy security validation (AST analysis, sandboxing).
    
    try:
        # For simplicity, we'll hash the code to create a filename.
        import hashlib
        code_hash = hashlib.sha256(code.encode()).hexdigest()
        file_path = CUSTOM_NODES_DIR / f"node_{code_hash}.py"
        
        if file_path.exists():
            # Node is already registered, we can assume.
            return {"message": "Node already exists and is registered."}
            
        file_path.write_text(code, encoding="utf-8")
        
        # Dynamically load the module.
        # The node's class should have the `@register_node` decorator,
        # so importing it will automatically register it.
        module_name = file_path.stem
        spec = importlib.util.spec_from_file_location(module_name, file_path)
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        
        # Find which node was just registered by looking for a new class
        # from the loaded module.
        newly_registered_type = None
        for node_type, node_class in NodeRegistry._implementations.items():
            if node_class.__module__ == module_name:
                newly_registered_type = node_type
                break
        
        if not newly_registered_type:
            # Cleanup the created file if registration fails
            os.remove(file_path)
            raise HTTPException(status_code=400, detail="Could not find a class with @register_node in the provided code.")

        return {
            "message": "Node registered successfully.",
            "node_type": newly_registered_type,
            "spec": NodeRegistry.get_spec(newly_registered_type).dict()
        }

    except Exception as e:
        import traceback
        raise HTTPException(
            status_code=500,
            detail={
                "error": str(e),
                "traceback": traceback.format_exc()
            }
        )