import pandas as pd
import ollama
import requests
import re


def get_code_blocks(text):
    # 使用正则表达式匹配 ```python 和 ``` 之间的内容
    pattern = r'```python(.*?)```'
    matches = re.findall(pattern, text, re.DOTALL)

    # 去除每个匹配项的前后空白字符
    return [match.strip() for match in matches]
def generate_and_execute_code(excel_path, process_instruction):
    # 1. 读取Excel表头
    try:
        df_header = pd.read_excel(excel_path, nrows=0)
        headers = df_header.columns.tolist()
    except Exception as e:
        return f"读取Excel失败: {str(e)}"

    # 2. 构造模型提示词
    prompt = f"""
    已知数据表包含以下字段：{headers}。
    用户需求：{process_instruction}。
    请生成满足需求的Python代码，要求：
    - 使用pandas处理数据，变量名用df 
    - 不要添加文件保存代码 
    - 直接打印处理结果 
    - 数据的读取用这个代码 df = pd.read_excel("input.xlsx")  # 读取完整数据
    返回格式：```python\n[代码]\n```
    """

    # 3. 调用Ollama模型生成代码
    try:
        response = requests.post(
            "http://localhost:11434/api/chat",
            json={
                "model": "deepseek-r1:14b",
                "messages": [{"role": "user", "content": prompt}],
                "stream": False
            },
            timeout=6000
        )
        response.raise_for_status()
        generated_code = response.json()['message']['content']
        # 提取代码块
        code_block=get_code_blocks(generated_code)
        print(code_block)
        file = open("test.py", "w", encoding="utf-8")
        file.write(code_block[0])
        file.close()

    except Exception as e:
        return f"模型调用失败: {str(e)}"

    # 4. 执行生成的代码
    try:

        exec(code_block[0])
        return "执行成功"
    except Exception as e:
        return f"代码执行错误: {str(e)}\n生成代码:\n{code_block[0]}"


# 使用示例
if __name__ == "__main__":
    result = generate_and_execute_code(
        excel_path="input.xlsx",
        #process_instruction="请分析销售数据，计算各区域季度环比增长率"
        #process_instruction="请分析销售数据，销售金额总和"
        #process_instruction="请分析销售数据，按照区域区分统计销售金额总和"
        process_instruction="请分析销售数据，华东的Q1与华南的Q2的销售额相加是多少？"
    )
    print(result)