import gradio as gr
import os
import shutil
import zipfile
import tempfile
import time

# Import the process_all function from our script
from process_all_apartments import process_all, DEFAULT_PROMPT

def process_zip_upload(zip_file, api_key, speed_mode):
    if zip_file is None:
        return None, "Please upload a ZIP file."
        
    # Create a unique temporary directory
    temp_dir = tempfile.mkdtemp(prefix="hf_apartments_")
    
    try:
        # Extract the uploaded zip file
        extract_dir = os.path.join(temp_dir, "extracted")
        os.makedirs(extract_dir, exist_ok=True)
        
        with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
            zip_ref.extractall(extract_dir)
            
        # Sometimes extracting a zip gives a single root folder. 
        # We need to find the directory that contains the 'Apartment_XX' folders.
        # We'll just pass the extract_dir to process_all. If the user zipped the folders directly, 
        # they will be in extract_dir. If they zipped a parent folder, process_all will just 
        # iterate over that single parent folder (which might fail if it doesn't contain Before/After).
        # To be safe, let's find the first directory that contains multiple subdirectories or "Before"/"After".
        
        target_dir = extract_dir
        # Simple heuristic: if extract_dir has exactly one folder and no files, go inside it
        items = os.listdir(extract_dir)
        if len(items) == 1 and os.path.isdir(os.path.join(extract_dir, items[0])):
            target_dir = os.path.join(extract_dir, items[0])
            
        # Set speed parameters
        if speed_mode == 'super-fast':
            max_size, fps = 360, 0.5
        elif speed_mode == 'fast':
            max_size, fps = 480, 1.0
        else:
            max_size, fps = 720, 2.0
            
        api_key = api_key.strip() if api_key else None
        
        # Process the videos
        process_all(
            base_dir=target_dir,
            prompt=DEFAULT_PROMPT,
            gemini_api_key=api_key,
            max_size=max_size,
            fps=fps,
            force=True  # Always force reprocessing in HF Spaces
        )
        
        # Now zip the processed target_dir back
        output_zip_path = os.path.join(temp_dir, "Processed_Apartments_Reports.zip")
        shutil.make_archive(output_zip_path.replace('.zip', ''), 'zip', target_dir)
        
        return output_zip_path, f"Successfully processed in '{speed_mode}' mode! Download your reports below."
        
    except Exception as e:
        import traceback
        error_msg = f"Error during processing: {str(e)}\n\n{traceback.format_exc()}"
        return None, error_msg

# Build the Gradio UI
with gr.Blocks(title="Apartment Damage Detection & Inspection") as demo:
    gr.Markdown("# 🏢 Automated Apartment Video Inspector")
    gr.Markdown("Upload a ZIP file containing your apartment folders (e.g., `Apartment_01`, `Apartment_02`). Each folder should contain a `Before` and `After` folder with the respective `.mp4` videos.")
    
    with gr.Row():
        with gr.Column():
            zip_input = gr.File(label="Upload ZIP file of Apartments", file_types=[".zip"])
            api_key_input = gr.Textbox(label="Gemini API Key (Optional but recommended)", type="password", placeholder="AI Studio Key for damage inspection...")
            speed_input = gr.Radio(choices=["super-fast", "fast", "normal"], value="fast", label="Processing Speed Mode")
            process_btn = gr.Button("🚀 Process Videos", variant="primary")
            
        with gr.Column():
            status_output = gr.Textbox(label="Status Log", lines=10)
            zip_output = gr.File(label="Download Processed Reports")
            
    process_btn.click(
        fn=process_zip_upload,
        inputs=[zip_input, api_key_input, speed_input],
        outputs=[zip_output, status_output]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)