import gradio as gr
import os
from src.processing.gemini_processor import GeminiProcessor
from src.analysis.coverage_generator import CoverageGenerator
from pathlib import Path
import logging

class ConsoleOutput:
    def __init__(self):
        self.messages = []
        
    def write(self, message):
        self.messages.append(str(message))
        return "\n".join(self.messages)
        
    def get_output(self):
        return "\n".join(self.messages)

def process_pdf(pdf_file, progress=gr.Progress()):
    if pdf_file is None:
        raise gr.Error("Please upload a PDF file")
    
    console = ConsoleOutput()
    
    # Set up logging to console
    logging.basicConfig(level=logging.DEBUG,
                       format='%(asctime)s - %(levelname)s - %(message)s')
    
    logger = logging.getLogger()
    logger.handlers = []  # Remove existing handlers
    
    class GradioHandler(logging.Handler):
        def emit(self, record):
            msg = self.format(record)
            console.write(msg)
            
    logger.addHandler(GradioHandler())
    
    try:
        processor = GeminiProcessor()
        coverage_gen = CoverageGenerator()
        
        # Clean screenplay
        progress(0.2, desc="Processing screenplay...")
        cleaned_path = Path("cleaned_screenplay_long.txt")
        success = processor.process_screenplay(pdf_file.name, str(cleaned_path))
        if not success:
            raise gr.Error("Failed to process screenplay")
        
        with open(cleaned_path, 'r') as f:
            cleaned_text = f.read()
        
        # Generate coverage
        progress(0.6, desc="Generating coverage...")
        success = coverage_gen.generate_coverage(cleaned_path)
        if not success:
            raise gr.Error("Failed to generate coverage")
            
        with open(Path("coverage.txt"), 'r') as f:
            coverage = f.read()
            
        progress(1.0, desc="Complete!")
        return cleaned_text, coverage, console.get_output()
        
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        console.write(error_msg)
        raise gr.Error(error_msg)

with gr.Blocks(title="Screenplay Coverage Generator") as demo:
    gr.Markdown("# Screenplay Coverage Generator")
    
    with gr.Row():
        file_input = gr.File(label="Upload Screenplay PDF", file_types=[".pdf"])
    
    with gr.Row():
        process_btn = gr.Button("Process Screenplay")
    
    with gr.Row():
        console = gr.Textbox(
            label="Console Output",
            lines=10,
            max_lines=30,
            autoscroll=True,
            container=True,
            show_copy_button=True
        )
    
    with gr.Tabs():
        with gr.TabItem("Cleaned Screenplay"):
            cleaned_output = gr.Textbox(
                label="Cleaned Screenplay",
                lines=10,
                max_lines=30,
                show_copy_button=True
            )
        with gr.TabItem("Coverage"):
            coverage_output = gr.Textbox(
                label="Coverage Document",
                lines=10,
                max_lines=30,
                show_copy_button=True
            )
    
    process_btn.click(
        fn=process_pdf,
        inputs=[file_input],
        outputs=[cleaned_output, coverage_output, console]
    )

if __name__ == "__main__":
    demo.launch()