# components/notebook_ui.py
import streamlit as st
from typing import Dict, List, Optional
import json

class NotebookUI:
    def __init__(self):
        if 'notebook_cells' not in st.session_state:
            st.session_state.notebook_cells = []

    def add_cell(self, cell_type: str, content: Dict, metadata: Optional[Dict] = None):
        """Add a new cell to the notebook"""
        cell = {
            'type': cell_type,
            'content': content,
            'metadata': metadata or {},
            'id': len(st.session_state.notebook_cells)
        }
        st.session_state.notebook_cells.append(cell)

    def render_cells(self):
        """Render all notebook cells"""
        for cell in st.session_state.notebook_cells:
            self._render_cell(cell)

    def _render_cell(self, cell: Dict):
        """Render a single cell"""
        with st.container():
            with st.expander(f"Cell {cell['id'] + 1}", expanded=True):
                if cell['type'] == 'query':
                    self._render_query_cell(cell)
                elif cell['type'] == 'analysis':
                    self._render_analysis_cell(cell)
                elif cell['type'] == 'document':
                    self._render_document_cell(cell)

    def _render_query_cell(self, cell: Dict):
        st.markdown("**Query:**")
        st.write(cell['content']['query'])
        st.markdown("**Response:**")
        st.write(cell['content']['response'])
        
        col1, col2 = st.columns([1, 4])
        with col1:
            if st.button("♻️ Regenerate", key=f"regen_{cell['id']}"):
                pass  # Implement regeneration logic
            if st.button("📋 Copy", key=f"copy_{cell['id']}"):
                st.toast("Copied to clipboard!")

    def _render_analysis_cell(self, cell: Dict):
        st.markdown(f"**{cell['metadata'].get('template', 'Analysis')}**")
        for section in cell['content']:
            st.markdown(f"### {section['title']}")
            st.write(section['content'])

    def _render_document_cell(self, cell: Dict):
        st.markdown(f"**Document: {cell['metadata'].get('filename', 'Unknown')}**")
        with st.expander("Document Content"):
            st.write(cell['content']['text'])
        st.markdown("### Metadata")
        st.json(cell['metadata'])