"""DataFrame-backed workbook for the pandas target. Shipped with generated code.

Each sheet is a ``pandas.DataFrame`` indexed by Excel row number with Excel
column letters as columns (object dtype, so a column can hold numbers, text,
blanks and Excel error values side by side, exactly as a sheet does).

``FrameGrid`` gives a frame the same cell-level API as ``xl.Grid``, so logic
blocks the vectoriser could not express run unchanged, cell by cell, over the
same frame. Vectorised blocks use the column helpers below, which keep
Excel's value semantics (blank is 0 in arithmetic, numeric text coerces,
error values propagate) while reading as column operations.
"""

from __future__ import annotations

import json
import math
from typing import Any, Callable, Iterator

import numpy as np
import pandas as pd

from . import xl
from .xl import Range, Unsupported, Val, XlError, col_index, col_letter, parse_a1, raw, split_qualified

# ---------------------------------------------------------------------------
# Scalar coercions reused element-wise
# ---------------------------------------------------------------------------


def _is_blank(v: Any) -> bool:
    """None, NaN, pd.NA and NaT all read as a blank cell."""
    if v is None:
        return True
    if isinstance(v, (str, XlError, Range)):
        return False
    return bool(pd.api.types.is_scalar(v) and pd.isna(v))


def _clean(v: Any) -> Any:
    """A frame cell as the runtime sees it.

    NaN becomes blank (None), and numpy scalars become Python values: a
    ``numpy.int64`` is not an ``int`` to the runtime, and a ``numpy.float64``
    divides by zero to ``inf`` instead of raising, which would turn Excel's
    ``#DIV/0!`` into ``#NUM!``.
    """
    v = raw(v)
    if isinstance(v, np.generic):
        v = v.item()
    return None if _is_blank(v) else v


def _num_scalar(v: Any) -> Any:
    v = _clean(v)
    if v is None:
        return 0
    if isinstance(v, XlError):
        return v
    return xl.to_number(v)


def _text_scalar(v: Any) -> Any:
    v = _clean(v)
    if isinstance(v, XlError):
        return v
    return xl.to_text(v)


def _fixed(a: Any) -> Any:
    """A non-vector argument as the runtime expects it: tables stay Range, scalars are cleaned."""
    if isinstance(a, Range):
        return a
    return _clean(a)


def _broadcast(a: pd.Series, index: pd.Index, columns: pd.Index) -> pd.DataFrame:
    """A Series indexed by row numbers repeats across columns; one indexed by letters repeats down rows."""
    if pd.api.types.is_integer_dtype(a.index):
        return pd.DataFrame({c: a.reindex(index).values for c in columns}, index=index, dtype=object)
    values = list(a.reindex(columns).values)
    return pd.DataFrame([values] * len(index), index=index, columns=columns, dtype=object)


def _apply(fn: Callable[..., Any], *args: Any) -> Any:
    """Apply a scalar Excel function element-wise over aligned vectors.

    Vectors are Series (a column block, indexed by row; or a row block, indexed
    by column letter) or DataFrames (a table block). Scalars broadcast; Range
    tables pass through whole.
    """
    frames = [a for a in args if isinstance(a, pd.DataFrame)]
    if frames:
        index, columns = frames[0].index, frames[0].columns
        aligned = []
        for a in args:
            if isinstance(a, pd.DataFrame):
                aligned.append(a.reindex(index=index, columns=columns))
            elif isinstance(a, pd.Series):
                aligned.append(_broadcast(a, index, columns))
            else:
                aligned.append(_fixed(a))
        data = {
            c: [fn(*[_clean(a.at[i, c]) if isinstance(a, pd.DataFrame) else a for a in aligned]) for i in index]
            for c in columns
        }
        return pd.DataFrame(data, index=index, dtype=object)[list(columns)]
    series = [a for a in args if isinstance(a, pd.Series)]
    if not series:
        return fn(*[_fixed(a) for a in args])
    index = series[0].index
    aligned = [a.reindex(index) if isinstance(a, pd.Series) else _fixed(a) for a in args]
    out = []
    for i in index:
        row_args = [_clean(a.at[i]) if isinstance(a, pd.Series) else a for a in aligned]
        out.append(fn(*row_args))
    return pd.Series(out, index=index, dtype=object)


