import re
import json
import subprocess
import sys
import requests
import io
import time
import tempfile
import pkg_resources
from pathlib import Path
from contextlib import redirect_stdout

class ProgramGenerator:
    def __init__(self):
        # Default model
        self.model = 'qwen2.5-coder:32b'
        self.available_models = [
            "mxbai-embed-large:latest",
            "qwq:latest",
            "qwen2-math:7b",
            "qwen2.5-coder:32b",
            "gemma3:27b",
            "gemma3:latest",
            "llama3:latest",
            "llama3.2:latest"
        ]
        self.api_url = "http://localhost:11434/api/generate"
        self.headers = {"Content-Type": "application/json"}
        self.temp_dir = Path(tempfile.gettempdir()) / "program_generator"
        self.temp_dir.mkdir(exist_ok=True)
        
    def select_model(self):
        """Permite utilizatorului să selecteze un model disponibil"""
        print("\nModele disponibile:")
        for i, model in enumerate(self.available_models, 1):
            print(f"{i}. {model}")
        
        while True:
            try:
                choice = input(f"\nSelectează un model (1-{len(self.available_models)}) [implicit: {self.model}]: ")
                if not choice:  # Utilizatorul a apăsat doar Enter
                    break
                
                choice = int(choice)
                if 1 <= choice <= len(self.available_models):
                    self.model = self.available_models[choice-1]
                    print(f"\nModel selectat: {self.model}")
                    break
                else:
                    print(f"Te rog selectează un număr între 1 și {len(self.available_models)}")
            except ValueError:
                print("Te rog introdu un număr valid")
        
        return self.model

    def check_ollama_connection(self):
        """Verifică dacă serverul Ollama este disponibil"""
        try:
            response = requests.get("http://localhost:11434/api/tags")
            if response.status_code == 200:
                return True
            return False
        except requests.exceptions.ConnectionError:
            return False
        except Exception as e:
            print(f"Eroare la verificarea conexiunii Ollama: {str(e)}")
            return False
            
    def refresh_available_models(self):
        """Actualizează lista de modele disponibile de la serverul Ollama"""
        try:
            response = requests.get("http://localhost:11434/api/tags")
            if response.status_code == 200:
                data = response.json()
                self.available_models = [model['name'] for model in data.get('models', [])]
                return True
            return False
        except Exception as e:
            print(f"Eroare la actualizarea modelelor: {str(e)}")
            return False

    def install_missing_modules(self, code):
        """Detectează și instalează module lipsă din cod"""
        required_modules = re.findall(r'import ([^ \n.,]+)', code)
        required_modules.extend(re.findall(r'from ([^ \n.,]+)', code))
        installed_modules = {dist.project_name for dist in pkg_resources.working_set}
        
        missing_modules = set(required_modules) - installed_modules
        for module in missing_modules:
            if module not in ['sys', 'os', 're', 'json', 'time']:  # Module standard
                print(f"\nInstalare modul {module}...")
                try:
                    subprocess.check_call([sys.executable, "-m", "pip", "install", module],
                                        stdout=subprocess.DEVNULL,
                                        stderr=subprocess.DEVNULL)
                    print(f"Modulul {module} a fost instalat cu succes!")
                except Exception as e:
                    print(f"Nu s-a putut instala modulul {module}: {str(e)}")

    def generate_filename(self, idea, temp=False):
        """Generează nume de fișier bazat pe idee"""
        filename = re.sub(r'[^a-zA-Z0-9\s]', '', idea.lower())
        filename = re.sub(r'\s+', '_', filename.strip())
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        filename = f"{filename[:30]}_{timestamp}.py"
        
        if temp:
            return self.temp_dir / filename
        return filename

    def format_response_line(self, line):
        try:
            data = json.loads(line)
            return data.get('response', '')
        except json.JSONDecodeError:
            return line

    def extract_python_code(self, text):
        patterns = [
            r"```python\n(.*?)\n```",
            r"```\n(.*?)\n```",
            r"```python(.*?)```",
            r"`{3,}(.*?)`{3,}"
        ]
        
        for pattern in patterns:
            match = re.search(pattern, text, re.DOTALL)
            if match:
                return match.group(1).strip()
        return text.strip()

    def create_prompt(self, idea, is_correction=False, code_to_fix=None):
        """Creează un prompt adecvat pentru modelul selectat"""
        model_type = self.model.split(':')[0].lower()
        
        if is_correction and code_to_fix:
            if "qwen" in model_type:
                return f"Corectează acest cod Python care a dat eroarea: {idea}\nCod de corectat:\n```python\n{code_to_fix}\n```\nReturnează doar codul corectat în format Python, fără explicații, în blocuri markdown ```python."
            elif "gemma" in model_type:
                return f"Eroare Python: {idea}\nCod cu eroare:\n```python\n{code_to_fix}\n```\nCorecteaza codul si returneaza-l in format markdown ```python, fara explicatii."
            elif "llama" in model_type:
                return f"Eroare de executie Python: {idea}\nCod cu bug:\n```python\n{code_to_fix}\n```\nReturnează doar codul Python corectat în format markdown ```python, fără alte explicații."
            else:
                return f"Corectează acest cod Python care a dat eroarea: {idea}\nCod de corectat:\n{code_to_fix}"
        else:
            if "qwen" in model_type:
                return f"Generează un program Python complet care să îndeplinească următoarea cerință: {idea}\nReturnează doar codul în format markdown ```python, fără explicații."
            elif "gemma" in model_type:
                return f"Scrie un program Python complet care: {idea}\nReturnează doar codul Python în blocuri markdown ```python, fără explicații."
            elif "llama" in model_type:
                return f"Scrie un program Python funcțional pentru: {idea}\nFurnizează doar codul Python în blocuri markdown ```python, fără alte explicații."
            else:
                return f"Generează un program Python care să îndeplinească următoarea cerință: {idea}"

    def generate_code(self, idea, is_correction=False, code_to_fix=None):
        print(f"\nGenerare cod în progres folosind modelul {self.model}...")
        sys.stdout.flush()
        
        prompt = self.create_prompt(idea, is_correction, code_to_fix)
        payload = {
            "model": self.model,
            "prompt": prompt,
            "stream": True
        }

        try:
            response = requests.post(self.api_url, json=payload, headers=self.headers, stream=True)
            if response.status_code == 200:
                full_response = ""
                for line in response.iter_lines():
                    if line:
                        text = self.format_response_line(line.decode('utf-8'))
                        if text:
                            full_response += text
                            sys.stdout.write('.')
                            sys.stdout.flush()
                
                print("\n")  # Linie nouă după indicatorul de progres
                code = self.extract_python_code(full_response)
                if code:
                    print("\nCod generat:")
                    print(code)
                    return code
            else:
                print(f"Eroare API Ollama: {response.status_code}, {response.text}")
            return ""

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

    def run_code_safely(self, code, temp_file):
        """Rulează codul într-un fișier temporar cu gestiunea erorilor"""
        try:
            # Adaugă gestiunea erorilor în cod
            safe_code = f"""
try:
{chr(10).join('    ' + line for line in code.split(chr(10)))}
except Exception as e:
    print(f'Eroare: {{str(e)}}')
"""
            # Scrie codul în fișierul temporar
            with open(temp_file, 'w', encoding='utf-8') as f:
                f.write(safe_code)
            
            # Rulează codul
            result = subprocess.run([sys.executable, temp_file], 
                                  capture_output=True, 
                                  text=True)
            return result
        except Exception as e:
            print(f"Eroare la rularea codului: {str(e)}")
            return None

    def save_code(self, code, idea):
        filename = self.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}")
            return True
        except Exception as e:
            print(f"Eroare la salvarea fișierului: {str(e)}")
            return False

    def run_code(self, code, original_idea):
        print("\nRulare program...")
        max_attempts = 3
        attempt = 0
        
        # Verifică și instalează module lipsă înainte de rulare
        self.install_missing_modules(code)
        
        temp_file = self.generate_filename(original_idea, temp=True)
        
        while attempt < max_attempts:
            result = self.run_code_safely(code, temp_file)
            
            if result and result.returncode == 0:
                print("\nIeșire program:")
                print(result.stdout)
                self.save_code(code, original_idea)
                return True
            
            if result:
                print(f"\nEroare detectată: {result.stderr}")
                print("Încercare de corectare a codului...")
                
                corrected_code = self.generate_code(
                    result.stderr,
                    is_correction=True,
                    code_to_fix=code
                )
                
                if corrected_code and corrected_code != code:
                    print("\nCod corectat:")
                    print(corrected_code)
                    code = corrected_code
                    attempt += 1
                    continue
            
            print("Nu s-a putut corecta codul")
            break

        try:
            if isinstance(temp_file, Path):
                temp_file.unlink()  # Șterge fișierul temporar
            else:
                Path(temp_file).unlink()
        except:
            pass
            
        return False

    def run(self):
        print("Generator de Programe Python v5.0 - Adaptat pentru modele Ollama")
        print("----------------------------------------------------------")
        
        if not self.check_ollama_connection():
            print("\nNu s-a putut conecta la serverul Ollama. Asigură-te că Ollama rulează și încearcă din nou.")
            return
            
        self.select_model()
        
        while True:
            idea = input("\nIntrodu ideea ta de program (sau 'exit' pentru ieșire, 'model' pentru schimbare model): ")
            if idea.lower() == 'exit':
                break
            
            if idea.lower() == 'model':
                self.select_model()
                continue
                
            if idea:
                code = self.generate_code(idea)
                if code:
                    self.run_code(code, idea)
            
            print("\n" + "="*50 + "\n")

def main():
    try:
        import requests
    except ImportError:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"])
    
    generator = ProgramGenerator()
    generator.run()

if __name__ == "__main__":
    main()