import streamlit as st
import os
import json
from transformers import AutoModelForCausalLM, AutoTokenizer

# UI Components for Streamlit
st.title("Ansible Code Reviewer")
uploaded_files = st.file_uploader("Upload your Ansible code files", type=['yml', 'yaml'], accept_multiple_files=True)

if uploaded_files:
    result = []
    model_name = "facebook/incoder-1B"  # Example model name
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)

    # Process each uploaded file
    for uploaded_file in uploaded_files:
        content = uploaded_file.read().decode("utf-8")
        # Here you could use the model to evaluate the content
        tokens = tokenizer(content, return_tensors="pt")
        review_output = model.generate(**tokens)
        review_text = tokenizer.decode(review_output[0], skip_special_tokens=True)

        # Store results
        result.append({
            "filename": uploaded_file.name,
            "review": review_text
        })

    # Save the results to a JSON file
    json_result = json.dumps(result, indent=4)
    with open("review_results.json", "w") as f:
        f.write(json_result)

    # Display results and download link
    st.json(result)
    st.download_button("Download Review Results", json_result, file_name="review_results.json")