"""Correlation/association sheet for selected columns."""

from __future__ import annotations

import contextlib
import math
import weakref
from collections.abc import Mapping, Sequence
from concurrent.futures import Future
from threading import Timer
from time import monotonic_ns
from typing import TYPE_CHECKING, Any, Literal

import polars as pl
from prompt_toolkit.eventloop import call_soon_threadsafe

from pulka.core.engine.contracts import EnginePayloadHandle, TableSlice
from pulka.core.engine.polars_adapter import (
    collect_lazyframe,
    unwrap_lazyframe_handle,
)
from pulka.core.jobs import JobRequest, JobResult, JobRunner
from pulka.core.plan import normalized_columns_key
from pulka.core.sheet import Sheet
from pulka.core.sheet_actions import SheetEnterAction
from pulka.core.viewer import Viewer
from pulka.sheets.data_sheet import DataSheet

if TYPE_CHECKING:  # pragma: no cover - typing only
    from pulka.command.registry import CommandContext
    from pulka.core.viewer.view_stack import ViewStack
    from pulka.tui.screen import Screen


_COR_DISPLAY_SCHEMA: dict[str, pl.DataType] = {
    "left": pl.Utf8,
    "right": pl.Utf8,
    "measure": pl.Utf8,
    "value": pl.Float64,
    "n": pl.Int64,
    "notes": pl.Utf8,
}

_COR_COLUMNS: tuple[str, ...] = tuple(_COR_DISPLAY_SCHEMA.keys())

_DEFAULT_MAX_COLUMNS = 24
_DEFAULT_SAMPLE_ROWS = 200_000
_DEFAULT_MAX_CATEGORIES = 200
_DEFAULT_MAX_CELLS = 20_000


def _empty_cor_frame() -> pl.DataFrame:
    return pl.DataFrame({name: [] for name in _COR_DISPLAY_SCHEMA}, schema=_COR_DISPLAY_SCHEMA)


def _is_numeric_dtype(dtype: pl.DataType | None) -> bool:
    if dtype is None:
        return False
    return bool(
        dtype.is_numeric()
        or dtype
        in {
            pl.Date,
            pl.Datetime,
            pl.Time,
            pl.Duration,
        }
    )


def _is_boolean_dtype(dtype: pl.DataType | None) -> bool:
    return dtype is not None and dtype == pl.Boolean


def _is_categorical_dtype(dtype: pl.DataType | None) -> bool:
    if dtype is None:
        return False
    with contextlib.suppress(Exception):
        if dtype in {pl.Utf8, pl.Categorical, pl.Enum}:
            return True
    return dtype == pl.Utf8


def _dtype_category(
    dtype: pl.DataType | None,
) -> Literal["numeric", "boolean", "categorical", "other"]:
    if _is_boolean_dtype(dtype):
        return "boolean"
    if _is_numeric_dtype(dtype):
        return "numeric"
    if _is_categorical_dtype(dtype):
        return "categorical"
    return "other"


def _pair_frame(
    df: pl.DataFrame,
    left: str,
    right: str,
    *,
    cast_left: pl.DataType | None = None,
    cast_right: pl.DataType | None = None,
) -> pl.DataFrame:
    cols: list[pl.Expr] = []
    left_expr: pl.Expr = pl.col(left)
    right_expr: pl.Expr = pl.col(right)
    if cast_left is not None:
        left_expr = left_expr.cast(cast_left, strict=False)
    if cast_right is not None:
        right_expr = right_expr.cast(cast_right, strict=False)
    cols.append(left_expr.alias(left))
    cols.append(right_expr.alias(right))
    return df.select(cols).drop_nulls()


def _pearson(df: pl.DataFrame, left: str, right: str) -> tuple[float | None, int, str | None]:
    frame = _pair_frame(df, left, right, cast_left=pl.Float64, cast_right=pl.Float64)
    n = int(frame.height)
    if n < 2:
        return None, n, "insufficient rows"
    try:
        value = frame.select(pl.corr(left, right).alias("value")).item(0, 0)
    except Exception as exc:
        return None, n, f"corr error: {exc}"[:80]
    if value is None:
        return None, n, "unavailable"
    try:
        value = float(value)
    except Exception:
        return None, n, "unavailable"
    if math.isnan(value):
        return None, n, "constant column"
    return value, n, None


