import os
import shutil
from huggingface_hub import InferenceClient
from datetime import datetime
import ast
import re

# Initialize Hugging Face client
client = InferenceClient('HuggingFaceH4/zephyr-7b-beta')

# List of common functions/classes and their corresponding modules
common_functions = {
    'print': 'builtins',
    'len': 'builtins',
    'str': 'builtins',
    'int': 'builtins',
    'float': 'builtins',
    'list': 'builtins',
    'dict': 'builtins',
    'set': 'builtins',
    'time': 'time',
    'sleep': 'time',
    'datetime': 'datetime',
    'date': 'datetime',
    'now': 'datetime',
    'json': 'json',
    'loads': 'json',
    'dumps': 'json',
    'os': 'os',
    'path': 'os.path',
    're': 're',
    'match': 're',
    'search': 're',
    'requests': 'requests',
    'get': 'requests',
    'post': 'requests',
    'math': 'math',
    'sqrt': 'math',
    'ceil': 'math',
    'floor': 'math',
}

def format_script(script):
    """Format the extracted script for better readability."""
    return script.replace('\\n', '\n').replace('\"', '"')

def check_script_for_issues(script):
    """Check the script for common issues like missing modules and undefined variables."""
    issues = []
    suggested_imports = set()
    
    # Check for undefined functions/classes
    used_functions = set()
    for function in common_functions.keys():
        if function in script:
            used_functions.add(function)

    # Check for missing module imports
    for func in used_functions:
        if func in common_functions:
            module = common_functions[func]
            if f'import {module}' not in script and f'from {module}' not in script:
                issues.append(f"Missing module import for function '{func}' from module '{module}'.")
                suggested_imports.add(module)

    # Check for undefined variables
    try:
        # Parse the script to get all variable assignments
        tree = ast.parse(script)
        defined_vars = {node.targets[0].id for node in ast.walk(tree) if isinstance(node, ast.Assign)}
        used_vars = {node.id for node in ast.walk(tree) if isinstance(node, ast.Name) and not isinstance(node.ctx, ast.Store)}
        
        undefined_vars = used_vars - defined_vars
        for var in undefined_vars:
            issues.append(f"Undefined variable: {var}")
    except SyntaxError:
        issues.append("Syntax error in the script.")

    return issues, suggested_imports

def generate_fixed_script(script, suggested_imports):
    """Generate a fixed version of the script with suggested imports."""
    fixed_script = "\n".join(f"import {module}" for module in suggested_imports) + "\n\n" + script
    return fixed_script

def run_ai_chat(task_context, fixed_script):
    """Run the AI chat to get a new response based on the fixed script."""
    current_datetime = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    system_prompt = (
        f'You are an assistant helping with the \'{task_context}\' task. '
        f'The current date and time is {current_datetime}. Provide guidance, answer questions, or assist as needed.'
    )
    messages = [
        {'role': 'system', 'content': system_prompt},
        {'role': 'user', 'content': f"Here is a fixed Python script:\n```python\n{fixed_script}\n```"}  # Send the fixed script to the AI
    ]

    print("**AI Response:**")
    response = ''
    for reply in client.chat_completion(messages, stream=True, max_tokens=512):
        if hasattr(reply, 'choices') and reply.choices:
            content = reply.choices[0].delta.content
            if content:
                response += content
                print(content, end='', flush=True)

    print()  # Print a newline after the response is complete
    return response

def process_python_files(input_directory, cleaned_directory):
    """Process all Python files in the input directory."""
    for filename in os.listdir(input_directory):
        if filename.endswith('.py'):  # Check for Python files
            file_path = os.path.join(input_directory, filename)
            print(f"Processing file: {file_path}")

            with open(file_path, 'r', encoding='utf-8') as file:
                code = file.read()

            issues, suggested_imports = check_script_for_issues(code)
            if issues:
                print("\nIssues found in the code snippet:")
                for issue in issues:
                    print(f"- {issue}")
                
                # Generate the fixed version of the script
                fixed_script = generate_fixed_script(code, suggested_imports)
                print("\nSuggested fixed version of the script:\n")
                print(fixed_script)

                # Get a new response from the AI based on the fixed script
                ai_response = run_ai_chat("Fix the following Python script", fixed_script)

                # Save the AI's response to the cleaned directory
                cleaned_file_path = os.path.join(cleaned_directory, filename)
                with open(cleaned_file_path, 'w', encoding='utf-8') as cleaned_file:
                    cleaned_file.write(ai_response)
                print(f"Saved cleaned version of '{filename}' to cleaned directory.")
            else:
                print("\nNo issues found in the code snippet.")

            # Move the original processed file to the completed directory
            shutil.move(file_path, os.path.join(cleaned_directory, filename))
            print(f"Moved '{filename}' to completed directory.")

if __name__ == '__main__':
    input_directory = r'J:\oops\project_longtime\chat\extracted_content\cleaned_python\python'
    cleaned_directory = r'J:\oops\project_longtime\chat\pages\ai\completed_clean'
    process_python_files(input_directory, cleaned_directory)


make sure the new response is saved with a new filename as well please. so this only needs to save the new generated fixed script please not the old one