import math
import os
import re
from datetime import datetime
from pathlib import Path
from typing import List

import numpy as np
from PyQt5 import QtWidgets, QtCore

from hdsemg_pipe._log.log_config import logger
from hdsemg_shared.fileio.file_io import EMGFile, Grid
from hdsemg_pipe.state.global_state import global_state
from hdsemg_pipe.ui_elements.theme import Colors, Spacing, BorderRadius, Fonts, Styles


class AssociationDialog(QtWidgets.QDialog):
    def __init__(self, file_paths, parent=None):
        super().__init__(parent)
        logger.info("Initializing AssociationDialog with %d files", len(file_paths))
        self.file_paths = file_paths
        self.emg_files: List[EMGFile] = []
        self.load_files()
        self.init_ui()
        self.setWindowTitle("Grid Association Tool")
        self.resize(1000, 600)

    def load_files(self):
        logger.debug("Starting file loading process")
        for fp in self.file_paths:
            try:
                logger.info("Loading file: %s", fp)
                emg = EMGFile.load(filepath=fp)
                self.emg_files.append(emg)
                logger.debug("Extracted %d grids from %s", len(emg.grids), Path(fp).name)
                for grid in emg.grids:
                    logger.debug("Added grid %s from %s", grid.grid_key, emg.file_name)
            except Exception as e:
                logger.error(f"Failed to load {fp}: {str(e)}", exc_info=True)
                QtWidgets.QMessageBox.warning(self, "Loading Error", f"Failed to load {fp}:\n{str(e)}")
            logger.info("Total files loaded: %d", len(self.emg_files))

    def init_ui(self):
        """Initialize modern UI with GitHub-style design."""
        # Set dialog background
        self.setStyleSheet(f"""
            QDialog {{
                background-color: {Colors.BG_SECONDARY};
            }}
        """)

        # Create list widgets with modern styling
        self.available_list = QtWidgets.QListWidget()
        self.selected_list = QtWidgets.QListWidget()
        self.selected_list.setDragDropMode(QtWidgets.QAbstractItemView.InternalMove)

        # Style both lists
        list_style = f"""
            QListWidget {{
                background-color: {Colors.BG_PRIMARY};
                border: 1px solid {Colors.BORDER_DEFAULT};
                border-radius: {BorderRadius.MD};
                padding: {Spacing.SM}px;
                font-size: {Fonts.SIZE_BASE};
            }}
            QListWidget::item {{
                padding: {Spacing.SM}px;
                border-radius: {BorderRadius.SM};
            }}
            QListWidget::item:hover {{
                background-color: {Colors.GRAY_100};
            }}
            QListWidget::item:selected {{
                background-color: {Colors.BLUE_100};
                color: {Colors.BLUE_900};
            }}
        """
        self.available_list.setStyleSheet(list_style)
        self.selected_list.setStyleSheet(list_style)

        # Populate available list
        for emg in self.emg_files:
            for grid in emg.grids:
                # Build item text with muscle info if available
                muscle_info = f" - {grid.muscle}" if grid.muscle else ""
                item_text = f"{grid.rows}x{grid.cols} Grid ({len(grid.ref_indices)} refs){muscle_info} - {emg.file_name}"
                item = QtWidgets.QListWidgetItem(item_text)

                # Store grid data as dictionary with all necessary info for association
                grid_data = {
                    'rows': grid.rows,
                    'cols': grid.cols,
                    'ied_mm': grid.ied_mm,
                    'electrodes': grid.electrodes,
                    'emg_indices': grid.emg_indices,
                    'ref_indices': grid.ref_indices,
                    'grid_key': grid.grid_key,
                    'muscle': grid.muscle,  # Include muscle information
                    'data': emg.data,
                    'time': emg.time,
                    'description': emg.description,
                    'sf': emg.sampling_frequency,
                    'file_name': emg.file_name,
                    'file_path': emg.file_name,  # This is actually just the filename
                }
                item.setData(QtCore.Qt.UserRole, grid_data)
                self.available_list.addItem(item)

        # Transfer buttons with icon-style
        self.add_btn = QtWidgets.QPushButton("Add →")
        self.remove_btn = QtWidgets.QPushButton("← Remove")
        self.add_btn.setStyleSheet(Styles.button_secondary())
        self.remove_btn.setStyleSheet(Styles.button_secondary())
        self.add_btn.clicked.connect(self.add_selected)
        self.remove_btn.clicked.connect(self.remove_selected)

        # Name input with modern styling
        self.name_edit = QtWidgets.QLineEdit()
        self.name_edit.setPlaceholderText("Enter association name...")
        self.name_edit.setStyleSheet(Styles.input_field())

        # Action buttons
        self.add_assoc_btn = QtWidgets.QPushButton("+ Add Another")
        self.add_assoc_btn.setToolTip("Save current association and add another")
        self.add_assoc_btn.setStyleSheet(Styles.button_secondary())
        self.save_close_btn = QtWidgets.QPushButton("Save & Close")
        self.save_close_btn.setStyleSheet(Styles.button_primary())

        self.add_assoc_btn.clicked.connect(lambda: self.save_association(close_dialog=False))
        self.save_close_btn.clicked.connect(lambda: self.save_association(close_dialog=True))

        # Saved files list
        self.saved_files_list = QtWidgets.QListWidget()
        self.saved_files_list.setToolTip("Successfully saved associations")
        self.saved_files_list.setStyleSheet(list_style)

        # Create labels with modern styling
        def create_section_label(text):
            label = QtWidgets.QLabel(text)
            label.setStyleSheet(f"""
                QLabel {{
                    color: {Colors.TEXT_PRIMARY};
                    font-size: {Fonts.SIZE_LG};
                    font-weight: {Fonts.WEIGHT_SEMIBOLD};
                    margin-bottom: {Spacing.SM}px;
                }}
            """)
            return label

        # Main layout with proper spacing
        main_layout = QtWidgets.QVBoxLayout()
        main_layout.setContentsMargins(Spacing.LG, Spacing.LG, Spacing.LG, Spacing.LG)
        main_layout.setSpacing(Spacing.LG)

        # Top section: Grid selection
        grid_layout = QtWidgets.QHBoxLayout()
        grid_layout.setSpacing(Spacing.MD)

        # Left panel (Available grids)
        left_panel = QtWidgets.QVBoxLayout()
        left_panel.setSpacing(Spacing.SM)
        left_panel.addWidget(create_section_label("Available Grids"))
        left_panel.addWidget(self.available_list)

        # Button panel (center)
        btn_panel = QtWidgets.QVBoxLayout()
        btn_panel.addStretch()
        btn_panel.addWidget(self.add_btn)
        btn_panel.addSpacing(Spacing.SM)
        btn_panel.addWidget(self.remove_btn)
        btn_panel.addStretch()

        # Right panel (Selected grids)
        right_panel = QtWidgets.QVBoxLayout()
        right_panel.setSpacing(Spacing.SM)

        selected_label_layout = QtWidgets.QVBoxLayout()
        selected_label = create_section_label("Selected Grids")
        order_hint = QtWidgets.QLabel("(Order matters - drag to reorder)")
        order_hint.setStyleSheet(f"""
            QLabel {{
                color: {Colors.TEXT_SECONDARY};
                font-size: {Fonts.SIZE_SM};
                font-style: italic;
            }}
        """)
        selected_label_layout.addWidget(selected_label)
        selected_label_layout.addWidget(order_hint)
        selected_label_layout.setSpacing(0)

        right_panel.addLayout(selected_label_layout)
        right_panel.addWidget(self.selected_list)

        grid_layout.addLayout(left_panel, 1)
        grid_layout.addLayout(btn_panel, 0)
        grid_layout.addLayout(right_panel, 1)

        main_layout.addLayout(grid_layout)

        # Middle section: Association name input
        name_group = QtWidgets.QGroupBox("Association Details")
        name_group.setStyleSheet(Styles.groupbox())
        name_layout = QtWidgets.QVBoxLayout()
        name_layout.setSpacing(Spacing.SM)

        name_label = QtWidgets.QLabel("Association Name:")
        name_label.setStyleSheet(f"color: {Colors.TEXT_PRIMARY}; font-weight: {Fonts.WEIGHT_MEDIUM};")
        name_layout.addWidget(name_label)
        name_layout.addWidget(self.name_edit)

        name_group.setLayout(name_layout)
        main_layout.addWidget(name_group)

        # Bottom section: Saved associations
        saved_group = QtWidgets.QGroupBox("Saved Associations")
        saved_group.setStyleSheet(Styles.groupbox())
        saved_layout = QtWidgets.QVBoxLayout()
        saved_layout.setSpacing(Spacing.SM)
        saved_layout.addWidget(self.saved_files_list)
        saved_group.setLayout(saved_layout)
        main_layout.addWidget(saved_group)

        # Button row at bottom
        button_row = QtWidgets.QHBoxLayout()
        button_row.setSpacing(Spacing.SM)
        button_row.addWidget(self.add_assoc_btn)
        button_row.addStretch()
        button_row.addWidget(self.save_close_btn)

        main_layout.addLayout(button_row)

        self.setLayout(main_layout)

    def add_selected(self):
        selected_item = self.available_list.currentItem()
        if selected_item:
            logger.debug("Adding grid to selection: %s", selected_item.text())
            new_item = selected_item.clone()
            self.selected_list.addItem(new_item)
            logger.info("Current selection count: %d", self.selected_list.count())

            # Auto-suggest muscle name if all selected grids have the same muscle
            self.update_name_suggestion()

    def remove_selected(self):
        current_row = self.selected_list.currentRow()
        if current_row >= 0:
            item = self.selected_list.item(current_row)
            logger.debug("Removing grid from selection: %s", item.text())
            self.selected_list.takeItem(current_row)
            logger.info("Current selection count: %d", self.selected_list.count())

            # Update name suggestion after removal
            self.update_name_suggestion()

    def update_name_suggestion(self):
        """
        Auto-suggest association name based on muscle name if all selected grids
        have the same muscle. Only updates if the name field is empty or contains
        a previous auto-suggestion.
        """
        # Only auto-suggest if there are selected grids
        if self.selected_list.count() == 0:
            return

        # Collect all muscle names from selected grids
        muscles = []
        for i in range(self.selected_list.count()):
            item = self.selected_list.item(i)
            grid_data = item.data(QtCore.Qt.UserRole)
            muscle = grid_data.get('muscle')
            if muscle:
                muscles.append(muscle)

        # Only suggest if we have muscle info and all muscles are the same
        if len(muscles) > 0 and len(set(muscles)) == 1:
            # All selected grids have the same muscle
            suggested_name = muscles[0]

            # Only update if:
            # 1. Field is empty, OR
            # 2. Field contains a previous muscle name (to allow updates when adding/removing)
            current_name = self.name_edit.text().strip()

            # Check if current name is empty or looks like a previous muscle suggestion
            # (simple heuristic: single word, capitalized, or matches a known muscle pattern)
            is_likely_auto_suggestion = (
                not current_name or  # Empty
                (current_name and not ' ' in current_name and current_name[0].isupper())  # Single capitalized word
            )

            if is_likely_auto_suggestion:
                self.name_edit.setText(suggested_name)
                logger.debug("Auto-suggested association name: %s", suggested_name)
        else:
            # Mixed muscles or no muscle info - clear suggestion only if it was auto-generated
            current_name = self.name_edit.text().strip()
            if current_name and not ' ' in current_name and current_name[0].isupper():
                # Looks like it might have been auto-suggested, but now it's not valid
                # Don't clear it - let user decide
                pass

    def save_association(self, close_dialog=True):
        logger.info("Starting save association process")
        selected_grids = []
        for i in range(self.selected_list.count()):
            item = self.selected_list.item(i)
            selected_grids.append(item.data(QtCore.Qt.UserRole))

        assoc_name = self.name_edit.text().strip()

        if not selected_grids and not assoc_name and close_dialog:
            logger.info("No grids or association name provided. Closing dialog.")
            self.accept()
            return

        logger.debug("Selected grids count: %d", len(selected_grids))
        if not selected_grids:
            logger.warning("Save attempted with no grids selected")
            QtWidgets.QMessageBox.warning(self, "Error", "No grids selected for association!")
            return

        if not assoc_name:
            logger.warning("Save attempted with empty association name")
            QtWidgets.QMessageBox.warning(self, "Error", "Please enter an association name!")
            return

        # --- Überprüfe, dass Zeitvektoren und Samplingfrequenzen übereinstimmen ---
        base_time = selected_grids[0]['time']
        base_sf = selected_grids[0]['sf']
        for grid in selected_grids[1:]:
            if not np.array_equal(base_time, grid['time']):
                logger.error("Time vector mismatch detected for file: %s", grid['file_name'])
                QtWidgets.QMessageBox.critical(self, "Time Error", "Time vectors mismatch between selected grids!")
                return
            if base_sf != grid['sf']:
                logger.error("Sampling frequency mismatch detected for file: %s", grid['file_name'])
                QtWidgets.QMessageBox.critical(self, "Sampling Frequency Error",
                                               "Sampling frequency mismatch between selected grids!")
                return

        # --- Berechne die neue Grid-Größe anhand der Gesamtzahl der EMG-Kanäle ---
        total_emg_channels = sum(len(grid['emg_indices']) for grid in selected_grids)
        new_rows, new_cols = compute_new_grid_size(total_emg_channels)
        logger.info("New combined EMG grid dimensions will be: %dx%d", new_rows, new_cols)

        # --- Bestimme die IED für das kombinierte Grid ---
        # Wenn alle Grids die gleiche IED haben, verwende diese
        # Sonst verwende den kleinsten gemeinsamen Nenner (z.B. Durchschnitt)
        ieds = [grid['ied_mm'] for grid in selected_grids]
        if len(set(ieds)) == 1:
            # Alle haben gleiche IED
            combined_ied = ieds[0]
            logger.info("All grids have same IED: %dmm", combined_ied)
        else:
            # Verschiedene IEDs - verwende den kleinsten oder Durchschnitt
            combined_ied = min(ieds)  # Kleinster IED (konservativer Ansatz)
            logger.warning("Mixed IEDs detected (%s). Using smallest IED: %dmm for combined grid",
                         set(ieds), combined_ied)

        try:
            # --- Initialisiere Container für die kombinierten Daten ---
            combined_emg_data = None
            combined_ref_data = None
            combined_emg_desc = []
            combined_ref_desc = []
            combined_grid_info = {}
            emg_current_index = 0  # Für die fortlaufende Kanalindizierung im neuen Grid
            ref_current_index = 0
            combined_time = base_time
            combined_sf = base_sf

            # Regex zum Aktualisieren der Grid-Dimension und IED im Beschreibungsstring
            pattern = re.compile(r"(HD)(\d{2})(MM)(\d{2})(\d{2})")

            for grid in selected_grids:
                logger.debug("Processing grid from file: %s", grid['file_name'])

                # --- EMG- und Referenzdaten extrahieren ---
                emg_data_slice = grid['data'][:, grid['emg_indices']]
                ref_data_slice = grid['data'][:, grid['ref_indices']]
                logger.debug("Extracted %d EMG channels and %d reference channels from %s",
                             len(grid['emg_indices']), len(grid['ref_indices']), grid['file_name'])

                # --- Erzeuge aktualisierte EMG-Beschreibungen mit einheitlicher IED und neuer Grid-Dimension ---
                emg_desc_list = []
                for idx in grid['emg_indices']:
                    desc_str = extract_description(grid, idx)
                    new_dims = f"{new_rows:02d}{new_cols:02d}"  # z.B. "1208" für ein 12x8-Grid
                    new_ied = f"{combined_ied:02d}"  # z.B. "08" für 8mm
                    # Ersetze HDXXMMYYZZ mit HD{new_ied}MM{new_dims}
                    # z.B. HD08MM0408 -> HD08MM1208 oder HD10MM0408 -> HD08MM1208
                    updated_desc = pattern.sub(rf"HD{new_ied}MM{new_dims}", desc_str)
                    emg_desc = f"{updated_desc}-{Path(grid['file_path']).stem}"
                    emg_desc_list.append(emg_desc)
                logger.debug("Updated EMG descriptions for %s: %s", grid['file_name'], emg_desc_list[:3])

                # --- Referenzbeschreibungen erstellen ---
                ref_desc_list = []
                for idx in grid['ref_indices']:
                    ref_desc = f"refSig-{Path(grid['file_path']).stem}-{grid['description'][idx]}"
                    ref_desc_list.append(ref_desc)
                logger.debug("Reference descriptions for %s: %s", grid['file_name'], ref_desc_list)

                # --- Daten kombinieren ---
                if combined_emg_data is None:
                    combined_emg_data = emg_data_slice
                    combined_ref_data = ref_data_slice
                else:
                    combined_emg_data = np.hstack((combined_emg_data, emg_data_slice))
                    combined_ref_data = np.hstack((combined_ref_data, ref_data_slice))
                logger.debug("Combined EMG data shape now: %s", combined_emg_data.shape)
                logger.debug("Combined reference data shape now: %s", combined_ref_data.shape)

                combined_emg_desc.extend(emg_desc_list)
                combined_ref_desc.extend(ref_desc_list)

                # --- Update der Grid-Info für das jeweilige Grid ---
                grid_key = f"{grid['rows']}x{grid['cols']}"
                suffix = 1
                while grid_key in combined_grid_info:
                    grid_key = f"{grid['rows']}x{grid['cols']}_{suffix}"
                    suffix += 1
                combined_grid_info[grid_key] = {
                    "original_rows": grid['rows'],
                    "original_cols": grid['cols'],
                    "allocated_emg_channels": len(emg_desc_list),
                    "emg_indices": list(range(emg_current_index, emg_current_index + len(grid['emg_indices']))),
                    "ied_mm": grid['ied_mm'],
                    "electrodes": grid['electrodes']
                }
                logger.debug("Updated grid info for %s: %s", grid_key, combined_grid_info[grid_key])
                emg_current_index += len(grid['emg_indices'])

                # --- Update der Referenz-Signale ---
                if "reference_signals" not in combined_grid_info:
                    combined_grid_info["reference_signals"] = []
                combined_grid_info["reference_signals"].append({
                    "file_name": grid['file_name'],
                    "indices": list(range(ref_current_index, ref_current_index + len(grid['ref_indices'])))
                })
                logger.debug("Added reference signal info for file %s", grid['file_name'])
                ref_current_index += len(grid['ref_indices'])

            # --- Gesamtinformation zur neuen EMG-Grid-Dimension hinzufügen ---
            combined_grid_info["combined_emg_grid"] = {"rows": new_rows, "cols": new_cols}

            # --- Endgültige Kombination ---
            combined_data = np.hstack((combined_emg_data, combined_ref_data))
            combined_description = combined_emg_desc + combined_ref_desc
            combined_description = np.array([[d] for d in combined_description], dtype=object)
            logger.debug("Final combined data shape: %s", combined_data.shape)
            logger.debug("Final combined description length: %d", len(combined_description))

            # --- Dateien speichern ---
            workfolder = global_state.get_associated_grids_path()
            filename = format_filename(assoc_name)
            save_path = os.path.join(workfolder, filename)
            logger.info("Saving association '%s' to MAT file: %s", assoc_name, save_path)

            # Calculate file size (approximate, will be updated after save)
            file_size = combined_data.nbytes + combined_time.nbytes

            # Create EMGFile with all required parameters
            assoc_emg_file = EMGFile(
                data=combined_data,
                time=combined_time,
                description=combined_description,
                sf=combined_sf,
                file_name=filename,
                file_size=file_size,
                file_type="mat"
            )

            assoc_emg_file.save(save_path)

            logger.info("MAT file saved successfully at %s", save_path)

            save_association_json(save_path, assoc_name, selected_grids, combined_grid_info)
            logger.info("JSON metadata file saved successfully for association '%s'", assoc_name)

            self.saved_files_list.addItem(filename)
            global_state.associated_files.append(save_path)
            QtWidgets.QMessageBox.information(self, "Success",
                                              f"Association saved successfully!\n{filename}")

            if not close_dialog:
                self.selected_list.clear()
                self.name_edit.clear()
            else:
                logger.info("Closing association dialog for '%s'", assoc_name)
                self.accept()

        except Exception as e:
            logger.error("Failed to save association: %s", str(e), exc_info=True)
            QtWidgets.QMessageBox.critical(self, "Error", f"Failed to save association:\n{str(e)}")