def _correlation_ratio_eta2(
    df: pl.DataFrame,
    categorical: str,
    numeric: str,
    *,
    max_categories: int,
) -> tuple[float | None, int, str | None]:
    frame = _pair_frame(df, categorical, numeric, cast_right=pl.Float64)
    n = int(frame.height)
    if n < 2:
        return None, n, "insufficient rows"

    cat_series = frame.get_column(categorical)
    try:
        unique_count = int(cat_series.n_unique())
    except Exception:
        unique_count = 0
    if unique_count <= 1:
        return None, n, "constant category"
    if unique_count > max_categories:
        return None, n, f"too many categories ({unique_count})"

    overall_mean = frame.get_column(numeric).mean()
    if overall_mean is None:
        return None, n, "unavailable"

    grouped = frame.group_by(categorical).agg(
        pl.len().alias("n_i"),
        pl.col(numeric).mean().alias("mean_i"),
    )
    between_ss = grouped.select(
        (pl.col("n_i") * (pl.col("mean_i") - float(overall_mean)) ** 2).sum().alias("between")
    ).item(0, 0)
    total_ss = frame.select(
        ((pl.col(numeric) - float(overall_mean)) ** 2).sum().alias("total")
    ).item(0, 0)
    if between_ss is None or total_ss is None:
        return None, n, "unavailable"
    try:
        between = float(between_ss)
        total = float(total_ss)
    except Exception:
        return None, n, "unavailable"
    if total <= 0:
        return None, n, "constant numeric"
    eta2 = between / total
    return float(max(0.0, min(1.0, eta2))), n, None


def _cramers_v(
    df: pl.DataFrame,
    left: str,
    right: str,
    *,
    max_categories: int,
    max_cells: int,
) -> tuple[float | None, int, str | None]:
    frame = _pair_frame(df, left, right, cast_left=pl.Utf8, cast_right=pl.Utf8)
    n = int(frame.height)
    if n < 2:
        return None, n, "insufficient rows"

    left_levels = frame.get_column(left).unique()
    right_levels = frame.get_column(right).unique()
    r = int(left_levels.len())
    c = int(right_levels.len())
    if r <= 1 or c <= 1:
        return None, n, "constant category"
    if r > max_categories or c > max_categories:
        return None, n, f"too many categories ({r}x{c})"
    if r * c > max_cells:
        return None, n, f"too many cells ({r * c})"

    counts = frame.group_by([left, right]).len().rename({"len": "count"})
    row_sums = frame.group_by(left).len().rename({"len": "row_total"})
    col_sums = frame.group_by(right).len().rename({"len": "col_total"})

    all_cells = left_levels.to_frame(name=left).join(right_levels.to_frame(name=right), how="cross")
    joined = (
        all_cells.join(counts, on=[left, right], how="left")
        .join(row_sums, on=left, how="left")
        .join(col_sums, on=right, how="left")
        .with_columns(pl.col("count").fill_null(0))
    )
    joined = joined.with_columns(
        (pl.col("row_total") * pl.col("col_total") / float(n)).alias("expected")
    )
    chi2 = joined.select(
        (((pl.col("count") - pl.col("expected")) ** 2) / pl.col("expected")).sum().alias("chi2")
    ).item(0, 0)
    if chi2 is None:
        return None, n, "unavailable"
    try:
        chi2_value = float(chi2)
    except Exception:
        return None, n, "unavailable"
    if chi2_value < 0:
        chi2_value = 0.0
    denom = float(n) * float(min(r - 1, c - 1))
    if denom <= 0:
        return None, n, "unavailable"
    v = math.sqrt(chi2_value / denom)
    return float(max(0.0, min(1.0, v))), n, None


