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

class ProgramGenerator:
    def __init__(self):
        self.model = 'qwen2.5-coder'
        self.api_url = "http://localhost:11434/api/generate"
        self.headers = {"Content-Type": "application/json"}

    def install_package(self, 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!")
            return True
        except Exception as e:
            print(f"Nu s-a putut instala modulul {package_name}: {str(e)}")
            return False

    def generate_filename(self, idea):
        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"
        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):
        # Încearcă mai multe pattern-uri comune pentru cod Python
        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()  # Returnează textul curat dacă nu găsește marcaje

    def generate_code(self, idea, is_correction=False):
        print("\nGenerare cod în progres...")
        sys.stdout.flush()
        
        prompt = idea if is_correction else f"Generează un program Python care să îndeplinească următoarea cerință: {idea}"
        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
                
                code = self.extract_python_code(full_response)
                if code:
                    print("\nCod generat:")
                    print(code)
                    return code
            return ""

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

    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
        
        while attempt < max_attempts:
            try:
                f = io.StringIO()
                with redirect_stdout(f):
                    exec_globals = {}
                    exec(code, exec_globals)
                
                output = f.getvalue()
                print("\nIeșire program:")
                print(output)
                
                self.save_code(code, original_idea)
                return True

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

            except Exception as e:
                print(f"\nEroare detectată: {str(e)}")
                print("Încercare de corectare a codului...")
                
                corrected_code = self.generate_code(
                    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
                    attempt += 1
                    continue
                else:
                    print("Nu s-a putut corecta codul")
                    break

        return False

    def run(self):
        print("Generator de Programe Python v2.0")
        print("--------------------------------")
        while True:
            idea = input("\nIntrodu ideea ta de program (sau 'exit' pentru ieșire): ")
            if idea.lower() == 'exit':
                break
                
            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()