def _frame_range(df: pd.DataFrame) -> Range:
    rows = [[_clean(v) for v in row] for row in df.itertuples(index=False, name=None)]
    return Range(rows)


# ---------------------------------------------------------------------------
# Column helpers used by generated vectorised code
# ---------------------------------------------------------------------------


def num(x: Any) -> Any:
    """Treat as numbers the way Excel arithmetic does (blank is 0, "12" is 12)."""
    if isinstance(x, (pd.Series, pd.DataFrame)):
        return x.map(_num_scalar)
    return _num_scalar(x)


def text(x: Any) -> Any:
    """Treat as text the way ``&`` does (blank is "", 12 is "12")."""
    if isinstance(x, (pd.Series, pd.DataFrame)):
        return x.map(_text_scalar)
    return _text_scalar(x)


def div(a: Any, b: Any) -> Any:
    return _apply(xl.div, a, b)


def power(a: Any, b: Any) -> Any:
    return _apply(xl.power, a, b)


def concat(*parts: Any) -> Any:
    return _apply(xl.concat, *parts)


def eq(a: Any, b: Any) -> Any:
    return _apply(xl.eq, a, b)


def ne(a: Any, b: Any) -> Any:
    return _apply(xl.ne, a, b)


def lt(a: Any, b: Any) -> Any:
    return _apply(xl.lt, a, b)


def le(a: Any, b: Any) -> Any:
    return _apply(xl.le, a, b)


def gt(a: Any, b: Any) -> Any:
    return _apply(xl.gt, a, b)


def ge(a: Any, b: Any) -> Any:
    return _apply(xl.ge, a, b)


def where(condition: Any, when_true: Any = True, when_false: Any = False) -> Any:
    """Excel IF, column-wise."""
    return _apply(xl.IF, condition, when_true, when_false)


def iferror(value: Any, value_if_error: Any) -> Any:
    return _apply(xl.IFERROR, value, value_if_error)


def isna(value: Any) -> Any:
    return _apply(xl.ISNA, value)


def and_(*args: Any) -> Any:
    return _apply(xl.AND, *args)


def or_(*args: Any) -> Any:
    return _apply(xl.OR, *args)


def not_(value: Any) -> Any:
    return _apply(xl.NOT, value)


def _as_table(table: Any) -> Any:
    return _frame_range(table) if isinstance(table, pd.DataFrame) else table


def vlookup(keys: Any, table: Any, col: int, exact: bool = False) -> Any:
    """VLOOKUP of a whole block of keys against a fixed table."""
    rng = _as_table(table)
    return _apply(lambda k: xl.VLOOKUP(k, rng, col, not exact), keys)


def hlookup(keys: Any, table: Any, row: int, exact: bool = False) -> Any:
    rng = _as_table(table)
    return _apply(lambda k: xl.HLOOKUP(k, rng, row, not exact), keys)


# Every other Excel function the runtime implements is available element-wise
# under a lower-case name (``sumifs``, ``eomonth``, ``countif`` ...). Names that
# would shadow a Python builtin get a trailing underscore (``round_``, ``abs_``).
_BUILTIN_CLASH = {"abs", "int", "round", "sum", "min", "max", "len", "and", "or", "not", "if"}
_HAND_WRITTEN = {"IF", "IFERROR", "ISNA", "AND", "OR", "NOT", "EXP", "LN", "LOG", "SQRT", "ABS", "INT", "ROUND", "POWER", "SUM", "MIN", "MAX", "AVERAGE", "VLOOKUP", "HLOOKUP", "UNSUPPORTED", "TODAY", "NOW", "RAND", "RANDBETWEEN", "INDIRECT", "OFFSET", "TRUE", "FALSE", "NA"}


def _vectorised(fn: Callable[..., Any], excel_name: str) -> Callable[..., Any]:
    def helper(*args: Any) -> Any:
        return _apply(fn, *args)

    helper.__name__ = helper.__qualname__ = _helper_name(excel_name)
    helper.__doc__ = f"Excel {excel_name}, applied element-wise over a block (scalars broadcast, tables whole)."
    return helper


def _helper_name(excel_name: str) -> str:
    name = excel_name.lower()
    return name + "_" if name in _BUILTIN_CLASH else name


