
import os
import sys


from HeGFlow.graph.tool_node import Tool_node
import abc
from typing import Dict, Any, Union, Optional
from abc import ABC, abstractmethod
from openai import AsyncOpenAI 
import io
import contextlib 
import sys
from openai import OpenAI
import re




def get_code_generation_prompt(task_description: str) -> str:
    
    return f"""You are a senior Python programming expert, skilled in generating concise, efficient, and complete Python code based on natural language tasks.

    Your code should be directly executable and should not rely on reading external files.
    If the task requires computation and output, ensure the code ultimately prints the result using the print() function.
    If the task requires function definitions, provide the function definition directly without including calling examples.

    Task description:
    {task_description}

    Output the complete Python code directly, without any explanations or additional text. The code must be self-contained and executable. Important: You must ensure that your entire output is nothing but executable code. Do not include any extra invalid characters, as this would cause significant harm.
    """

def _execute_code_safely(code: str, timeout: int = 5) -> Dict[str, Any]:
    
    output_capture = io.StringIO()
    error_capture = io.StringIO()
    
    
    sandbox_globals = {"__builtins__": {
        "print": print,
        "len": len,
        "str": str,
        "int": int,
        "float": float,
        "list": list,
        "dict": dict,
        "tuple": tuple,
        "set": set,
        "range": range,
        "sum": sum,
        "min": min,
        "max": max,
        
    }}
    sandbox_locals = {}

    try:
    
        with contextlib.redirect_stdout(output_capture), \
             contextlib.redirect_stderr(error_capture):
            
   
            exec(code, sandbox_globals, sandbox_locals)
            
        stdout = output_capture.getvalue()
        stderr = error_capture.getvalue()
        
        if stderr:
            return {"output": stdout, "error": stderr, "success": False}
        return {"output": stdout, "error": None, "success": True}
        
    except Exception as e:
        stderr = error_capture.getvalue()
        return {"output": output_capture.getvalue(), "error": f"Execution Error: {e}\n{stderr}", "success": False}

class Coder(Tool_node):
    def __init__(self):
        super().__init__(
            name="Coder",
            
        )

    async def _execute_tool(self, task:str) -> str:
        
        print(f"[CoderTool] Generating code for task: '{task}'...")
        code_prompt = get_code_generation_prompt(task)
        try:
            client = OpenAI(api_key="", base_url="")
            response = client.chat.completions.create(
            model="deepseek-chat",
            messages=[
                {"role": "system", "content": "You are a helpful AI assistant that writes Python code."},
                {"role": "user", "content": code_prompt},
            ],
            stream=False
                )
            generated_code = response.choices[0].message.content
            generated_code = generated_code.lstrip("```python\n").rstrip("\n```")
            print(f"[CoderTool] Generated code:\n{generated_code}\n---")
            

            execution_result = _execute_code_safely(generated_code)
            
            return f"code: {generated_code}, execution_output: {execution_result.get('output', '')}, execution_error: {execution_result.get('error', None)}, execution_success: {execution_result.get('success', False)}"
            
        except Exception as e:
            return f"code: '', execution_output: '', execution_error: Code generation or API error: {e}, execution_success: False"

        
        return input

    def _format_result_to_natural_language(self, raw_result: Any, task_description: str) -> str:
        
        pass
    
