import gradio as gr
from src.invoice_processor import InvoiceProcessor
from pathlib import Path

class InvoiceUI:
    def __init__(self):
        self.processor = InvoiceProcessor()
        self.current_json_output = None  # Store processed JSON for downloads

    def process_invoice(self, file_obj, fields_text: str):
        """Process the invoice and return JSON preview."""
        try:
            if not file_obj or not fields_text.strip():
                raise ValueError("Please provide both a file and fields to extract")
            
            # Parse fields
            user_fields = [f.strip() for f in fields_text.split(',')]
            
            # Process invoice
            invoice_text = self.processor.file_processor.process_file(file_obj.name)
            self.current_json_output = self.processor.process_invoice(invoice_text, user_fields)
            
            return self.current_json_output
            
        except Exception as e:
            raise gr.Error(str(e))

    def save_output(self, output_format: str):
        """Save current JSON output in specified format and return download path."""
        try:
            if self.current_json_output is None:
                raise gr.Error("Please process an invoice first")
            
            # Create output name using timestamp for uniqueness
            from datetime import datetime
            output_name = f"invoice_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            
            # Save in requested format
            output_path = self.processor.save_to_format(
                self.current_json_output,
                output_name,
                output_format
            )
            
            return output_path
            
        except Exception as e:
            raise gr.Error(str(e))

    def create_interface(self) -> gr.Blocks:
        """Create and return the Gradio interface."""
        with gr.Blocks(title="Invoice Processor") as app:
            gr.Markdown("# Invoice Data Processor")
            
            with gr.Row():
                with gr.Column():
                    # Input components
                    file_input = gr.File(
                        label="Upload Invoice File",
                        file_types=[".pdf", ".doc", ".docx", ".xlsx", ".txt", ".csv", ".xls"]
                    )
                    fields_input = gr.Textbox(
                        label="Fields to Extract (comma-separated)",
                        placeholder="e.g., invoice_number, date, amount, product_name"
                    )
                    process_btn = gr.Button("Process Invoice")
            
            with gr.Row():
                # JSON Preview
                json_output = gr.JSON(label="Extracted Data Preview")
            
            with gr.Row():
                with gr.Column():
                    # Download options
                    format_select = gr.Radio(
                        choices=['json', 'xml', 'csv', 'excel'],
                        label="Select Output Format",
                        value='json'
                    )
                    download_btn = gr.Button("Download in Selected Format")
                    download_output = gr.File(label="Download Processed File")
            
            # Connect components
            process_btn.click(
                fn=self.process_invoice,
                inputs=[file_input, fields_input],
                outputs=[json_output]
            )
            
            download_btn.click(
                fn=self.save_output,
                inputs=[format_select],
                outputs=[download_output]
            )
        
        return app

def main():
    """Initialize and launch the Gradio interface"""
    ui = InvoiceUI()
    app = ui.create_interface()
    app.launch(debug=True)

if __name__ == "__main__":
    main()