def compute_association_table(
    df: pl.DataFrame,
    columns: Sequence[str],
    schema: Mapping[str, pl.DataType],
    *,
    sample_rows: int | None,
    max_categories: int = _DEFAULT_MAX_CATEGORIES,
    max_cells: int = _DEFAULT_MAX_CELLS,
) -> pl.DataFrame:
    """Compute pairwise association measures for ``columns`` within ``df``."""

    records: list[dict[str, Any]] = []
    sampled_note = None
    if sample_rows is not None and df.height >= sample_rows:
        sampled_note = f"first {sample_rows} rows"

    for idx, left in enumerate(columns):
        for right in columns[idx + 1 :]:
            left_dtype = schema.get(left, df.schema.get(left))
            right_dtype = schema.get(right, df.schema.get(right))
            left_kind = _dtype_category(left_dtype)
            right_kind = _dtype_category(right_dtype)

            measure: str | None = None
            value: float | None = None
            n: int = 0
            notes: str | None = None

            if left_kind == "numeric" and right_kind == "numeric":
                measure = "pearson"
                value, n, notes = _pearson(df, left, right)
            elif left_kind == "boolean" and right_kind == "boolean":
                measure = "phi"
                value, n, notes = _pearson(df, left, right)
            elif left_kind == "categorical" and right_kind == "categorical":
                measure = "cramers_v"
                value, n, notes = _cramers_v(
                    df, left, right, max_categories=max_categories, max_cells=max_cells
                )
            elif left_kind == "categorical" and right_kind in {"numeric", "boolean"}:
                measure = "eta2"
                value, n, notes = _correlation_ratio_eta2(
                    df,
                    left,
                    right,
                    max_categories=max_categories,
                )
            elif right_kind == "categorical" and left_kind in {"numeric", "boolean"}:
                measure = "eta2"
                value, n, notes = _correlation_ratio_eta2(
                    df,
                    right,
                    left,
                    max_categories=max_categories,
                )
            elif "categorical" in (left_kind, right_kind):
                measure = "cramers_v"
                value, n, notes = _cramers_v(
                    df, left, right, max_categories=max_categories, max_cells=max_cells
                )
            else:
                measure = "unavailable"
                notes = "unsupported dtype"

            if sampled_note and notes:
                notes = f"{notes}; {sampled_note}"
            elif sampled_note and not notes:
                notes = sampled_note

            records.append(
                {
                    "left": left,
                    "right": right,
                    "measure": measure,
                    "value": value,
                    "n": n,
                    "notes": notes or "",
                }
            )

    if not records:
        return _empty_cor_frame()

    result = pl.DataFrame(records, schema=_COR_DISPLAY_SCHEMA)
    with contextlib.suppress(Exception):
        result = result.sort(
            by=[pl.col("value").abs().fill_null(-1.0), pl.col("left"), pl.col("right")],
            descending=[True, False, False],
        )
    return result


def compute_association_df_for_sheet(
    base_sheet: Sheet,
    selected_columns: Sequence[str],
    *,
    sample_rows: int,
    max_categories: int,
    max_cells: int,
) -> pl.DataFrame:
    schema = getattr(base_sheet, "schema", {}) or {}
    lf_handle = getattr(base_sheet, "lf", None)
    lf = None
    if isinstance(lf_handle, EnginePayloadHandle):
        with contextlib.suppress(Exception):
            lf = unwrap_lazyframe_handle(lf_handle)
    if lf is None and hasattr(base_sheet, "to_lazy"):
        with contextlib.suppress(Exception):
            candidate = base_sheet.to_lazy()
            if isinstance(candidate, EnginePayloadHandle):
                lf = unwrap_lazyframe_handle(candidate)
    if lf is None:
        raise ValueError("Correlation requires a Polars-backed sheet")

    limited = lf.select(list(selected_columns)).limit(int(sample_rows))
    data = collect_lazyframe(limited)
    if not data.columns:
        return _empty_cor_frame()
    coerced_schema: dict[str, pl.DataType] = {
        name: schema.get(name, data.schema.get(name, pl.Utf8)) for name in selected_columns
    }
    return compute_association_table(
        data,
        list(selected_columns),
        coerced_schema,
        sample_rows=int(sample_rows),
        max_categories=int(max_categories),
        max_cells=int(max_cells),
    )


