import ollama
import subprocess
import json
import os
import requests
from duckduckgo_search import DDGS

# Coloque sua chave da OpenWeatherMap aqui
OPENWEATHER_KEY = "SUA_CHAVE_AQUI"

# Ferramentas disponíveis
tools = [
    {
        "type": "function",
        "function": {
            "name": "rodar_comando",
            "description": "Executa um comando no terminal e retorna a saída",
            "parameters": {
                "type": "object",
                "properties": {
                    "cmd": {"type": "string", "description": "Comando shell"}
                },
                "required": ["cmd"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "buscar_web",
            "description": "Faz uma busca na web e retorna links e descrições",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Termo de busca"}
                },
                "required": ["query"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "ler_arquivo",
            "description": "Lê o conteúdo de um arquivo do disco",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "Caminho do arquivo"}
                },
                "required": ["path"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "salvar_arquivo",
            "description": "Cria ou sobrescreve um arquivo com conteúdo fornecido",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "Nome do arquivo"},
                    "conteudo": {"type": "string", "description": "Conteúdo a salvar"}
                },
                "required": ["path", "conteudo"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "listar_arquivos",
            "description": "Lista arquivos e pastas de um diretório",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "Caminho da pasta"}
                },
                "required": ["path"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "obter_clima",
            "description": "Obtém o clima atual de uma cidade via OpenWeatherMap",
            "parameters": {
                "type": "object",
                "properties": {
                    "cidade": {"type": "string", "description": "Nome da cidade"}
                },
                "required": ["cidade"]
            }
        }
    }
]

# Execução das ferramentas
def executar_tool(name, args):
    try:
        if name == "rodar_comando":
            return subprocess.getoutput(args["cmd"])

        elif name == "buscar_web":
            resultados = DDGS().text(args["query"], max_results=5)
            return json.dumps(list(resultados), indent=2, ensure_ascii=False)

        elif name == "ler_arquivo":
            if not os.path.exists(args["path"]):
                return "Arquivo não encontrado."
            with open(args["path"], "r", encoding="utf-8", errors="ignore") as f:
                return f.read()

        elif name == "salvar_arquivo":
            with open(args["path"], "w", encoding="utf-8") as f:
                f.write(args["conteudo"])
            return f"Arquivo '{args['path']}' salvo com sucesso."

        elif name == "listar_arquivos":
            if not os.path.exists(args["path"]):
                return "Diretório não encontrado."
            return "\n".join(os.listdir(args["path"]))

        elif name == "obter_clima":
            url = f"http://api.openweathermap.org/data/2.5/weather?q={args['cidade']}&appid={OPENWEATHER_KEY}&units=metric&lang=pt"
            resp = requests.get(url).json()
            if "main" in resp:
                return f"{resp['name']}: {resp['main']['temp']}°C, {resp['weather'][0]['description']}"
            return f"Erro ao obter clima: {resp}"

        else:
            return "Ferramenta desconhecida."
    except Exception as e:
        return f"Erro: {e}"

# Loop de chat
mensagens = []

while True:
    pergunta = input("Você: ")
    mensagens.append({"role": "user", "content": pergunta})

    resposta = ollama.chat(
        model="llama3.1:8b",
        messages=mensagens,
        tools=tools
    )

    # Se houver chamada de ferramenta
    if "message" in resposta and "tool_calls" in resposta["message"]:
        for call in resposta["message"]["tool_calls"]:
            nome = call["function"]["name"]
            args_raw = call["function"]["arguments"]
            if isinstance(args_raw, str):
                argumentos = json.loads(args_raw)
            else:
                argumentos = args_raw

            print(f"[Ollama chamou '{nome}' com {argumentos}]")
            resultado = executar_tool(nome, argumentos)

            mensagens.append({
                "role": "tool",
                "name": nome,
                "content": resultado
            })

            resposta_final = ollama.chat(
                model="llama3.1:8b",
                messages=mensagens
            )
            print("Ollama:", resposta_final["message"]["content"])
    else:
        print("Ollama:", resposta["message"]["content"])