import gradio as gr
import os
import shutil
from transformers import AutoTokenizer, AutoModelForCausalLM

# Model details and token
model_name = "Salesforce/codegen-350M-mono"
token = "token"  # Replace with your actual token

# Load tokenizer and model with authentication token
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=token)
model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=token)

# Function to explain the Java code
def explain_code(code):
    inputs = tokenizer(code, return_tensors="pt")
    outputs = model.generate(inputs['input_ids'], max_length=500)
    explanation = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return explanation

# Function to extract Java files from a directory
def extract_java_files(src_dir, dest_dir=None):
    java_files = []
    for root, dirs, files in os.walk(src_dir):
        for file in files:
            if file.endswith('.java'):
                java_file_path = os.path.join(root, file)
                java_files.append(java_file_path)
                # If a destination directory is provided, copy the file
                if dest_dir:
                    if not os.path.exists(dest_dir):
                        os.makedirs(dest_dir)
                    shutil.copy(java_file_path, os.path.join(dest_dir, file))
    return java_files

# Function to read and process uploaded file
def process_uploaded_file(file):
    code = file.read().decode("utf-8")
    return explain_code(code)

# Create a folder for storing extracted Java files (optional)
src_directory = "path_to_your_directory"  # replace with your folder path
dest_directory = "extracted_java_files"  # Folder where Java files will be copied
java_files = extract_java_files(src_directory, dest_directory)

# Gradio Interface for selecting uploaded Java file or explaining code
def explain_from_uploaded_or_folder(file=None):
    if file:
        code = file.read().decode("utf-8")
    else:
        # If no file uploaded, use the first Java file from the extracted folder (or prompt user)
        if java_files:
            with open(java_files[0], "r") as f:
                code = f.read()
        else:
            return "No Java files found."

    # Process and return explanation of the code
    return explain_code(code)

# Create the Gradio interface
iface = gr.Interface(
    fn=explain_from_uploaded_or_folder,
    inputs=[
        gr.File(label="Upload Java File", file_types=[".java"]),
        gr.Button("Use First Extracted Java File", visible=not bool(java_files))
    ],
    outputs="text",
    live=True
)

# Launch the interface
iface.launch()