def _select_columns_from_summary_view(viewer: Viewer, sheet: Sheet) -> list[str]:
    selected_ids = set(getattr(viewer, "_selected_row_ids", set()))
    selected: list[str] = []
    if selected_ids:
        for row_id in selected_ids:
            if isinstance(row_id, str):
                selected.append(row_id)
                continue
            with contextlib.suppress(Exception):
                name = sheet.get_value_at(int(row_id), "column")
                if isinstance(name, str):
                    selected.append(name)
    return list(dict.fromkeys(selected))


def _coerce_selected_columns(
    context: CommandContext,
    *,
    require_summary_selection: bool,
    max_columns: int,
) -> list[str] | None:
    viewer = context.viewer
    sheet = viewer.sheet

    args = getattr(context, "args", None)
    _ = args  # placeholder to placate linters for older contexts

    if require_summary_selection:
        selected = _select_columns_from_summary_view(viewer, sheet)
    else:
        selected = _select_columns_from_summary_view(viewer, sheet)

    if not selected:
        return None
    if len(selected) > max_columns:
        return selected[: max_columns + 1]
    return selected


class CorrelationSheet(DataSheet):
    """Sheet implementation for correlation / association tables."""

    is_correlation_view = True
    hide_transforms_panel_by_default = True
    hide_insight_panel_by_default = True

    def __init__(
        self,
        base_sheet: Sheet,
        selected_columns: Sequence[str],
        *,
        runner: JobRunner,
        sample_rows: int = _DEFAULT_SAMPLE_ROWS,
        max_categories: int = _DEFAULT_MAX_CATEGORIES,
        max_cells: int = _DEFAULT_MAX_CELLS,
        association_df: pl.DataFrame | None = None,
    ) -> None:
        if runner is None:  # pragma: no cover - defensive guard
            msg = "CorrelationSheet requires a JobRunner instance"
            raise ValueError(msg)
        frame = association_df if association_df is not None else _empty_cor_frame()
        super().__init__(
            frame.lazy(),
            columns=list(_COR_COLUMNS),
            runner=runner,
        )

        self.source_sheet = base_sheet
        self.selected_columns = tuple(selected_columns)
        self._sample_rows = int(sample_rows)
        self._max_categories = int(max_categories)
        self._max_cells = int(max_cells)
        self._pending_future: Future[JobResult] | None = None
        self.is_insight_soft_disabled = True

        self._job_sheet_id: str | None = getattr(base_sheet, "sheet_id", None)
        plan_hash = ""
        job_context = getattr(base_sheet, "job_context", None)
        if callable(job_context):
            with contextlib.suppress(Exception):
                sheet_id, _generation, plan_hash = job_context()
                self._job_sheet_id = sheet_id

        colsig = normalized_columns_key(self.selected_columns)
        self._job_tag = f"cor:cols={colsig}:plan={plan_hash}"
        self._preserve_jobs_from = self._job_sheet_id

        if association_df is None:
            self._ensure_job()

    @classmethod
    def from_dataframe(
        cls,
        base_sheet: Sheet,
        df: pl.DataFrame,
        *,
        selected_columns: Sequence[str],
        runner: JobRunner,
        sample_rows: int = _DEFAULT_SAMPLE_ROWS,
        max_categories: int = _DEFAULT_MAX_CATEGORIES,
        max_cells: int = _DEFAULT_MAX_CELLS,
    ) -> CorrelationSheet:
        return cls(
            base_sheet,
            selected_columns,
            runner=runner,
            sample_rows=sample_rows,
            max_categories=max_categories,
            max_cells=max_cells,
            association_df=df,
        )

    def fetch_slice(self, row_start: int, row_count: int, columns: Sequence[str]) -> TableSlice:
        table_slice = super().fetch_slice(row_start, row_count, columns)
        if {"left", "right"} <= set(table_slice.column_names):
            with contextlib.suppress(Exception):
                left_values = table_slice.column("left").values
                right_values = table_slice.column("right").values
                table_slice.row_ids = tuple(zip(left_values, right_values, strict=False))
        return table_slice

    def enter_action(self, viewer: Viewer) -> SheetEnterAction | None:
        selected_ids = set(getattr(viewer, "_selected_row_ids", set()))
        columns: list[str] = []
        if selected_ids:
            for row_id in selected_ids:
                if isinstance(row_id, tuple) and len(row_id) >= 2:
                    left, right = row_id[0], row_id[1]
                    if isinstance(left, str):
                        columns.append(left)
                    if isinstance(right, str):
                        columns.append(right)
                    continue
                with contextlib.suppress(Exception):
                    idx = int(row_id)
                    left = self.get_value_at(idx, "left")
                    right = self.get_value_at(idx, "right")
                    if isinstance(left, str):
                        columns.append(left)
                    if isinstance(right, str):
                        columns.append(right)

        if not columns:
            with contextlib.suppress(Exception):
                left = self.get_value_at(viewer.cur_row, "left")
                right = self.get_value_at(viewer.cur_row, "right")
                if isinstance(left, str):
                    columns.append(left)
                if isinstance(right, str):
                    columns.append(right)

        deduped = tuple(dict.fromkeys(name for name in columns if isinstance(name, str)))
        if not deduped:
            return None
        return SheetEnterAction(kind="apply-selection", columns=deduped, pop_viewer=False)

    # Job plumbing -----------------------------------------------------

    def _ensure_job(self) -> None:
        sheet_id = self._job_sheet_id
        if sheet_id is None:
            return
        gen = self.job_runner.current_generation(sheet_id)
        if self.job_runner.get(sheet_id, self._job_tag) is not None:
            return
        if self._pending_future is not None:
            return

        selected = list(self.selected_columns)
        base_sheet = self.source_sheet
        runner = self.job_runner
        tag = self._job_tag
        sample_rows = self._sample_rows
        max_categories = self._max_categories
        max_cells = self._max_cells

        def _fn(job_gen: int) -> pl.DataFrame | None:
            if runner.current_generation(sheet_id) != job_gen:
                return None
            return compute_association_df_for_sheet(
                base_sheet,
                selected,
                sample_rows=sample_rows,
                max_categories=max_categories,
                max_cells=max_cells,
            )

        future = runner.enqueue(JobRequest(sheet_id=sheet_id, generation=gen, tag=tag, fn=_fn))
        self._pending_future = future