def compute_new_grid_size(total_channels):
    """
    Berechnet neue Grid-Dimensionen (rows, cols), sodass rows * cols == total_channels.
    Es wird der Divisor gewählt, der der Quadratwurzel von total_channels am nächsten kommt.
    """
    # Ermittele alle Divisoren von total_channels
    divisors = [i for i in range(1, total_channels + 1) if total_channels % i == 0]
    sqrt_val = math.sqrt(total_channels)
    best_divisor = min(divisors, key=lambda x: abs(x - sqrt_val))
    new_cols = best_divisor
    new_rows = total_channels // new_cols
    return new_rows, new_cols


import json


def save_association_json(save_file_path, assoc_name, selected_grids, combined_grid_info):
    """
    Save association metadata to a JSON file.

    Parameters:
        save_file_path (str): The file path used for the .mat file.
        assoc_name (str): The name of the association.
        selected_grids (list): List of grid dictionaries with original grid info.
        combined_grid_info (dict): The combined grid info structure used in the MAT file.
    """
    metadata = {
        "association_name": assoc_name,
        "timestamp": datetime.now().isoformat(),
        "grids": [],
        "combined_grid_info": combined_grid_info
    }

    for grid in selected_grids:
        grid_data = {
            "file_name": grid['file_name'],
            "file_path": grid['file_path'],
            "rows": grid['rows'],
            "cols": grid['cols'],
            "emg_count": len(grid['emg_indices']),
            "ref_count": len(grid['ref_indices']),
            "ied_mm": grid['ied_mm'],
            "electrodes": grid['electrodes'],
            "muscle": grid.get('muscle')  # Include muscle information if available
        }
        metadata["grids"].append(grid_data)

    # Create a JSON file path based on the .mat file path.
    json_file_path = save_file_path.replace(".mat", ".json")

    with open(json_file_path, "w") as json_file:
        json.dump(metadata, json_file, indent=4)

    logger.info("JSON metadata file saved successfully: %s", json_file_path)


