import gradio as gr
import os
import hashlib
import json
import subprocess

# --- Setup vault folders ---
for folder in ["etc", "data", "config", "modules"]:
    os.makedirs(folder, exist_ok=True)

# --- Helper functions ---
def compute_sha256(file_path):
    sha256 = hashlib.sha256()
    with open(file_path, "rb") as f:
        while chunk := f.read(8192):
            sha256.update(chunk)
    return sha256.hexdigest()

def generate_hash_chain(seed, length=100):
    chain = [seed]
    current = seed
    for _ in range(length-1):
        current = hashlib.sha256(current.encode()).hexdigest()
        chain.append(current)
    return chain

def save_permanently_if_requested(save_perm):
    if save_perm:
        try:
            subprocess.run(["git", "add", "."], check=True)
            subprocess.run(["git", "commit", "-m", "GradioVault user save"], check=True)
            subprocess.run(["git", "push"], check=True)
            return "Saved permanently to repo."
        except subprocess.CalledProcessError:
            return "Failed to save permanently. Check Git permissions."
    return ""

# --- Main processing function ---
def process_image(uploaded_file, save_perm):
    if uploaded_file is None:
        return "No file uploaded."
    
    # Save image in vault
    save_path = os.path.join("data", uploaded_file.name)
    uploaded_file.save(save_path)

    # Compute checksum and hash chain
    checksum = compute_sha256(save_path)
    hash_chain = generate_hash_chain(checksum)
    
    # Save hash chain to a JSON file
    hash_file = os.path.join("data", f"{uploaded_file.name}_hash_chain.json")
    with open(hash_file, "w") as f:
        json.dump(hash_chain, f, indent=2)
    
    # Optionally save permanently
    git_msg = save_permanently_if_requested(save_perm)
    
    return f"Uploaded {uploaded_file.name}\nSHA-256: {checksum}\nHash chain saved.\n{git_msg}"

# --- Gradio UI ---
with gr.Blocks() as app:
    gr.Markdown("## GradioVault: Upload Images & Generate Hash Chain")
    
    with gr.Row():
        uploaded_file = gr.File(label="Upload Image", file_types=[".png", ".jpg", ".jpeg"])
        save_perm = gr.Checkbox(label="Save permanently to repo?")
    
    output = gr.Textbox(label="Output", lines=10)
    
    submit_btn = gr.Button("Upload & Process")
    submit_btn.click(process_image, inputs=[uploaded_file, save_perm], outputs=output)

app.launch()