VECTOR_FUNCS: dict[str, str] = {
    "IF": "where",
    "IFERROR": "iferror",
    "ISNA": "isna",
    "AND": "and_",
    "OR": "or_",
    "NOT": "not_",
    "EXP": "exp",
    "LN": "ln",
    "LOG": "log",
    "SQRT": "sqrt",
    "ABS": "abs_",
    "INT": "int_",
    "ROUND": "round_",
    "POWER": "power",
    "SUM": "sum_",
    "MIN": "min_",
    "MAX": "max_",
    "AVERAGE": "average",
}
for _excel_name, _fn in xl.FUNCTIONS.items():
    if _excel_name in _HAND_WRITTEN:
        continue
    _helper = _helper_name(_excel_name)
    globals()[_helper] = _vectorised(_fn, _excel_name)
    VECTOR_FUNCS[_excel_name] = _helper


def exp(x: Any) -> Any:
    return _apply(xl.EXP, x)


def ln(x: Any) -> Any:
    return _apply(xl.LN, x)


def log(x: Any, base: Any = 10) -> Any:
    return _apply(xl.LOG, x, base)


def sqrt(x: Any) -> Any:
    return _apply(xl.SQRT, x)


def abs_(x: Any) -> Any:
    return _apply(xl.ABS, x)


def int_(x: Any) -> Any:
    return _apply(xl.INT, x)


def round_(x: Any, digits: Any = 0) -> Any:
    return _apply(xl.ROUND, x, digits)


def sum_(*args: Any) -> Any:
    return _apply(xl.SUM, *args)


def min_(*args: Any) -> Any:
    return _apply(xl.MIN, *args)


def max_(*args: Any) -> Any:
    return _apply(xl.MAX, *args)


def average(*args: Any) -> Any:
    return _apply(xl.AVERAGE, *args)


def _target_rows(rows: slice) -> range:
    """``rows`` is an inclusive Excel row span, as in ``df.loc[r0:r1]``."""
    return range(rows.start, rows.stop + 1)


def _ensure_columns(df: pd.DataFrame, letters: Iterable[str]) -> None:
    for letter in letters:
        if letter not in df.columns:
            df[letter] = pd.Series(None, index=df.index, dtype=object)


def _at(df: pd.DataFrame, row: int | None, letter: str | None) -> Any:
    if row is None or letter is None or letter not in df.columns or row not in df.index:
        return None
    return df.at[row, letter]


def _offset_letter(letter: str, offset: int) -> str | None:
    """The column ``offset`` places to the left (positive) of ``letter``; None when off the sheet."""
    idx = col_index(letter) - offset
    return col_letter(idx) if idx >= 1 else None


def across(df: pd.DataFrame, row: int, letters: list[str], offset: int = 0) -> pd.Series:
    """One row over a row block's columns, read ``offset`` columns to the left (positive).

    The result is indexed by the block's own letters, so expressions over a row
    block align column by column and assign back with ``df.loc[row, letters]``.
    """
    _ensure_columns(df, letters)
    values = [_at(df, row, _offset_letter(letter, offset)) for letter in letters]
    return pd.Series(values, index=pd.Index(letters), dtype=object)


def cells(
    df: pd.DataFrame,
    rows: slice,
    letters: list[str],
    drow: int = 0,
    dcol: int = 0,
    fixed_row: int | None = None,
    fixed_col: str | None = None,
) -> pd.DataFrame:
    """A table block's view of another cell pattern, shaped like the block itself.

    Each target cell (r, c) reads (r - drow, c - dcol); a fixed row or column
    pins that axis instead (an absolute reference in the formula). Values come
    back as a DataFrame indexed by the block's rows and letters, so arithmetic
    between references aligns cell for cell and assigns back with
    ``df.loc[rows, letters]``.
    """
    _ensure_columns(df, letters)
    target_rows = list(_target_rows(rows))
    data: dict[str, list[Any]] = {}
    for letter in letters:
        src_col = fixed_col if fixed_col is not None else _offset_letter(letter, dcol)
        data[letter] = [_at(df, fixed_row if fixed_row is not None else r - drow, src_col) for r in target_rows]
    return pd.DataFrame(data, index=pd.Index(target_rows), dtype=object)[letters]