class _CorrelationUiHandle:
    """Polls the job runner for correlation results and refreshes the viewer."""

    _tick_interval_ns = int(0.15 * 1e9)

    def __init__(self, sheet: CorrelationSheet, screen: Screen | None) -> None:
        self._sheet_ref: weakref.ReferenceType[CorrelationSheet] | None = weakref.ref(sheet)
        self._screen_ref: weakref.ReferenceType[Screen] | None = (
            weakref.ref(screen) if screen is not None else None
        )
        self._last_result_ts: int = getattr(sheet, "_result_ts", 0)
        self._finished = False
        self._timer: Timer | None = None
        self._last_tick_ns: int = monotonic_ns()
        if self._screen_ref is not None:
            self._schedule_tick()

    def _cancel_timer(self) -> None:
        timer = self._timer
        if timer is not None:
            timer.cancel()
            self._timer = None

    def _schedule_tick(self) -> None:
        if self._tick_interval_ns <= 0:
            return
        delay_s = self._tick_interval_ns / 1_000_000_000
        if delay_s <= 0:
            return
        self._cancel_timer()
        timer = Timer(delay_s, self._invalidate_screen)
        timer.daemon = True
        timer.start()
        self._timer = timer

    def _invalidate_screen(self) -> None:
        screen_ref = self._screen_ref
        if screen_ref is None:
            return
        screen = screen_ref()
        if screen is None:
            return
        app = getattr(screen, "app", None)
        if app is None:
            return
        try:
            call_soon_threadsafe(app.invalidate)
        except Exception:
            with contextlib.suppress(Exception):
                app.invalidate()

    def notify_ready(self, _future: Future[JobResult] | None = None) -> None:
        self._invalidate_screen()

    def consume_update(self, viewer: Viewer) -> bool:
        sheet = self._sheet_ref() if self._sheet_ref is not None else None
        if sheet is None or viewer.sheet is not sheet:
            self._finished = True
            self._cancel_timer()
            return True

        sheet_id = getattr(sheet, "_job_sheet_id", None)
        tag = getattr(sheet, "_job_tag", None)
        if sheet_id is None or tag is None:
            self._finished = True
            self._cancel_timer()
            return True

        result = sheet.job_runner.get(sheet_id, tag)
        pending = sheet._pending_future
        if result is None or result.ts_ns <= self._last_result_ts:
            if pending is not None and pending.done():
                with contextlib.suppress(Exception):
                    completed = pending.result()
                    if isinstance(completed, JobResult):
                        result = completed
            if result is None or result.ts_ns <= self._last_result_ts:
                if pending is None:
                    self._finished = True
                    self._cancel_timer()
                    viewer.status_message = None
                    return True
                now = monotonic_ns()
                if now - self._last_tick_ns >= self._tick_interval_ns:
                    self._last_tick_ns = now
                    self._schedule_tick()
                viewer.status_message = "computing correlations…"
                return False

        self._last_result_ts = result.ts_ns

        if result.error is not None:
            viewer.status_message = f"cor error: {result.error!s}"[:120]
            self._finished = True
            self._cancel_timer()
            return True

        value = result.value
        if not isinstance(value, pl.DataFrame):
            viewer.status_message = "correlation unavailable"
            self._finished = True
            self._cancel_timer()
            return True

        source_sheet = getattr(sheet, "source_sheet", sheet)
        replacement = CorrelationSheet.from_dataframe(
            source_sheet,
            value,
            selected_columns=getattr(sheet, "selected_columns", ()),
            runner=sheet.job_runner,
            sample_rows=getattr(sheet, "_sample_rows", _DEFAULT_SAMPLE_ROWS),
            max_categories=getattr(sheet, "_max_categories", _DEFAULT_MAX_CATEGORIES),
            max_cells=getattr(sheet, "_max_cells", _DEFAULT_MAX_CELLS),
        )
        viewer.replace_sheet(replacement, source_path=None)
        with contextlib.suppress(Exception):
            total_rows = int(len(replacement))
            viewer._total_rows = total_rows
            viewer._row_count_stale = False
            viewer._row_count_future = None
            viewer._row_count_display_pending = False
            viewer.mark_status_dirty()
        viewer.status_message = None
        self._finished = True
        self._cancel_timer()
        return True

    def cancel(self) -> None:
        self._finished = True
        self._cancel_timer()
        self._sheet_ref = None
        self._screen_ref = None


