import re
import json
import subprocess
import sys
import requests
import io
from contextlib import redirect_stdout

def install_package(package_name):
    print(f"\nInstalare modul {package_name}...")
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", package_name], 
                            stdout=subprocess.DEVNULL, 
                            stderr=subprocess.DEVNULL)
        print(f"Modulul {package_name} a fost instalat cu succes!")
    except Exception as e:
        print(f"Nu s-a putut instala modulul {package_name}")

def generate_filename(idea):
    # Convertește ideea într-un nume de fișier valid
    filename = re.sub(r'[^a-zA-Z0-9\s]', '', idea.lower())
    filename = re.sub(r'\s+', '_', filename.strip())
    # Limitează lungimea numelui fișierului
    filename = filename[:50] if len(filename) > 50 else filename
    return f"{filename}.py"

def format_response_line(line):
    try:
        data = json.loads(line)
        if 'response' in data:
            return data['response']
        else:
            return ""
    except json.JSONDecodeError:
        return line

def generate_code_from_idea(idea, is_correction=False):
    url = "http://localhost:11434/api/generate"
    headers = {"Content-Type": "application/json"}
    
    prompt = idea if is_correction else f"Generează un program Python care să îndeplinească următoarea cerință: {idea}"
    
    payload = {
        "model": "qwen2.5-coder",
        "prompt": prompt,
        "stream": True
    }

    print("\nGenerare cod în progres...")
    sys.stdout.flush()
    
    full_response = ""

    try:
        response = requests.post(url, json=payload, headers=headers, stream=True)

        if response.status_code == 200:
            for line in response.iter_lines():
                if line:
                    text = format_response_line(line.decode('utf-8'))
                    if text:
                        full_response += text

            # Extrage codul Python dintre delimitatori
            match = re.search(r"```python\n(.*?)\n```", full_response, re.DOTALL)
            if not match:
                match = re.search(r"```\n(.*?)\n```", full_response, re.DOTALL)
            
            if match:
                code = match.group(1).strip()
                print("\nCod generat:")
                print(code)
                return code
            else:
                return ""
                
        else:
            return ""

    except requests.exceptions.ConnectionError:
        print("Nu s-a putut conecta la serverul local Ollama")
        return ""
    except Exception as e:
        print(f"Eroare la generarea codului: {str(e)}")
        return ""

def save_code_to_file(code, idea):
    filename = generate_filename(idea)
    try:
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(code)
        print(f"\nCod salvat cu succes în fișierul: {filename}")
    except Exception as e:
        print(f"Eroare la salvarea fișierului: {str(e)}")

def run_code(code, original_idea):
    print("\nRulare program...")
    while True:
        try:
            # Captează output-ul
            f = io.StringIO()
            with redirect_stdout(f):
                exec_globals = {}
                exec(code, exec_globals)
            
            output = f.getvalue()
            print("\nIeșire program:")
            print(output)
            
            # Salvează codul doar dacă a rulat cu succes
            save_code_to_file(code, original_idea)
            break

        except ModuleNotFoundError as e:
            module_name = re.search(r"No module named '(\w+)'", str(e))
            if module_name:
                install_package(module_name.group(1))
                continue
            else:
                print(f"Eroare: {str(e)}")
                break

        except Exception as e:
            print(f"\nEroare detectată: {str(e)}")
            print("Încercare de corectare a codului...")
            corrected_code = generate_code_from_idea(
                f"Corectează acest cod Python care a dat eroarea {str(e)}. Cod de corectat:\n{code}",
                is_correction=True
            )
            
            if corrected_code and corrected_code != code:
                print("\nCod corectat:")
                print(corrected_code)
                code = corrected_code
                continue
            else:
                print("Nu s-a putut corecta codul")
                break

def main():
    while True:
        idea = input("\nIntrodu ideea ta de program (sau 'exit' pentru ieșire): ")
        if idea.lower() == 'exit':
            break
            
        if idea:
            code = generate_code_from_idea(idea)
            if code:
                run_code(code, idea)
        
        print("\n" + "="*50 + "\n")

if __name__ == "__main__":
    try:
        import requests
    except ImportError:
        install_package("requests")
    main()