def col(df: pd.DataFrame, letter: str, rows: slice) -> pd.Series:
    """Column ``letter`` of a sheet frame over the block's rows.

    The frame may belong to another sheet with fewer rows, so the column is
    reindexed onto the block's rows (missing rows read as blank).
    """
    if letter not in df.columns:
        df[letter] = pd.Series(None, index=df.index, dtype=object)
    return df[letter].reindex(_target_rows(rows))


def shifted(df: pd.DataFrame, letter: str, offset: int, rows: slice) -> pd.Series:
    """Column ``letter`` read ``offset`` rows earlier (positive) or later (negative).

    Row ``r`` of the result holds the source's row ``r - offset``, whatever the
    source frame's own extent.
    """
    if letter not in df.columns:
        df[letter] = pd.Series(None, index=df.index, dtype=object)
    target = _target_rows(rows)
    source = df[letter].reindex([r - offset for r in target])
    source.index = pd.Index(target)
    return source


# ---------------------------------------------------------------------------
# Frames and the book
# ---------------------------------------------------------------------------


def _letters(n: int) -> list[str]:
    return [col_letter(i) for i in range(1, n + 1)]


class FrameGrid:
    """``xl.Grid``'s cell API over one sheet's DataFrame."""

    def __init__(self, title: str, df: pd.DataFrame):
        self.title = title
        self.df = df

    @staticmethod
    def _letter(col: int | str) -> str:
        return col_letter(col) if isinstance(col, int) else col.upper()

    def _ensure(self, row: int, letter: str) -> None:
        if letter not in self.df.columns:
            self.df[letter] = pd.Series(None, index=self.df.index, dtype=object)
        if row not in self.df.index:
            start = int(self.df.index.max()) + 1 if len(self.df.index) else 1
            extra = pd.DataFrame(index=range(start, row + 1), columns=self.df.columns, dtype=object)
            self.df = pd.concat([self.df, extra])

    def get(self, row: int, col: int | str) -> Any:
        letter = self._letter(col)
        if letter not in self.df.columns or row not in self.df.index:
            return None
        return _clean(self.df.at[row, letter])

    def set(self, row: int, col: int | str, value: Any) -> None:
        letter = self._letter(col)
        v = raw(value)
        self._ensure(row, letter)
        self.df.at[row, letter] = 0 if v is None else v

    def rng(self, r0: int, c0: int | str, r1: int, c1: int | str) -> Range:
        c0i, c1i = col_index(c0), col_index(c1)
        rows = [[self.get(r, c) for c in range(c0i, c1i + 1)] for r in range(r0, r1 + 1)]
        return Range(rows, f"{self.title}!{col_letter(c0i)}{r0}:{col_letter(c1i)}{r1}")

    def __getitem__(self, key: Any) -> Any:
        if isinstance(key, str):
            r0, c0, r1, c1 = parse_a1(key)
            if (r0, c0) == (r1, c1):
                return Val(self.get(r0, c0))
            return self.rng(r0, c0, r1, c1)
        row, col = key
        return Val(self.get(row, col))

    def __setitem__(self, key: Any, value: Any) -> None:
        if isinstance(key, str):
            r0, c0, r1, c1 = parse_a1(key)
            if (r0, c0) != (r1, c1):
                raise ValueError("can only assign to a single cell")
            self.set(r0, c0, value)
            return
        row, col = key
        self.set(row, col, value)

    def items(self) -> Iterator[tuple[tuple[int, int], Any]]:
        for letter in self.df.columns:
            col = col_index(letter)
            for row, v in self.df[letter].items():
                if not _is_blank(v):
                    yield (int(row), col), v

    def __len__(self) -> int:
        return sum(1 for _ in self.items())

    def block(self, a1: str) -> pd.DataFrame:
        """A fixed rectangle of the sheet as a DataFrame (lookup tables, SUM ranges).

        Always the full rectangle: rows beyond the frame's extent come back
        blank via ``reindex`` rather than by growing the frame, which would
        replace the DataFrame object under a function still holding it.
        """
        r0, c0, r1, c1 = parse_a1(a1)
        letters = _letters(c1)[c0 - 1 :]
        for letter in letters:
            if letter not in self.df.columns:
                self.df[letter] = pd.Series(None, index=self.df.index, dtype=object)
        return self.df[letters].reindex(range(r0, r1 + 1))


