import re
from langchain import LLMChain, PromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.llms import OpenAI
from langchain.agents import tool, Tool
# from langchain.utilities import PythonREPL

import sys
from io import StringIO
from typing import Dict, Optional

from pydantic import BaseModel, Field

from models import llm


class PythonREPL(BaseModel):
    """Simulates a standalone Python REPL."""

    # globals: Optional[Dict] = Field(default_factory=dict, alias="_globals")
    # locals: Optional[Dict] = Field(default_factory=dict, alias="_locals")

    def run(self, command: str) -> str:
        """Run command with own globals/locals and returns anything printed."""
        old_stdout = sys.stdout
        sys.stdout = mystdout = StringIO()
        try:
            code_content = command
            if('```python' in command): 
                start = command.find('```python') + len('```python')
                end = command.rfind('```')
                code_content = command[start:end].strip()
            elif("```" in command): 
                start = command.find('```') + len('```')
                end = command.rfind('```')
                code_content = command[start:end].strip()
            exec(code_content, globals(), globals())
            sys.stdout = old_stdout
            output = mystdout.getvalue()
        except Exception as e:
            sys.stdout = old_stdout
            output = str(e)
        return output
    

generate_python_code = """
Please write Python script to fulfill the following requirement:

---
{input}
---

Only output the code section with code block, without __name__ guard.
"""

generate_python_code_promopt = PromptTemplate(input_variables=["input"], template=generate_python_code,)

generate_code_chain = LLMChain(llm = llm(temperature=0.1), prompt=generate_python_code_promopt, output_key="code")


@tool("Generate and Excute Python Code ", return_direct=True)
def generate_and_excute_python_code(input: str) -> str:
    '''useful for when you need to generate python code and excute it'''
    answer_code = generate_code_chain.run(input)
    python_repl = PythonREPL()
    result = python_repl.run(answer_code)
    print(result)
    return f"""
code:
```
{answer_code}
```

execute result:
---
{result}
---
    """

python_repl = PythonREPL()   
repl_tool = Tool(
    name="python_repl",
    description="A Python shell. Use this to execute python commands. Input should be a valid python command. If you want to see the output of a value, you should print it out with `print(...)`.",
    func=python_repl.run
)

if __name__ == "__main__":
    input = """
我有一个json文件url为: https://artwork-assets-staging-sbux.starbucks.com.cn/accountavatars.json
并按照如下Example进行格式转换
文件格式为:
```
{
'artworks': {
        'file1.png': {
            'middle@1x': '***', 
            'middle@2x': '***', 
            'middle@3x': '***'
        },
        'file2.png': {
            'middle@1x': '***', 
            'middle@2x': '***', 
            'middle@3x': '***'
        }
    }
}
```
输出格式:
```
curl https://active.stg.starbucks.com.cn/accountAvatar/file1.png
curl https://active.stg.starbucks.com.cn/accountAvatar/file2.png
```
"""

    result = generate_and_excute_python_code(input)
    print(result)