import gradio as gr
import cv2
import numpy as np
from src.ocr_service import DocumentOCRService
from src.document_config import HealthcareProcess, DocumentType
import json

# Initialize OCR service
ocr_service = DocumentOCRService()

def process_document(image, document_type, process_type):
    """Process document and return results."""
    try:
        # Convert image to numpy array
        image_np = np.array(image)
        
        # Convert to BGR for OpenCV
        if len(image_np.shape) == 3 and image_np.shape[2] == 3:
            image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
        
        # Process document
        result = ocr_service.process_document(image_np, document_type, process_type)
        
        # Format the output
        output = {
            "Extracted Data": result["extracted_data"],
            "Gemini Analysis": result["gemini_analysis"]["analysis"],
            "Validation Result": result["validation_result"]
        }
        
        return json.dumps(output, indent=2)
    except Exception as e:
        return f"Error processing document: {str(e)}"

def get_requirements(process_type):
    """Get document requirements for a process."""
    try:
        process = HealthcareProcess(process_type)
        requirements = ocr_service.get_process_requirements(process)
        return json.dumps(requirements, indent=2)
    except ValueError as e:
        return f"Error: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="TakeCare OCR Service") as demo:
    gr.Markdown("# TakeCare OCR Service")
    gr.Markdown("Upload your healthcare documents for processing and validation.")
    
    with gr.Tab("Process Document"):
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(type="pil", label="Upload Document")
                document_type = gr.Dropdown(
                    choices=[dt.value for dt in DocumentType],
                    label="Document Type"
                )
                process_type = gr.Dropdown(
                    choices=[pt.value for pt in HealthcareProcess],
                    label="Process Type (Optional)"
                )
                process_btn = gr.Button("Process Document")
            
            with gr.Column():
                output = gr.Textbox(label="Results", lines=20)
        
        process_btn.click(
            fn=process_document,
            inputs=[image_input, document_type, process_type],
            outputs=output
        )
    
    with gr.Tab("View Requirements"):
        with gr.Row():
            with gr.Column():
                req_process_type = gr.Dropdown(
                    choices=[pt.value for pt in HealthcareProcess],
                    label="Select Process Type"
                )
                view_req_btn = gr.Button("View Requirements")
            
            with gr.Column():
                requirements_output = gr.Textbox(label="Document Requirements", lines=20)
        
        view_req_btn.click(
            fn=get_requirements,
            inputs=[req_process_type],
            outputs=requirements_output
        )

# Launch the app
if __name__ == "__main__":
    demo.launch()