import pandas as pd
import numpy as np
import math

from pandaview.css_style import _get_base_styles, quick_css

try:  # Optional: available in notebooks
    from IPython.display import display, HTML  # type: ignore

    _HAS_DISPLAY = True
except Exception:  # pragma: no cover - environment without IPython
    display = None  # type: ignore[assignment]
    HTML = None  # type: ignore[assignment]
    _HAS_DISPLAY = False

try:  # Optional: richer widget-based UI
    from ipywidgets import widgets  # type: ignore

    _HAS_WIDGETS = True
except Exception:  # pragma: no cover - environment without ipywidgets
    widgets = None  # type: ignore[assignment]
    _HAS_WIDGETS = False

def _get_preview_df(df, rows):
    """Creates a subset of the dataframe for preview (first few rows + last row)."""
    row_count = df.shape[0]
    
    if row_count > 2 and row_count > rows:
        # Slice logic from original code: from 0 to -2 with step, then add last row
        step = math.ceil(row_count / (rows - 1))
        df_subset = pd.concat([
            df.iloc[0:-2:step],
            df.iloc[[-1]]
        ])
    else:
        df_subset = df.copy()
        
    return df_subset.reset_index()

def _analyze_column(df, col_name, is_index=False):
    """
    Analyzes a column to determine its type, uniqueness, and count.
    Returns a dictionary with metadata.
    """
    if is_index:
        series = df.index
    else:
        series = df[col_name]
        
    try:
        col_type = pd.api.types.infer_dtype(series)
    except Exception:
        col_type = 'error'

    return {
        'type': col_type,
        'series': series
    }

def _generate_view_html(df_subset, df_full):
    """
    Helper function to generate the HTML for a specific view (expanded or collapsed).
    """
    row_count, col_count = df_full.shape
    styles = _get_base_styles()
    col_types_map = {}
    
    # Replicating original iteration logic cleanly
    # df_subset has [index, col1, col2, ...]
    
    for i, col_name in enumerate(df_subset.columns):
        css_col_idx = i
        
        col_selector = f".col_heading.level0.col{css_col_idx}"
        is_index = (i == 0)
        
        # Analyze
        full_series = df_full.index if is_index else df_full[col_name]
        
        try:
            col_type = pd.api.types.infer_dtype(full_series)
            col_types_map[col_name] = col_type
        except Exception:
            styles.append(dict(selector=col_selector, props=[('color', 'red')]))
            continue

        # Style logic based on type/uniqueness/NaNs
        if col_type not in ['list', 'mixed']:
            unique_count = full_series.nunique()
        else:
            unique_count = 0
            
        if is_index:
            count = len(full_series)
        else:
            count = full_series.count()

        # Apply styles
        # 1. Uniqueness (Underline)
        if unique_count == row_count:
            styles.append(dict(selector=col_selector, props=[('text-decoration', 'underline')]))
            continue
            
        # 2. NaN values (Orange)
        if count != row_count:
            styles.append(dict(selector=col_selector, props=[('color', 'orange')]))
            continue

    is_expanded = len(df_subset) >= row_count

    caption_html = f'''
    <div class='btm_caption'>
        <b>{row_count} rows * {col_count} Columns :</b> {list(df_full.columns)}
        <br>
        Showing ({len(df_subset)} / {row_count} rows){' [expanded]' if is_expanded else ' [preview]'}
    </div>
    '''

    styler = df_subset.style.set_table_attributes('class="myTable"')\
        .set_table_styles(styles)\
        .set_caption(caption_html)\
        .hide(axis=0)

    html_render = styler.to_html()

    # Inject Column Types into HTML
    for col, col_type in col_types_map.items():
        html_render = html_render.replace(f'{col}</th>', f'{col}<p class="dtype">{col_type}</p></th>')
        
    return html_render


    
def see(df, rows=2):
    """
    Enhanced dataframe display function.
    Shows shape, columns, types, and a sample of data with visual cues for uniqueness and missing values.
    """
    try:
        # 1. Validation
        if not isinstance(df, pd.DataFrame):
            raise Exception("Input is not a DataFrame")
        if not df.columns.is_unique:
            raise Exception("Dataframe contains duplicated column names. Which is not managed")
        if len(df.columns) > 200:
            raise Exception(f"Number of columns is too damn high ! ({len(df.columns)})")

        # If display stack is not available, fall back to a plain text preview
        if not _HAS_DISPLAY:
            print(df.head(rows))
            return

        # If ipywidgets is not available, fall back to a simple static HTML preview
        if not _HAS_WIDGETS:
            df_small = _get_preview_df(df, rows)
            html_small = _generate_view_html(df_small, df)
            display(HTML(quick_css() + html_small))  # type: ignore[func-returns-value]
            return

        # 2. Prepare Data for Small View
        row_count, _ = df.shape
        current_rows = rows

        out = widgets.Output()

        def _render(current_rows_val):
            df_view = _get_preview_df(df, current_rows_val)
            html = _generate_view_html(df_view, df)
            display(HTML(quick_css() + html))  # type: ignore[func-returns-value]

        with out:
            _render(current_rows)

        def _on_expand_clicked(button):
            nonlocal current_rows
            if current_rows >= row_count:
                button.disabled = True
                return
            current_rows = min(current_rows * 10, row_count)
            out.clear_output()
            with out:
                _render(current_rows)
            if current_rows >= row_count:
                button.disabled = True
                button.description = "Expanded (max)"

        expand_btn = widgets.Button(
            description="Expand",
            button_style="",
            layout=widgets.Layout(width="auto")
        )
        expand_btn.on_click(_on_expand_clicked)

        # If initial preview already shows all rows, disable the button immediately
        if current_rows >= row_count:
            expand_btn.disabled = True
            expand_btn.description = "Expanded (max)"

        container = widgets.VBox([out, expand_btn])
        display(container)

    except Exception as e:
        print(f'Error on xm df.see:{e}')
        display(df)
