import os
import requests
import subprocess
import sys
import re
import time
from colorama import init, Fore, Style

init()

def get_ollama_host():
    return input("Introduce adresa OLLAMA_HOST: ")

def get_ai_response(host, message, model="llama3.3"):
    try:
        response = requests.post(
            f"{host}/api/chat",
            json={
                "model": model,
                "messages": [{"role": "user", "content": message}],
                "stream": False
            }
        )
        response.raise_for_status()
        return response.json()['message']['content']
    except Exception as e:
        return f"Error: {str(e)}"

def extract_code(text):
    # Caută blocuri de cod Python între ```python și ```
    code_blocks = re.findall(r'```python\n(.*?)\n```', text, re.DOTALL)
    if code_blocks:
        return code_blocks[0]
    # Caută blocuri de cod între ``` și ```
    code_blocks = re.findall(r'```(.*?)```', text, re.DOTALL)
    if code_blocks:
        return code_blocks[0]
    return text

def install_module(module):
    print(f"{Fore.YELLOW}Instalare modul {module}...{Style.RESET_ALL}")
    subprocess.check_call([sys.executable, "-m", "pip", "install", module])

def save_code(code, filename):
    with open(filename, 'w', encoding='utf-8') as f:
        f.write(code)
    print(f"{Fore.GREEN}Cod salvat în {filename}{Style.RESET_ALL}")

def test_code(code, requirements):
    print(f"{Fore.CYAN}Testare cod...{Style.RESET_ALL}")
    
    # Salvare temporară
    temp_file = "temp_test.py"
    save_code(code, temp_file)
    
    try:
        # Rulare cod
        result = subprocess.run([sys.executable, temp_file], 
                              capture_output=True, 
                              text=True,
                              timeout=30)
        
        if result.returncode == 0:
            print(f"{Fore.GREEN}Codul rulează cu succes!{Style.RESET_ALL}")
            print("\nIeșire program:")
            print(result.stdout)
            
            # Verificare cerințe
            feedback_prompt = f"""Analizează următoarele:
            Cod: {code}
            Ieșire: {result.stdout}
            Cerințe: {requirements}
            
            Codul îndeplinește toate cerințele? Răspunde doar cu DA sau NU, urmat de o scurtă explicație dacă este NU."""
            
            return True, result.stdout
        else:
            print(f"{Fore.RED}Eroare la rulare:{Style.RESET_ALL}")
            print(result.stderr)
            return False, result.stderr
            
    except subprocess.TimeoutExpired:
        print(f"{Fore.RED}Timeout - codul rulează prea mult timp{Style.RESET_ALL}")
        return False, "Timeout"
    except Exception as e:
        if "ModuleNotFoundError" in str(e):
            module = str(e).split("'")[1]
            install_module(module)
            return test_code(code, requirements)  # Reîncearcă după instalare
        print(f"{Fore.RED}Eroare: {str(e)}{Style.RESET_ALL}")
        return False, str(e)
    finally:
        if os.path.exists(temp_file):
            os.remove(temp_file)

def ai_code_generator():
    host = get_ollama_host()
    requirements = input("Descrie proiectul/scriptul Python dorit: ")
    
    attempt = 1
    while True:
        print(f"\n{Fore.YELLOW}Încercare #{attempt}{Style.RESET_ALL}")
        
        # Generare cod
        prompt = f"""Generează cod Python pentru următoarele cerințe: {requirements}
        Răspunde doar cu codul, fără explicații. Folosește doar module Python standard sau populare."""
        
        response = get_ai_response(host, prompt)
        code = extract_code(response)
        
        # Testare cod
        success, output = test_code(code, requirements)
        
        if success:
            # Verificare cu AI dacă codul îndeplinește cerințele
            verify_prompt = f"""Verifică dacă acest cod îndeplinește cerințele:
            Cod: {code}
            Ieșire: {output}
            Cerințe: {requirements}
            
            Răspunde doar cu DA sau NU."""
            
            verification = get_ai_response(host, verify_prompt).strip().upper()
            
            if verification == "DA":
                filename = input("\nIntroduce numele fișierului pentru salvare (ex: script.py): ")
                save_code(code, filename)
                print(f"{Fore.GREEN}Proiect generat și testat cu succes!{Style.RESET_ALL}")
                break
        
        # Generare feedback pentru îmbunătățire
        feedback_prompt = f"""Analizează următoarele:
        Cod: {code}
        Ieșire/Eroare: {output}
        Cerințe: {requirements}
        
        Ce trebuie modificat pentru a îndeplini cerințele? Răspunde concis."""
        
        feedback = get_ai_response(host, feedback_prompt)
        print(f"\n{Fore.YELLOW}Feedback: {feedback}{Style.RESET_ALL}")
        
        attempt += 1
        if attempt > 5:
            retry = input(f"\n{Fore.YELLOW}Am făcut 5 încercări. Continuăm? (da/nu): {Style.RESET_ALL}").lower()
            if retry != 'da':
                break
            attempt = 1

if __name__ == "__main__":
    ai_code_generator()