import ollama
from playwright.sync_api import sync_playwright
import sqlite3
import subprocess
import os
import re
import sys

# Lista modelelor necesare
MODELS_NEEDED = ["llama3.1", "codellama"]

# Inițializează clientul Ollama
ollama_client = ollama.Client()

# Conectare la baza de date SQLite pentru memorie
conn = sqlite3.connect('agent_memory.db')
cursor = conn.cursor()
cursor.execute('''CREATE TABLE IF NOT EXISTS solutions
                  (task TEXT PRIMARY KEY, code TEXT)''')
conn.commit()

# Funcție pentru a verifica și instala modelele necesare
def check_and_install_models():
    installed_models = [model['name'] for model in ollama_client.list()['models']]
    for model in MODELS_NEEDED:
        if model not in installed_models:
            print(f"Instalez modelul {model}...")
            subprocess.run(["ollama", "pull", model])

# Funcție pentru a verifica și porni Ollama dacă nu rulează
def start_ollama():
    try:
        ollama_client.list()
    except Exception:
        print("Ollama nu rulează. Pornesc Ollama...")
        subprocess.run(["ollama", "serve"])

# Funcție pentru a extrage codul din răspunsul modelului
def extract_code(response):
    match = re.search(r"```python(.*?)```", response, re.DOTALL)
    if match:
        return match.group(1).strip()
    return response  # Dacă nu găsește bloc de cod, returnează tot răspunsul

# Funcție pentru a verifica dacă e nevoie de căutare pe internet
def needs_internet_search(response):
    return "caută pe internet" in response.lower()

# Funcție pentru a verifica dacă e nevoie de generare de cod
def needs_code_generation(response):
    return "generează cod" in response.lower()

# Funcție pentru căutări pe internet cu Playwright
def search_internet(query):
    with sync_playwright() as p:
        browser = p.chromium.launch()
        page = browser.new_page()
        page.goto("https://www.google.com")
        page.fill("input[name='q']", query)
        page.press("input[name='q']", "Enter")
        page.wait_for_selector("h3")
        results = page.eval_on_selector_all("h3", "elements => elements.map(e => e.textContent)")
        browser.close()
        return results

# Funcție pentru a genera cod cu un model specializat
def generate_code(task):
    response = ollama_client.generate(model="codellama", prompt=task)
    code = extract_code(response['generated_text'])
    return code

# Funcție pentru a executa codul într-un mediu izolat
def execute_code(code):
    with open('temp_script.py', 'w') as f:
        f.write(code)
    
    result = subprocess.run([sys.executable, 'temp_script.py'], capture_output=True, text=True)
    return result.stdout, result.stderr, result.returncode

# Funcția principală a agentului
def main():
    start_ollama()  # Asigură că Ollama rulează
    check_and_install_models()  # Verifică și instalează modelele necesare
    
    while True:
        task = input("Introdu sarcina: ")
        if not task:
            break
        
        # Verifică soluții existente
        cursor.execute("SELECT code FROM solutions WHERE task=?", (task,))
        stored_solution = cursor.fetchone()
        if stored_solution:
            print("Folosesc soluția stocată.")
            code = stored_solution[0]
            stdout, stderr, returncode = execute_code(code)
            if returncode == 0:
                print("Rezultat:", stdout)
            else:
                print("Eroare:", stderr)
        else:
            # Procesează sarcina cu modelul principal
            response = ollama_client.generate(model="llama3.1", prompt=task)
            initial_response = response['generated_text']
            print("Răspuns inițial:", initial_response)
            
            if needs_internet_search(initial_response):
                search_results = search_internet(task)
                print("Rezultate căutare:", search_results)
            elif needs_code_generation(initial_response):
                code = generate_code(task)
                print("Cod generat:", code)
                
                # Execută și corectează codul
                stdout, stderr, returncode = execute_code(code)
                if returncode == 0:
                    print("Execuție reușită:", stdout)
                    cursor.execute("INSERT INTO solutions (task, code) VALUES (?, ?)", (task, code))
                    conn.commit()
                else:
                    print("Eroare:", stderr)
                    new_prompt = f"{task}\nEroare: {stderr}"
                    code = generate_code(new_prompt)
                    print("Cod corectat:", code)
                    stdout, stderr, returncode = execute_code(code)
                    if returncode == 0:
                        print("Execuție reușită după corectare:", stdout)
                        cursor.execute("INSERT INTO solutions (task, code) VALUES (?, ?)", (task, code))
                        conn.commit()
                    else:
                        print("Eroare persistentă:", stderr)
            else:
                print("Răspuns:", initial_response)

if __name__ == "__main__":
    main()