def register_correlation_ui_handle(
    cor_sheet: CorrelationSheet,
    viewer: Viewer,
    screen: Screen | None,
) -> None:
    if screen is None:
        return
    handle = _CorrelationUiHandle(cor_sheet, screen)
    screen.register_job(viewer, handle)
    pending = getattr(cor_sheet, "_pending_future", None)
    if pending is not None:

        def _deliver_result() -> None:
            handle.consume_update(viewer)
            with contextlib.suppress(Exception):
                viewer.ui_hooks.invalidate()

        def _on_done(_future: Future[JobResult]) -> None:
            handle.notify_ready()
            try:
                viewer.ui_hooks.call_soon(_deliver_result)
            except Exception:
                _deliver_result()

        with contextlib.suppress(Exception):
            pending.add_done_callback(_on_done)
    else:
        handle.notify_ready()


def _activate_viewer(
    base_viewer: Viewer,
    derived_viewer: Viewer,
    *,
    view_stack: ViewStack | None,
) -> Viewer:
    if view_stack is not None:
        view_stack.push(derived_viewer)
        return view_stack.active or derived_viewer
    base_viewer.replace_sheet(derived_viewer.sheet, source_path=None)
    base_viewer.status_message = getattr(derived_viewer, "status_message", None)
    return base_viewer


