import panel as pn
import pandas as pd
import numpy as np
from hdf5 import HDF5Reader

class DataTable:
    """데이터 테이블 표시 담당 클래스"""
    
    def __init__(self, reader: HDF5Reader):
        self.reader = reader
        self.current_df = None
        self.visible_columns = []
        self.all_columns = []
        self.filtered_indices = None  # 현재 필터링된 인덱스 저장
        
        # Tag 업데이트 콜백
        self.on_tag_update_callback = None
        
        self._create_widgets()
        self._setup_callbacks()
    
    def _create_widgets(self):
        """위젯 생성"""
        # 컬럼 선택 체크박스 (사이드바용)
        self.column_selector = pn.widgets.CheckBoxGroup(
            name='Columns to Display',
            value=[],
            options=[],
            inline=False,
            sizing_mode='stretch_width'
        )
        
        # 테이블 위 컨트롤들 - Tag 전용
        self.tag_label = pn.pane.HTML(
            "<span style='font-weight:bold; font-size:20px;'>Tag:</span>",
            width=35
        )

        
        self.input_field = pn.widgets.TextInput(
            name='',
            placeholder='Enter tag value',
            width=150
        )
        
        self.accept_all_button = pn.widgets.Button(
            name='Accept All',
            button_type='primary',
            width=100
        )
        
        self.save_button = pn.widgets.Button(
            name='Save',
            button_type='warning',
            width=80
        )
        
        self.export_button = pn.widgets.Button(
            name='Export',
            button_type='success',
            width=80
        )
        
        # Particles 카운트 라벨
        self.particle_count_label = pn.pane.Markdown(
            "## Particles Table", margin=(5, 10, 0, 10)
        )
        
        # Tabulator with pagination
        self.data_table = pn.widgets.Tabulator(
            pd.DataFrame(),
            sizing_mode='stretch_both',
            theme='modern',
            show_index=True,
            pagination='local',
            page_size=1000,
            layout='fit_data_stretch',
            formatters={
                'Image': {'type': 'html'}
            },
            height=700,
            widths={
                'Image': 200
            },
            # Index 이름을 'Id'로 변경
            titles={
                'index': 'Id'
            },
            configuration={
                'indexTitle': 'Id'
            }
        )
    
    def _setup_callbacks(self):
        """콜백 설정"""
        self.column_selector.param.watch(self._on_column_change, 'value')
        self.data_table.param.watch(lambda e: self._update_table_width_by_page(), 'page')
        self.accept_all_button.on_click(self._on_accept_all)
        self.save_button.on_click(self._on_save)
        self.export_button.on_click(self._on_export)
    
    def _on_accept_all(self, event):
        """Accept All 버튼 클릭 - 현재 테이블의 Tag 컬럼에 입력값 적용"""
        if self.current_df is None:
            return
        
        input_value = self.input_field.value
        
        if 'Tag' not in self.current_df.columns:
            return
        
        # 입력값이 비어있으면 무시
        if not input_value or input_value.strip() == '':
            return
        
        # 현재 테이블에 표시된 인덱스들
        if self.filtered_indices is not None:
            target_indices = self.filtered_indices
        else:
            target_indices = self.current_df.index.tolist()
        
        # Tag 컬럼에 값 적용
        self.current_df.loc[target_indices, 'Tag'] = input_value
        
        # 테이블 새로고침
        self._refresh_table()
        
        print(f"[DataTable] Applied '{input_value}' to Tag column for {len(target_indices)} rows")
        
        # Tag 업데이트 콜백 호출 (그래프의 Tag 드롭다운 업데이트)
        if self.on_tag_update_callback:
            self.on_tag_update_callback()
    
    def _on_save(self, event):
        """Save 버튼 클릭 - HDF5 파일에 Tag 저장"""
        if self.current_df is None or 'Tag' not in self.current_df.columns:
            print("[DataTable] No data or Tag column to save")
            pn.state.notifications.warning("No Tag data to save", duration=3000)
            return
        
        tag_data = self.current_df['Tag'].tolist()
        success, message = self.reader.save_tags(tag_data)
        
        if success:
            print(f"[DataTable] {message}")
            pn.state.notifications.success("Tags saved to HDF5 file", duration=3000)
        else:
            print(f"[DataTable] Save failed: {message}")
            pn.state.notifications.error(f"Save failed: {message}", duration=5000)
    
    def _on_export(self, event):
        """Export 버튼 클릭 - Export 팝업 표시"""
        print("[DataTable] Export button clicked")
        
        # Export 콜백 호출 (main_ui에서 처리)
        if hasattr(self, 'on_export_callback') and self.on_export_callback:
            self.on_export_callback()
    
    def _refresh_table(self):
        """현재 상태로 테이블 새로고침"""
        if self.current_df is None:
            return
        
        # 현재 필터링된 인덱스가 있으면 그것을 사용, 없으면 전체 데이터
        if self.filtered_indices is not None:
            display_df = self.current_df.loc[self.filtered_indices]
        else:
            display_df = self.current_df
        
        if len(self.visible_columns) > 0:
            display_columns = ['Image'] if 'Image' in display_df.columns else []
            display_columns.extend(self.visible_columns)
            self.data_table.value = display_df[display_columns]
        else:
            if 'Image' in display_df.columns:
                self.data_table.value = display_df[['Image']]
            else:
                self.data_table.value = display_df
    
    def _update_count_label(self):
        """Particles 카운트 라벨 업데이트"""
        self.particle_count_label.object = "## Particles"
        # if self.current_df is None:
        #     self.particle_count_label.object = "## Particles"
        #     return
        
        # if self.filtered_indices is not None:
        #     count = len(self.filtered_indices)
        #     total = len(self.current_df)
        #     self.particle_count_label.object = f"## Particles ({count:,} / {total:,})"
        # else:
        #     total = len(self.current_df)
        #     self.particle_count_label.object = f"## Particles ({total:,})"
    
    def _on_column_change(self, event):
        """표시할 컬럼 변경 시 - 필터링 상태 유지"""
        self.visible_columns = event.new
        
        if self.current_df is None:
            return
        
        # 현재 필터링된 인덱스가 있으면 그것을 사용, 없으면 전체 데이터
        if self.filtered_indices is not None:
            display_df = self.current_df.loc[self.filtered_indices]
        else:
            display_df = self.current_df
        
        if len(self.visible_columns) > 0:
            # Image 컬럼은 항상 포함, 나머지는 선택된 것만
            display_columns = ['Image'] if 'Image' in display_df.columns else []
            display_columns.extend(self.visible_columns)
            self.data_table.value = display_df[display_columns]
        else:
            # visible_columns가 비어있으면 Image만 표시
            if 'Image' in display_df.columns:
                self.data_table.value = display_df[['Image']]
            else:
                self.data_table.value = pd.DataFrame()
    
    def _update_table_width_by_page(self):
        """현재 페이지의 이미지 너비에 맞춰 테이블 너비 조정"""
        if self.reader.particle_rects is None:
            return
        
        current_page = self.data_table.page
        page_size = self.data_table.page_size
        start_idx = current_page * page_size
        end_idx = start_idx + page_size
        
        page_widths = self.reader.particle_rects['width'][start_idx:end_idx]
        
        if len(page_widths) > 0:
            max_width = int(np.max(page_widths))
            image_col_width = max_width + 20
            
            self.data_table.widths = {
                **(self.data_table.widths or {}),
                'Image': image_col_width
            }
    
    def update_data(self, df: pd.DataFrame):
        """데이터 업데이트"""
        self.current_df = df
        self.filtered_indices = None  # 전체 데이터로 초기화
        self.data_table.value = df
        
        # 첫 페이지 너비 조정
        if self.reader.particle_rects is not None:
            page_size = self.data_table.page_size
            first_page_widths = self.reader.particle_rects['width'][:page_size]
            
            if len(first_page_widths) > 0:
                max_width = int(np.max(first_page_widths))
                image_col_width = max_width + 20
                
                self.data_table.widths = {
                    **(self.data_table.widths or {}),
                    'Image': image_col_width
                }
        
        # 컬럼 목록 업데이트 (Image만 제외, Tag는 포함)
        available_columns = [col for col in df.columns if col != 'Image']
        self.all_columns = available_columns
        
        # 체크박스 옵션 및 기본값 설정
        self.column_selector.options = available_columns
        self.column_selector.value = available_columns
        self.visible_columns = available_columns
        
        self._update_count_label()
    
    def filter_data(self, indices: list):
        """선택된 인덱스의 데이터만 표시"""
        print(f"[DataTable] filter_data called with {len(indices)} indices")
        
        if self.current_df is None:
            print(f"[DataTable] current_df is None, returning")
            return
        
        # 필터링된 인덱스 저장
        self.filtered_indices = indices
        
        # 빈 리스트면 빈 테이블 표시
        if len(indices) == 0:
            print(f"[DataTable] Empty indices, showing empty table")
            self.data_table.value = pd.DataFrame()
            self.data_table.page = 1
            self._update_count_label()
            return
        
        # 선택된 인덱스로 필터링
        filtered_df = self.current_df.loc[indices]
        print(f"[DataTable] Filtered df shape: {filtered_df.shape}")
        print(f"[DataTable] visible_columns: {self.visible_columns}")
        
        # 표시할 컬럼 결정
        if len(self.visible_columns) > 0:
            # Image 컬럼은 항상 포함, 나머지는 선택된 것만
            display_columns = ['Image'] if 'Image' in filtered_df.columns else []
            display_columns.extend(self.visible_columns)
            print(f"[DataTable] Displaying columns: {display_columns}")
            self.data_table.value = filtered_df[display_columns]
        else:
            # visible_columns가 비어있으면 Image만 표시
            print(f"[DataTable] visible_columns is empty, showing Image only")
            if 'Image' in filtered_df.columns:
                self.data_table.value = filtered_df[['Image']]
            else:
                self.data_table.value = filtered_df
        
        # 첫 페이지로 이동
        self.data_table.page = 1
        self._update_count_label()
        print(f"[DataTable] Table updated, new value shape: {self.data_table.value.shape}")
    
    def reset_filter(self):
        """필터 초기화 - 전체 데이터 표시"""
        if self.current_df is None:
            return
        
        # 필터링 인덱스 초기화
        self.filtered_indices = None
        
        # 전체 데이터로 복원
        if len(self.visible_columns) > 0:
            # Image 컬럼은 항상 포함, 나머지는 선택된 것만
            display_columns = ['Image'] if 'Image' in self.current_df.columns else []
            display_columns.extend(self.visible_columns)
            self.data_table.value = self.current_df[display_columns]
        else:
            if 'Image' in self.current_df.columns:
                self.data_table.value = self.current_df[['Image']]
            else:
                self.data_table.value = self.current_df
        
        # 첫 페이지로 이동
        self.data_table.page = 1
        self._update_count_label()
    
    def get_control_buttons(self):
        """Save, Export 버튼 반환 (외부에서 사용)"""
        return pn.Row(
            self.save_button,
            self.export_button,
            align='center',
            margin=0
        )
    
    def get_layout(self):
        """테이블 레이아웃 반환"""
        # Tag 입력 그룹 (가운데 배치)
        tag_group = pn.Row(
            self.tag_label,
            self.input_field,
            self.accept_all_button,
            align='center',
            margin=0
        )
        
        # 테이블 위 컨트롤 바
        controls_row = pn.Row(
            self.particle_count_label,
            pn.layout.HSpacer(),
            tag_group,
            pn.layout.HSpacer(),
            sizing_mode='stretch_width',
            align='center',
            margin=(5, 0, 5, 0)
        )
        
        return pn.Column(
            controls_row,
            self.data_table,
            sizing_mode='stretch_both',
            margin=0
        )