import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load a Hugging Face model (e.g., GPT-Neo)
model_name = "EleutherAI/gpt-neo-1.3B"  # Choose a model from Hugging Face
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Function to analyze and fix shell scripts
def analyze_and_fix_shell_script(script_content):
    # Create the prompt for GPT-Neo to analyze and fix the shell script
    prompt = f"""
    I have the following shell script. Please identify any errors, inefficiencies, or improvements that can be made. Provide an explanation of each issue and then suggest an improved version of the script:
    
    Script:
    {script_content}
    
    Please return the improved script and highlight the changes you made.
    """
    
    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
    outputs = model.generate(**inputs, max_length=1024, num_return_sequences=1)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Gradio interface to upload shell script and process
def upload_and_fix(file):
    # Handle both string and byte file formats
    script_content = file if isinstance(file, str) else file.decode("utf-8")
    
    # Call the GPT model to analyze and fix the shell script
    fixed_script = analyze_and_fix_shell_script(script_content)
    
    return fixed_script

# Create a Gradio interface
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            gr.Markdown("## Upload Shell Script for Analysis and Fixing")
            file_input = gr.File(label="Upload Shell Script (.sh)")
            output_text = gr.Textbox(label="Fixed Shell Script", lines=20)
            submit_btn = gr.Button("Analyze and Fix")
        
        # Define the button action
        submit_btn.click(upload_and_fix, inputs=file_input, outputs=output_text)
    
    # Launch the app
    demo.launch()