def open_correlation_viewer(
    base_viewer: Viewer,
    selected_columns: Sequence[str],
    *,
    session: Any | None = None,
    view_stack: ViewStack | None = None,
    screen: Screen | None = None,
    runner: JobRunner,
    association_df: pl.DataFrame | None = None,
) -> Viewer:
    base_sheet = getattr(base_viewer.sheet, "source_sheet", base_viewer.sheet)
    cor_sheet = CorrelationSheet(
        base_sheet,
        selected_columns,
        runner=runner,
        association_df=association_df,
    )
    derived_viewer = Viewer(
        cor_sheet,
        viewport_rows=base_viewer._viewport_rows_override,
        viewport_cols=base_viewer._viewport_cols_override,
        source_path=None,
        session=session or base_viewer.session,
        runner=runner,
    )
    derived_viewer.status_message = "computing correlations…"
    active = _activate_viewer(base_viewer, derived_viewer, view_stack=view_stack)
    register_correlation_ui_handle(cor_sheet, active, screen)
    return active


def _cor_cmd(context: CommandContext, args: list[str]) -> None:
    viewer = context.viewer
    sheet = viewer.sheet
    base_sheet = getattr(sheet, "source_sheet", sheet)
    runner = getattr(base_sheet, "job_runner", None) or getattr(viewer, "job_runner", None)
    if not isinstance(runner, JobRunner):
        runner = viewer.job_runner  # type: ignore[assignment]

    fallback_used = False
    max_columns = _DEFAULT_MAX_COLUMNS
    session = getattr(context, "session", None)
    summary_config = getattr(getattr(session, "config", None), "summary", None)
    configured_max = getattr(summary_config, "correlation_max_columns", None)
    if isinstance(configured_max, int) and configured_max >= 2:
        max_columns = configured_max
    if args:
        selected = [a for a in args if isinstance(a, str) and a.strip()]
    else:
        selected = _select_columns_from_summary_view(viewer, sheet)

    if len(selected) < 2:
        selected = list(getattr(base_sheet, "columns", []))
        fallback_used = True

    selected_lookup = [name for name in selected if name in getattr(base_sheet, "columns", [])]
    selected_lookup = list(dict.fromkeys(selected_lookup))
    if len(selected_lookup) < 2:
        viewer.status_message = "no matching columns selected"
        return
    if len(selected_lookup) > max_columns:
        selected_lookup = selected_lookup[:max_columns]
        fallback_used = True

    ui = getattr(context, "ui", None) or getattr(context, "screen", None)
    screen = ui if ui is not None else None
    session = session or viewer.session
    view_stack = getattr(context, "view_stack", None)
    if view_stack is None and session is not None:
        view_stack = getattr(session, "view_stack", None)

    association_df: pl.DataFrame | None = None
    if screen is None:
        try:
            association_df = compute_association_df_for_sheet(
                base_sheet,
                selected_lookup,
                sample_rows=_DEFAULT_SAMPLE_ROWS,
                max_categories=_DEFAULT_MAX_CATEGORIES,
                max_cells=_DEFAULT_MAX_CELLS,
            )
        except Exception as exc:
            viewer.status_message = f"cor error: {exc}"[:120]
            return

    try:
        active = open_correlation_viewer(
            viewer,
            selected_lookup,
            session=session,
            view_stack=view_stack,
            screen=screen,
            runner=runner,
            association_df=association_df,
        )
        context.viewer = active
        context.sheet = active.sheet
        if screen is None:
            if fallback_used:
                active.status_message = f"correlations: first {len(selected_lookup)} columns"
            else:
                active.status_message = f"correlations: {len(selected_lookup)} columns"
    except Exception as exc:  # pragma: no cover - guardrail
        viewer.status_message = f"cor error: {exc}"[:120]
    if screen is not None:
        screen.refresh(skip_metrics=True)


__all__ = [
    "CorrelationSheet",
    "_cor_cmd",
    "compute_association_table",
    "open_correlation_viewer",
    "register_correlation_ui_handle",
]
