import gradio as gr
import pandas as pd
import PIL.Image
import io
import re
import os
from google import genai

# --- 1. The JS Hack to Force the Rear Native Camera ---
js_code = """
function() {
    // Continuously check for the upload button and force the camera intent
    setInterval(function() {
        let file_inputs = document.querySelectorAll('input[type="file"]:not([capture])');
        file_inputs.forEach(function(input) {
            input.setAttribute('capture', 'environment');
            input.setAttribute('accept', 'image/*');
        });
    }, 500);
}
"""

theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="slate",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
)

css = """
footer {display: none !important;}
.gradio-container {border-radius: 12px; box-shadow: 0 4px 6px rgba(0,0,0,0.1);}
"""

def extract_tables(image, api_key_input, progress=gr.Progress()):
    progress(0.1, desc="Initializing Gemini Vision...")
    
    api_key = api_key_input or os.environ.get("GEMINI_API_KEY")
    
    if not api_key:
        raise gr.Error("⚠️ Please enter your Gemini API Key in the box.")
    if image is None:
        raise gr.Error("⚠️ Please capture or upload a log sheet first.")

    try:
        progress(0.4, desc="Analyzing log sheet image...")
        client = genai.Client(api_key=api_key)
        
        prompt = """
        Carefully extract all the data tables from this dashboard image. 
        Format the output strictly as CSV (Comma Separated Values).
        For each table you find, provide the title of the table, and then enclose the CSV data inside ```csv ``` code blocks.
        Ensure column headers are included in the CSV.
        """
        
        response = client.models.generate_content(
            model="gemini-3-flash-preview", 
            contents=[prompt, image],
        )
        
        progress(0.8, desc="Parsing extracted tables...")
        csv_blocks = re.findall(r'```csv\n(.*?)\n```', response.text, re.DOTALL)
        
        if not csv_blocks:
            raise gr.Error(f"⚠️ Could not detect any tables. Raw Output:\n{response.text}")
            
        extracted_dfs = []
        for csv_data in csv_blocks:
            df = pd.read_csv(io.StringIO(csv_data))
            extracted_dfs.append(df)
            
        progress(1.0, desc="Done!")
        return extracted_dfs, gr.update(value="✅ **Extraction Complete!**", visible=True)
        
    except Exception as e:
        raise gr.Error(f"❌ An error occurred: {str(e)}")

# --- Build the UI ---
with gr.Blocks(title="Industrial Data Extractor") as demo:
    gr.HTML("<h2 style='text-align: center; color: #4338ca; margin-bottom: 0px;'>🏭 Industrial Data Extractor</h2>")
    gr.HTML("<p style='text-align: center; color: #64748b; margin-top: 5px;'>Capture log sheets to digitize tabular data.</p>")
    
    table_state = gr.State([])
    
    api_key_ui = gr.Textbox(
        label="Authentication", 
        type="password", 
        placeholder="Enter Gemini API Key here to begin..."
    )
    
    # We use "upload" only. The JS will hijack it and turn it into a camera button.
    image_input = gr.Image(
        type="pil", 
        label="Tap here to open Rear Camera", 
        sources=["upload"]
    )
    
    extract_btn = gr.Button("Extract Tables", variant="primary", size="lg")
    status_ui = gr.Markdown("", visible=False)
    
    gr.Markdown("---")
    gr.Markdown("### 📥 Extracted Data")
    
    @gr.render(inputs=table_state)
    def show_tables(dfs):
        if not dfs:
            gr.HTML("<div style='text-align: center; color: #94a3b8;'><i>Data will appear here.</i></div>")
        for i, df in enumerate(dfs):
            gr.Markdown(f"**Table {i+1}**")
            gr.Dataframe(value=df, interactive=False, wrap=True)

    extract_btn.click(
        fn=extract_tables, 
        inputs=[image_input, api_key_ui], 
        outputs=[table_state, status_ui]
    )
    
    # Inject the JS hack into the app
    demo.load(js=js_code)

demo.launch(theme=theme, css=css)