import nbformat
import gradio as gr
from transformers import pipeline


class NotebookEnhancer:
    def __init__(self):
        # Initialize Hugging Face models
        self.title_generator = pipeline(
            "summarization", model="facebook/bart-large-cnn"
        )
        self.summary_generator = pipeline(
            "summarization", model="sshleifer/distilbart-cnn-12-6"
        )

    def generate_title(self, code):
        """Generate a concise title for a code cell"""
        # Limit input length to match model constraints
        max_length = 1024
        truncated_code = code[:max_length] if len(code) > max_length else code

        result = self.title_generator(
            truncated_code, max_length=10, min_length=3, do_sample=False
        )
        title = result[0]["summary_text"].strip()
        # Format as a markdown title
        return f"## {title.capitalize()}"

    def generate_summary(self, code):
        """Generate a detailed summary for a code cell"""
        # Limit input length to match model constraints
        max_length = 1024
        truncated_code = code[:max_length] if len(code) > max_length else code

        result = self.summary_generator(
            truncated_code, max_length=100, min_length=30, do_sample=True
        )
        return result[0]["summary_text"].strip()

    def enhance_notebook(self, notebook_content):
        """Add title and summary markdown cells before each code cell"""
        # Load the notebook
        notebook = nbformat.reads(notebook_content, as_version=4)

        # Create a new notebook
        enhanced_notebook = nbformat.v4.new_notebook()
        enhanced_notebook.metadata = notebook.metadata

        # Process each cell
        i = 0
        while i < len(notebook.cells):
            cell = notebook.cells[i]

            # For code cells, add title and summary markdown cells
            if cell.cell_type == "code" and cell.source.strip():
                # Generate title
                title = self.generate_title(cell.source)
                title_cell = nbformat.v4.new_markdown_cell(title)
                enhanced_notebook.cells.append(title_cell)

                # Generate summary
                summary = self.generate_summary(cell.source)
                summary_cell = nbformat.v4.new_markdown_cell(summary)
                enhanced_notebook.cells.append(summary_cell)

            # Add the original cell
            enhanced_notebook.cells.append(cell)
            i += 1

        # Convert back to string
        return nbformat.writes(enhanced_notebook)


def process_notebook(file):
    """Process an uploaded notebook file"""
    enhancer = NotebookEnhancer()

    # Read uploaded file
    notebook_content = file.decode("utf-8")

    # Process the notebook
    enhanced_notebook = enhancer.enhance_notebook(notebook_content)

    # Save to temp file
    output_path = "enhanced_notebook.ipynb"
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(enhanced_notebook)

    return output_path


def build_gradio_interface():
    """Create and launch the Gradio interface"""
    with gr.Blocks(title="Notebook Enhancer") as demo:
        gr.Markdown("# Jupyter Notebook Enhancer")
        gr.Markdown(
            """
        Upload a Jupyter notebook to enhance it with automatically generated titles and summaries for each code cell.
        
        This tool uses Hugging Face models to:
        1. Generate concise titles for code cells
        2. Create explanatory summaries of what the code does
        """
        )

        with gr.Row():
            with gr.Column():
                file_input = gr.File(label="Upload Jupyter Notebook (.ipynb)")
                process_btn = gr.Button("Enhance Notebook")

            with gr.Column():
                output = gr.File(label="Enhanced Notebook")

        process_btn.click(fn=process_notebook, inputs=file_input, outputs=output)

    return demo


# This will be the entry point when running the script
if __name__ == "__main__":
    demo = build_gradio_interface()
    demo.launch()