def sanitize_filename(filename, replace_with="_"):
    # Define forbidden characters for Windows & Linux
    forbidden_chars = r'<>:"/\\|?*\x00-\x1F'  # Control characters and special ones
    filename = re.sub(f"[{re.escape(forbidden_chars)}]", replace_with, filename)

    # Remove leading/trailing spaces and dots
    filename = filename.strip().strip(".")

    # Ensure filename is not empty or just dots
    if not filename or filename in {".", ".."}:
        filename = "associated_grids"

    return filename


def extract_description(grid, idx):
    # Try to extract the description string robustly.
    try:
        if isinstance(grid['description'], np.ndarray):
            if grid['description'].ndim == 2:
                desc_candidate = grid['description'][idx, 0]
            else:
                desc_candidate = grid['description'][idx]
        else:
            desc_candidate = grid['description'][idx]
        try:
            desc_str = desc_candidate.item()
        except AttributeError:
            desc_str = desc_candidate

        # Convert bytes to string if necessary
        if isinstance(desc_str, bytes):
            desc_str = desc_str.decode('utf-8', errors='ignore')
        # Ensure it's a string
        elif not isinstance(desc_str, str):
            desc_str = str(desc_str)

    except Exception as e:
        logger.error("Failed to extract description for index %d in file %s: %s",
                     idx, grid['file_name'], str(e))
        desc_str = ""

    return desc_str


def format_filename(filename):
    # Replace spaces with underscores and make it lowercase
    filename = filename.replace(" ", "_").lower()

    # Sanitize the filename
    filename = sanitize_filename(filename)

    # Append a timestamp
    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

    return f"{filename}_{timestamp}.mat"