class FrameBook:
    """A workbook as DataFrames, with the same book-level API as ``xl.Book``."""

    def __init__(self) -> None:
        self.sheets: dict[str, FrameGrid] = {}
        self.names: dict[str, str] = {}
        self.frozen: dict[str, Any] = {}

    def frozen_value(self, name: str) -> Any:
        """See ``xl.Book.frozen_value``: volatile functions frozen to the file's last-saved time."""
        if name in self.frozen:
            return self.frozen[name]
        return Unsupported(f"volatile function {name}() has no frozen value (the file carries no last-saved date)")

    def __getitem__(self, title: str) -> FrameGrid:
        grid = self.sheets.get(title)
        if grid is None:
            grid = self.sheets[title] = FrameGrid(title, pd.DataFrame(index=range(1, 2), columns=["A"], dtype=object))
        return grid

    def frame(self, title: str) -> pd.DataFrame:
        return self[title].df

    def at(self, title: str, a1: str) -> Any:
        """A single cell as a raw scalar (None when blank)."""
        r0, c0, _, _ = parse_a1(a1)
        return self[title].get(r0, c0)

    def table(self, title: str, a1: str) -> Range:
        """A fixed rectangle as a lookup or aggregation table (always the full rectangle)."""
        return _frame_range(self[title].block(a1))

    def _name_text(self, name: str, scope: str | None) -> str | None:
        text = self.names.get(f"{scope}!{name}") if scope is not None else None
        return text if text is not None else self.names.get(name)

    def name(self, name: str, scope: str | None = None) -> Any:
        text = self._name_text(name, scope)
        if text is None:
            return Val(Unsupported(f"defined name {name!r} is missing or broken"))
        sheet, area = split_qualified(text)
        r0, c0, r1, c1 = parse_a1(area)
        grid = self[sheet]
        if (r0, c0) == (r1, c1):
            return Val(grid.get(r0, c0))
        return grid.rng(r0, c0, r1, c1)

    def name_table(self, name: str, scope: str | None = None) -> Any:
        """A defined name as a table (Range), or a scalar for a single cell."""
        text = self._name_text(name, scope)
        if text is None:
            return Unsupported(f"defined name {name!r} is missing or broken")
        sheet, area = split_qualified(text)
        r0, c0, r1, c1 = parse_a1(area)
        if (r0, c0) == (r1, c1):
            return self[sheet].get(r0, c0)
        return _frame_range(self[sheet].block(area))

    # -- serialisation ---------------------------------------------------------
    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "FrameBook":
        book = cls()
        dims = data.get("dims", {})
        for title, cells in data.get("sheets", {}).items():
            max_row, max_col = dims.get(title, [1, 1])
            for a1 in cells:
                r, c, _, _ = parse_a1(a1)
                max_row, max_col = max(max_row, r), max(max_col, c)
            df = pd.DataFrame(index=range(1, max_row + 1), columns=_letters(max_col), dtype=object)
            for a1, value in cells.items():
                r, c, _, _ = parse_a1(a1)
                df.at[r, col_letter(c)] = xl.Book.decode_value(value)
            book.sheets[title] = FrameGrid(title, df)
        for title, (max_row, max_col) in dims.items():
            if title not in book.sheets:
                book.sheets[title] = FrameGrid(title, pd.DataFrame(index=range(1, max_row + 1), columns=_letters(max_col), dtype=object))
        book.names = dict(data.get("names", {}))
        book.frozen = dict(data.get("frozen", {}))
        return book

    @classmethod
    def from_json(cls, path: str) -> "FrameBook":
        with open(path, encoding="utf-8") as fh:
            return cls.from_dict(json.load(fh))

    def to_dict(self) -> dict[str, Any]:
        out: dict[str, Any] = {"sheets": {}, "names": dict(self.names), "frozen": dict(self.frozen)}
        for title, grid in self.sheets.items():
            out["sheets"][title] = {f"{col_letter(c)}{r}": xl.Book.encode_value(v) for (r, c), v in sorted(grid.items())}
        return out