import csv
from pathlib import Path
import ipywidgets as widgets
from IPython.display import display

class BPMarker:
    """
    A notebook-based tool for annotating boundary and 
    prominence labels for each element (word) in a list.

    Boundary annotations support three possible values:
        - 0: No boundary
        - 1: Phrase boundary
        - 2: Sentence boundary

    Prominence annotations support two possible values:
        - 0: Not prominent
        - 1: Prominent

    Although the tool is designed specifically for 
    boundary and prominence annotation, it can be adapted
    for other use cases where a similar categorical 
    interpretation is appropriate.
    """

    def __init__(self, *args, title="", header="Boundary & Prominence Marking Tool") -> None:
        if len(args) == 0:
            raise ValueError(f"Expected one positonal argument, either list of labels or a csv file created using BPMaker")

        # =========================
        # Internal state variables
        # =========================
        self.data: list[list[str, int | None, int | None]] = [] # [[word, boundary, prominence]]

        self.header: str = header
        self.title: str = title

        # ========================
        # Parsing the user input
        # ========================
        
        # If user passes a list of labels
        if isinstance(args[0], list):
            labels = args[0]
            for label in labels:
                self.data.append([label, 0, 0])

        # If user passes a csv file (Label, Boundary, Prominence columns)
        # This is like loading an already marked file, good for editing
        elif isinstance(args[0], str):
            csv_fp = Path(args[0])
            with csv_fp.open(mode="r", encoding="utf-8") as f:
                content = csv.reader(f)

                # ======================
                # Sanity check
                # ======================
                EXPECTED_HEADER = ["Label", "Boundary", "Prominence"]
                header = next(content)
                if header != EXPECTED_HEADER:
                    raise RuntimeError("Invalid CSV file, please use the csv file that had been generated by the BP Marker earlier")
                    
                # ================
                # Start extracting
                # ================
                for row in content:
                    label, boundary, prominence = row
                    self.data.append([label, int(boundary), int(prominence)])
            
        else:
            raise ValueError("Invalid argument")

        # ==========================
        # Generate marking layout
        # ==========================
        self._generate_marking_layout()

        # ==========================
        # Display the entire layout
        # ==========================
        self._generate_complete_layout()

    def _generate_marking_layout(self):
        """
        Generates the layout based on the configurations.
        """
        
        self.label_buttons: list[widgets.Button] = []
        self.gap_buttons: list[widgets.Button] = []

        # ==========================================
        # Create button widgets for marking 
        # prominence and boundary
        # ==========================================
        for i, (label, boundary, prominence) in enumerate(self.data):

            # ======================
            # Label button
            # ======================
            label_button = widgets.Button(
                description=str(label),
                layout=widgets.Layout(width='auto', margin='2px 0px'),
                style={'button_color': 'white'},
            )

            # Add styling to the label button based on its state
            if prominence == 1:
                label_button.style.text_color = "purple"
                label_button.style.font_weight = "bold"
            else:
                label_button.style.text_color = "black"
                label_button.style.font_weight = "normal"
            
            
            label_button.idx: int = i # To keep track of which element it belongs to
            label_button.on_click(self._on_click_label_button)

            self.label_buttons.append(label_button)

            # ======================
            # Gap button
            # ======================
            gap_button = widgets.Button(
                description='-',
                layout=widgets.Layout(width='20px', margin='2px 0px', padding='0px'),
                style={'button_color': 'white'},
            )

            gap_button.idx: int = i
            gap_button.on_click(self._on_click_gap_button)

            # Add styling to the gap button based on its state
            if boundary == 0:
                gap_button.description = "-"
                gap_button.style.text_color = "black"
                gap_button.style.font_weight = "normal"
            elif boundary == 1:
                gap_button.description = "|"
                gap_button.style.text_color = "red"
                gap_button.style.font_weight = "bold"
            elif boundary == 2:
                gap_button.description = "||"
                gap_button.style.text_color = "red"
                gap_button.style.font_weight = "bold"
            
            self.gap_buttons.append(gap_button)

    def _on_click_label_button(self, button: widgets.Button):
        """
        Updates the prominence value of the data at
        the selected idx when the button is clicked.
        """
        
        self.data[button.idx][2] += 1 # index 2 at each row of data is holding prominence value for that row
        self.data[button.idx][2] %= 2 # index 2 at each row of data is holding prominence value for that row
 

        if self.data[button.idx][2] == 1:
            button.style.text_color = "purple"
            button.style.font_weight = "bold"
        else:
            button.style.text_color = "black"
            button.style.font_weight = "normal"

    def _on_click_gap_button(self, button: widgets.Button):
        """
        Updates the boundary value of the data at
        the selected idx when the button is clicked.
        """
        
        self.data[button.idx][1] += 1 # index 1 at each row of data is holding boundary value for that row
        self.data[button.idx][1] %= 3 # index 1 at each row of data is holding boundary value for that row
        
        if self.data[button.idx][1] == 0:
            button.description = "-"
            button.style.text_color = "black"
            button.style.font_weight = "normal"
        elif self.data[button.idx][1] == 1:
            button.description = "|"
            button.style.text_color = "red"
            button.style.font_weight = "bold"
        elif self.data[button.idx][1] == 2:
            button.description = "||"
            button.style.text_color = "red"
            button.style.font_weight = "bold"
   

    def _generate_complete_layout(self):
        """
        Generate the entire layout in a formatted layout with heading and styled container.
        """
        all_widgets = []
        
        for i in range(len(self.data)):
            all_widgets.append(self.label_buttons[i])
            all_widgets.append(self.gap_buttons[i])
        
        # Create heading
        heading = widgets.HTML(
            value=f"<h3 style='margin: 0 0 8px 0; color: #ef8c03; text-align: center;'>{self.header}</h3>"
        )

        # Create title
        title = widgets.HTML(
            value=f"<h4 style='margin: 0 0 8px 0; color: black; text-align: left;'>{self.title}</h3>"
        )
        
        # Create description
        description = widgets.HTML(
                value="""
                <div style='display: flex; justify-content: right; gap: 20px; margin-bottom: 12px;'>
                    <table style='font-size: 12px; border-collapse: collapse;'>
                        <thead>
                            <tr>
                                <th colspan='2' style='padding: 4px 8px; border: 1px solid #ddd; background-color: #f0f0f0; font-weight: bold;'>Prominence</th>
                            </tr>
                        </thead>
                        <tbody>
                            <tr>
                                <td style='padding: 4px 8px; border: 1px solid #ddd;'><span style='color: #555;'>word</span></td>
                                <td style='padding: 4px 8px; border: 1px solid #ddd;'>Not Prominent</td>
                            </tr>
                            <tr>
                                <td style='padding: 4px 8px; border: 1px solid #ddd;'><span style='color: purple; font-weight: bold;'>word</span></td>
                                <td style='padding: 4px 8px; border: 1px solid #ddd;'>Prominent</td>
                            </tr>
                        </tbody>
                    </table>
                    
                    <table style='font-size: 12px; border-collapse: collapse;'>
                        <thead>
                            <tr>
                                <th colspan='2' style='padding: 4px 8px; border: 1px solid #ddd; background-color: #f0f0f0; font-weight: bold;'>Boundary</th>
                            </tr>
                        </thead>
                        <tbody>
                            <tr>
                                <td style='padding: 4px 8px; border: 1px solid #ddd;'><span style='color: #555;'>-</span></td>
                                <td style='padding: 4px 8px; border: 1px solid #ddd;'>No Boundary</td>
                            </tr>
                            <tr>
                                <td style='padding: 4px 8px; border: 1px solid #ddd;'><span style='color: red; font-weight: bold;'>|</span></td>
                                <td style='padding: 4px 8px; border: 1px solid #ddd;'>Phrase Boundary</td>
                            </tr>
                            <tr>
                                <td style='padding: 4px 8px; border: 1px solid #ddd;'><span style='color: red; font-weight: bold;'>||</span></td>
                                <td style='padding: 4px 8px; border: 1px solid #ddd;'>Sentence Boundary</td>
                            </tr>
                        </tbody>
                    </table>
                </div>
                """
            )
        
        # Layout: wrap words nicely
        words_box = widgets.HBox(
            all_widgets,
            layout=widgets.Layout(
                display='flex',
                flex_flow='row wrap',
                align_items='center',
                width='100%',
                padding='12px',
                border='1px solid #ef8c03',
                border_radius='8px',
                background_color='#f8f9fa'
            )
        )
        
        # Combine everything in a VBox with rounded container
        container = widgets.VBox(
            [heading, description, title, words_box],
            layout=widgets.Layout(
                padding='20px',
                border='2px solid #ef8c03',
                border_radius='12px',
                background_color='white',
                box_shadow='0 2px 8px rgba(0,0,0,0.1)',
                width='auto',
                max_width='900px'
            )
        )
        
        display(container)

    def get_markings(self) -> list:
        """Returns the currently marked data"""
        return self.data.copy()

    def save(self, path: str | Path, overwrite: bool = False):
        """
        Save the markings as a csv file using the csv module.
        """
        path = Path(path)
    
        # Ensure the parent directory exists
        path.parent.mkdir(parents=True, exist_ok=True)
    
        # Determine mode: 'w' for overwrite, 'x' for exclusive creation
        mode = "w" if overwrite else "x"
    
        try:
            # newline='' is required by the csv module to prevent double line breaks
            with path.open(mode=mode, encoding="utf-8", newline='') as f:
                writer = csv.writer(f)
                
                # Header
                writer.writerow(['Label', 'Boundary', 'Prominence'])
                
                # Write the data rows
                # This handles escaping commas and special characters automatically
                writer.writerows(self.get_markings())
    
            print(f"Successfully saved to {path}")
    
        except FileExistsError:
            print(f"Error: The file '{path}' already exists and overwrite is set to False.")
