from PySide6.QtWidgets import QMainWindow, QToolBar, QWidget, QVBoxLayout, QLabel, QPushButton, QLineEdit, QMenu, QHBoxLayout, QFrame, QSizePolicy, QTextEdit, QSplitter, QScrollArea, QWidgetAction
from PySide6.QtCore import Qt, Signal
from factum import config
from factum.ui.stylesheet import get_stylesheet

class FactumWindow(QMainWindow):
    def __init__(self, db_service):
        super().__init__()
        self.db_service = db_service
        self.config = config.settings
        self.create_ui()
        self.connect_ui_elements()
        
    def create_ui(self):
        self.setWindowTitle("Factum")
        self.setGeometry(100, 100, 800, 600)
        # load ai-generated stylesheet
        # we use setObjectName throughout to reference the stylesheet
        self.setStyleSheet(get_stylesheet())

        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # Create menu bar with branding toolbar
        self.create_menu_bar()
        
        self.main_layout = QHBoxLayout()
        # remove any margings and spacing to create a cohesive ui
        self.set_zero_spacing(self.main_layout)
        
        # add the left sidebar which has a QFrame as border
        left_sidebar = QFrame()
        left_sidebar.setObjectName("case_sidebar")
        left_sidebar.setProperty("class", "sidebar")
        left_sidebar.setLayout(self.create_casetree_layout())
        self.main_layout.addWidget(left_sidebar, 0)


        # add the center area which has a QFrame as border
        center_area = QFrame()
        center_area.setObjectName("note_area")
        center_area.setProperty("class", "centerArea")
        center_area.setLayout(self.create_central_area_layout())
        self.main_layout.addWidget(center_area, 1)

        # add the right sidebar which has a QFrame as border
        right_sidebar = QFrame()
        right_sidebar.setObjectName("right_sidebar")
        right_sidebar.setProperty("class", "sidebar")
        right_sidebar.setLayout(self.create_info_widget_layout())
        self.main_layout.addWidget(right_sidebar, 0)
        
        central_widget.setLayout(self.main_layout)
        self.create_bottom_toolbar()
        
    def create_casetree_layout(self):
        tree_layout = QVBoxLayout()
        self.set_zero_spacing(tree_layout)

        # Header (might change)
        header_toolbar = QToolBar()
        header_toolbar.setMovable(False)
        header_toolbar.setFloatable(False)
        header_toolbar.setContextMenuPolicy(Qt.ContextMenuPolicy.PreventContextMenu)
        header_toolbar.setObjectName("sidebarHeader")
        header_label = QLabel("  CASE STRUCTURE")
        header_label.setObjectName("sidebarHeaderTitle")
        header_toolbar.addWidget(header_label)
        tree_layout.addWidget(header_toolbar)
        
        # tree content area
        tree_content = QLabel("Case Tree Placeholder")
        tree_layout.addWidget(tree_content, 1)
        
        # button area with margings
        button_frame = QFrame()
        button_frame.setObjectName("sidebarButtonArea")
        button_frame.setContentsMargins(10, 10, 10, 10)
        button_layout = QVBoxLayout()
        self.set_zero_spacing(button_layout)
        
        # and space between buttons
        button_layout.setSpacing(8)
        new_case_button = QPushButton("+ New Case")
        new_case_button.setProperty("buttonStyle", "primary")
        delete_case_button = QPushButton("- Delete Case")
        delete_case_button.setProperty("buttonStyle", "danger")
        button_layout.addWidget(new_case_button)
        button_layout.addWidget(delete_case_button)
        button_frame.setLayout(button_layout)
        tree_layout.addWidget(button_frame)
        
        return tree_layout
    
    def create_central_area_layout(self):
        note_area_layout = QVBoxLayout()
        self.set_zero_spacing(note_area_layout)

        # filter toolbar setup
        filter_toolbar = QToolBar()
        filter_toolbar.setMovable(False)
        filter_toolbar.setFloatable(False)
        filter_toolbar.setContextMenuPolicy(Qt.ContextMenuPolicy.PreventContextMenu)
        filter_toolbar.setObjectName("filterBar")
        filter_label = QLabel("  Filter by tag:")
        filter_toolbar.addWidget(filter_label)
        note_area_layout.addWidget(filter_toolbar)
        
        # the area where notes are displayed (placeholder)
        note_display = QLabel("Notes display area")
        note_area_layout.addWidget(note_display, 1)

        # note input area with frame
        note_input_frame = QFrame()
        note_input_frame.setObjectName("noteInputArea")
        note_input_layout = QVBoxLayout()
        # text input is multiline
        note_input_text = QTextEdit()
        note_input_text.setObjectName("note_input")
        note_input_text.setPlaceholderText("Notes will be time-stamped. Add #tags if needed.")
        note_input_text.setMaximumHeight(96)
        note_input_layout.addWidget(note_input_text)
        
        # submit button toolbar with status text
        bottom_row = QHBoxLayout()
        bottom_row.setSpacing(8)
        bottom_row.setContentsMargins(0,16,0,4)
        status_label = QLabel("IoCs will be auto-extracted | GPG signing enabled") #TODO real status from current settings
        status_label.setObjectName("noteInputStatus")
        add_button = QPushButton("Add Note")
        add_button.setProperty("buttonStyle", "primary")
        bottom_row.addWidget(status_label)
        bottom_row.addStretch()
        bottom_row.addWidget(add_button)
        note_input_layout.addLayout(bottom_row)
        note_input_frame.setLayout(note_input_layout)
        note_area_layout.addWidget(note_input_frame)
        
        return note_area_layout

    def create_info_widget_layout(self):
        info_widget_layout = QVBoxLayout()
        self.set_zero_spacing(info_widget_layout)
        
        # the resize_separator enables the drag-resize
        self.info_splitter = QSplitter(Qt.Orientation.Vertical)
        self.info_splitter.setObjectName("infoSplitter")
        self.info_splitter.setChildrenCollapsible(False)  # this keeps the sections from disappearing
        
        # the investigative goals section can be collapsed
        goals_widget = self.create_collapsible_section("Investigative Goals", self.create_goals_content())
        self.info_splitter.addWidget(goals_widget)
        
        # pinned notes are always visible
        pinned_widget = self.create_section("Pinned Notes", self.create_pinned_content())
        self.info_splitter.addWidget(pinned_widget)
        
        # extracted iocs can be collapsed
        iocs_widget = self.create_collapsible_section("Extracted IoCs", self.create_iocs_content())
        self.info_splitter.addWidget(iocs_widget)
        
        # Set initial sizes (in pixels) for each section
        self.info_splitter.setSizes([150, 150, 150])
        
        # TODO: add a section for the ai assistant if we choose to implement it
        
        info_widget_layout.addWidget(self.info_splitter, 1)
        
        return info_widget_layout
    
    def create_collapsible_section(self, title, content_widget):
        section_frame = QFrame()
        section_frame.setObjectName("infoSection")
        section_layout = QVBoxLayout()
        self.set_zero_spacing(section_layout)
        
        # header has toggle button
        header = QFrame()
        header.setObjectName("infoSectionHeader")
        header_layout = QHBoxLayout()
        header_layout.setContentsMargins(8, 4, 8, 4)
        
        toggle_button = QPushButton("▼")
        toggle_button.setObjectName("toggleButton")
        toggle_button.setFixedWidth(24)
        toggle_button.setProperty("collapsed", False)
        # store the id references (python memory addresses) to retrieve later
        toggle_button.setProperty("content_widget", id(content_widget))
        toggle_button.setProperty("section_frame", id(section_frame))
        
        title_label = QLabel(title)
        title_label.setObjectName("infoSectionTitle")
        
        header_layout.addWidget(toggle_button)
        header_layout.addWidget(title_label)
        header_layout.addStretch()
        header.setLayout(header_layout)
        
        # Store widget references in instance dict and connect toggle function
        if not hasattr(self, 'toggle_widgets'):
            self.toggle_widgets = {}
        if not hasattr(self, 'section_frames'):
            self.section_frames = {}
        self.toggle_widgets[id(content_widget)] = content_widget
        self.section_frames[id(section_frame)] = section_frame
        toggle_button.clicked.connect(self.toggle_section)
        
        section_layout.addWidget(header)
        section_layout.addWidget(content_widget, 1)
        section_frame.setLayout(section_layout)
        
        return section_frame
    
    def create_section(self, title, content_widget):
        section_frame = QFrame()
        section_frame.setObjectName("infoSection")
        section_layout = QVBoxLayout()
        self.set_zero_spacing(section_layout)
        
        header = QFrame()
        header.setObjectName("infoSectionHeader")
        header_layout = QHBoxLayout()
        header_layout.setContentsMargins(8, 4, 8, 4)
        
        spacer = QWidget()
        spacer.setFixedWidth(24)
        # this height needs to be set to align with the toggle button height
        spacer.setFixedHeight(20)
        
        title_label = QLabel(title)
        title_label.setObjectName("infoSectionTitle")
        
        header_layout.addWidget(spacer)
        header_layout.addWidget(title_label)
        header_layout.addStretch()
        header.setLayout(header_layout)
        
        section_layout.addWidget(header)
        section_layout.addWidget(content_widget, 1)
        section_frame.setLayout(section_layout)
        
        return section_frame
    
    def create_goals_content(self):
        scroll_area = QScrollArea()
        scroll_area.setWidgetResizable(True)
        scroll_area.setObjectName("infoContent")
        
        content = QLabel()
        content.setObjectName("goalsContent")
        content.setText("will be implemented later") #TODO fetch from case data

        scroll_area.setWidget(content)
        return scroll_area
    
    def create_pinned_content(self):
        scroll_area = QScrollArea()
        scroll_area.setWidgetResizable(True)
        scroll_area.setObjectName("infoContent")
        
        content = QLabel("No pinned notes") #TODO fetch from database, notes with flag
        content.setObjectName("pinnedContent")
        content.setAlignment(Qt.AlignmentFlag.AlignTop | Qt.AlignmentFlag.AlignLeft)
        content.setWordWrap(True)
        content.setContentsMargins(8, 8, 8, 8)
        
        scroll_area.setWidget(content)
        return scroll_area
    
    def create_iocs_content(self):
        scroll_area = QScrollArea()
        scroll_area.setWidgetResizable(True)
        scroll_area.setObjectName("infoContent")

        content = QLabel("No IoCs extracted yet") #TODO fetch from database
        content.setObjectName("iocsContent")
        content.setAlignment(Qt.AlignmentFlag.AlignTop | Qt.AlignmentFlag.AlignLeft)
        content.setWordWrap(True)
        content.setContentsMargins(8, 8, 8, 8)
        
        scroll_area.setWidget(content)
        return scroll_area
    
    def toggle_section(self):
        # if the function is called, sender() returns the button which was clicked
        button = self.sender()
        # this is necessary to satisfy type checker
        if not isinstance(button, QPushButton):
            return
        
        widget_id = button.property("content_widget")
        section_frame_id = button.property("section_frame")
        # these were stored previously in create_collapsible_section
        content_widget = self.toggle_widgets.get(widget_id)
        section_frame = self.section_frames.get(section_frame_id)
        
        # this check is necessary to satisfy type checker
        if content_widget and section_frame:
            is_visible = content_widget.isVisible()
            # invert visibility
            new_visibility = not is_visible
            # set inverted visibility
            content_widget.setVisible(new_visibility)
            button.setText("▼" if new_visibility else "▶")
            button.setProperty("collapsed", not new_visibility)
            
            # this is proposed by AI to resize the section properly - seems functional
            if new_visibility:  # Expanding
                section_frame.setMinimumHeight(100)
                section_frame.setMaximumHeight(16777215)  # QWIDGETSIZE_MAX
                section_frame.setSizePolicy(QSizePolicy.Policy.Preferred, QSizePolicy.Policy.Expanding)
            else:  # Collapsing
                header_height = 32
                section_frame.setMinimumHeight(header_height)
                section_frame.setMaximumHeight(header_height)
                section_frame.setSizePolicy(QSizePolicy.Policy.Preferred, QSizePolicy.Policy.Fixed)

    def connect_ui_elements(self):
        pass
    
    def create_menu_bar(self):
        # the menu is actually a toolbar with nested menus 
        menu_toolbar = QToolBar("Menu Bar", self)
        menu_toolbar.setMovable(False)
        menu_toolbar.setFloatable(False)
        # by default there is an unnecessary context menu, we have to explicitly prevent it
        menu_toolbar.setContextMenuPolicy(Qt.ContextMenuPolicy.PreventContextMenu)
        menu_toolbar.setObjectName("menuToolbar")
        
        brand_label = QLabel(" FACTUM ")
        brand_label.setObjectName("appBrand")
        menu_toolbar.addWidget(brand_label)
        
        file_menu = QMenu("File", self)
        file_button = QPushButton("File")
        file_button.setObjectName("menuButton")
        file_button.setMenu(file_menu)
        menu_toolbar.addWidget(file_button)
        
        view_menu = QMenu("View", self)
        view_button = QPushButton("View")
        view_button.setObjectName("menuButton")
        view_button.setMenu(view_menu)
        menu_toolbar.addWidget(view_button)
        
        settings_menu = QMenu("Settings", self)
        settings_button = QPushButton("Settings")
        settings_button.setObjectName("menuButton")
        settings_button.setMenu(settings_menu)
        menu_toolbar.addWidget(settings_button)
        
        help_menu = QMenu("Help", self)
        help_button = QPushButton("Help")
        help_button.setObjectName("menuButton")
        help_button.setMenu(help_menu)
        menu_toolbar.addWidget(help_button)
        
        self.addToolBar(Qt.ToolBarArea.TopToolBarArea, menu_toolbar)
        
    def create_bottom_toolbar(self):
        bottom_toolbar = QToolBar("Bottom Toolbar", self)
        bottom_toolbar.setMovable(False)
        bottom_toolbar.setFloatable(False)
        # by default there is an unnecessary context menu, we have to explicitly prevent it
        bottom_toolbar.setContextMenuPolicy(Qt.ContextMenuPolicy.PreventContextMenu)
        
        viewing_label = QLabel("  Viewing: ")
        viewing_label.setProperty("statusLabel", "true")
        viewing_value = QLabel("No case selected")
        viewing_value.setProperty("statusValue", "true")
        
        bottom_toolbar.addWidget(viewing_label)
        bottom_toolbar.addWidget(viewing_value)
        bottom_toolbar.addSeparator()
        
        cli_label = QLabel("CLI Active: ")
        cli_label.setProperty("statusLabel", "true")
        cli_value = QLabel("None")
        cli_value.setProperty("statusValue", "true")
        
        bottom_toolbar.addWidget(cli_label)
        bottom_toolbar.addWidget(cli_value)
        
        # this adds a spacer to push the gpg status to the right
        spacer = QWidget()
        spacer.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Preferred)
        bottom_toolbar.addWidget(spacer)
        
        # TODO: real GPG status from settings
        gpg_label = QLabel("GPG Key: ")
        gpg_label.setProperty("statusLabel", "true")
        gpg_value = QLabel("Not configured")
        gpg_value.setProperty("statusValue", "true")
        
        bottom_toolbar.addWidget(gpg_label)
        bottom_toolbar.addWidget(gpg_value)
        
        bottom_toolbar.setContentsMargins(4, 0, 4, 0)
        
        self.addToolBar(Qt.ToolBarArea.BottomToolBarArea, bottom_toolbar)
        
    def set_zero_spacing(self, layout):
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)