#!/usr/bin/env python3
# Copyright (c) 2026 Lark Technologies Pte. Ltd.
# SPDX-License-Identifier: MIT
"""Recommend a Lark Sheet chart size before creating the chart object."""

from __future__ import annotations

import argparse
import json
import re
from typing import Any

from lark_chart_size_rules import SUPPORTED_CHART_TYPES, recommend_chart_size
from lark_sheet_read_cli import (
    LarkCliError,
    emit_error,
    emit_success,
    envelope_data,
    extract_sheets,
    run_sheets,
    sheet_identifier,
    sheet_title,
)

ACTION = "chart_size_advisor"


def _locator(target: str) -> dict[str, str]:
    return {"url": target} if target.startswith(("http://", "https://")) else {"spreadsheet_token": target}


def _split_range_refs(value: str) -> list[str]:
    refs: list[str] = []
    start = 0
    quoted = False
    index = 0
    while index < len(value):
        char = value[index]
        if char == "'":
            if quoted and index + 1 < len(value) and value[index + 1] == "'":
                index += 2
                continue
            quoted = not quoted
        elif char == "," and not quoted:
            ref = value[start:index].strip()
            if ref:
                refs.append(ref)
            start = index + 1
        index += 1
    tail = value[start:].strip()
    if tail:
        refs.append(tail)
    if not refs or quoted:
        raise ValueError(f"Invalid data range: {value!r}")
    return refs


def _parse_ref(value: str) -> tuple[str | None, str]:
    raw = value.strip()
    sheet_name = None
    cell_range = raw
    if "!" in raw:
        sheet_name, cell_range = raw.rsplit("!", 1)
        sheet_name = sheet_name.strip()
        if len(sheet_name) >= 2 and sheet_name[0] == sheet_name[-1] == "'":
            sheet_name = sheet_name[1:-1].replace("''", "'")
    cell_range = cell_range.replace("$", "")
    if not re.fullmatch(r"[A-Za-z]+\d+:[A-Za-z]+\d+", cell_range):
        raise ValueError(f"Invalid A1 range: {value!r}")
    return sheet_name, cell_range


def _cell_value(cell: Any) -> Any:
    if not isinstance(cell, dict):
        return None
    return cell.get("value", cell.get("raw_value"))


def _matrix(data: dict[str, Any]) -> list[list[Any]]:
    ranges = data.get("ranges")
    if not isinstance(ranges, list) or not ranges:
        return []
    result = ranges[0]
    cells = result.get("cells") if isinstance(result, dict) else None
    if not isinstance(cells, list):
        return []
    return [
        [_cell_value(cell) for cell in row]
        for row in cells
        if isinstance(row, list)
    ]


def _is_network_timeout(exc: LarkCliError) -> bool:
    text = str(exc).strip()
    try:
        payload = json.loads(text)
    except json.JSONDecodeError:
        payload = None
    if isinstance(payload, dict):
        error = payload.get("error")
        if isinstance(error, dict) and str(error.get("subtype") or "").lower() == "timeout":
            return True
    lowered = text.lower()
    return "server time out" in lowered or "timed out" in lowered


def _run_read(
    shortcut: str,
    *,
    stage: str,
    timeout: int,
    **kwargs: Any,
) -> dict[str, Any]:
    retried = False
    while True:
        try:
            return run_sheets(shortcut, timeout=timeout, **kwargs)
        except LarkCliError as exc:
            if not retried and _is_network_timeout(exc):
                retried = True
                continue
            suffix = " after one retry" if retried else ""
            raise LarkCliError(f"{stage} failed{suffix}: {exc}", cmd=exc.cmd) from exc


def _sheet_selector(
    sheets: list[dict[str, Any]],
    *,
    explicit_name: str | None,
    worksheet_id: str | None,
    worksheet_name: str | None,
) -> dict[str, str]:
    if explicit_name:
        return {"sheet_name": explicit_name}
    if worksheet_id:
        return {"sheet_id": worksheet_id}
    if worksheet_name:
        return {"sheet_name": worksheet_name}
    if len(sheets) != 1:
        raise LarkCliError("Unqualified data ranges require --worksheet-id or --worksheet-name")
    sheet_id = sheet_identifier(sheets[0])
    if sheet_id:
        return {"sheet_id": sheet_id}
    return {"sheet_name": sheet_title(sheets[0])}


def _needs_workbook_metadata(
    ranges: list[str | None],
    *,
    worksheet_id: str | None,
    worksheet_name: str | None,
) -> bool:
    if worksheet_id or worksheet_name:
        return False
    return any(
        _parse_ref(ref)[0] is None
        for value in ranges
        if value
        for ref in _split_range_refs(value)
    )


def _read_ranges(
    locator: dict[str, str],
    sheets: list[dict[str, Any]],
    value: str,
    *,
    worksheet_id: str | None,
    worksheet_name: str | None,
    timeout: int,
    stage_prefix: str = "data range",
) -> list[list[list[Any]]]:
    matrices: list[list[list[Any]]] = []
    for ref in _split_range_refs(value):
        explicit_name, cell_range = _parse_ref(ref)
        selector = _sheet_selector(
            sheets,
            explicit_name=explicit_name,
            worksheet_id=worksheet_id,
            worksheet_name=worksheet_name,
        )
        data = envelope_data(
            _run_read(
                "+cells-get",
                stage=f"{stage_prefix} {ref}",
                **locator,
                **selector,
                flags={"range": cell_range, "include": "value"},
                timeout=timeout,
            )
        )
        matrix = _matrix(data)
        if not matrix:
            raise LarkCliError(f"No cells returned for {ref}")
        matrices.append(matrix)
    return matrices


def _combine(matrices: list[list[list[Any]]], direction: str) -> list[list[Any]]:
    if direction == "column":
        row_count = len(matrices[0])
        if any(len(matrix) != row_count for matrix in matrices):
            raise ValueError("Column-direction ranges must contain the same number of rows")
        return [sum((matrix[row] for matrix in matrices), []) for row in range(row_count)]
    column_count = max((len(row) for row in matrices[0]), default=0)
    if any(max((len(row) for row in matrix), default=0) != column_count for matrix in matrices):
        raise ValueError("Row-direction ranges must contain the same number of columns")
    return sum(matrices, [])


def _parse_indexes(value: str | None, *, dimension_count: int, dim1_index: int) -> list[int]:
    indexes = (
        [int(item.strip()) for item in value.split(",") if item.strip()]
        if value
        else [index for index in range(1, dimension_count + 1) if index != dim1_index]
    )
    if not indexes or any(index < 1 or index > dimension_count for index in indexes):
        raise ValueError("--dim2-indexes contains an out-of-range dimension index")
    if dim1_index in indexes:
        raise ValueError("--dim1-index cannot also appear in --dim2-indexes")
    return indexes


def _numeric(value: Any) -> float | None:
    if isinstance(value, (int, float)) and not isinstance(value, bool):
        return float(value)
    if isinstance(value, str):
        text = value.strip().replace(",", "")
        if text.endswith("%"):
            text = text[:-1]
        try:
            return float(text)
        except ValueError:
            return None
    return None


def profile_matrix(
    matrix: list[list[Any]],
    *,
    direction: str,
    dim1_index: int,
    dim2_indexes: str | None,
    detached_headers: list[Any] | None = None,
) -> dict[str, Any]:
    if not matrix:
        raise ValueError("Data range is empty")
    dimension_count = max((len(row) for row in matrix), default=0) if direction == "column" else len(matrix)
    if dim1_index < 1 or dim1_index > dimension_count:
        raise ValueError("--dim1-index is outside the data range")
    selected = _parse_indexes(
        dim2_indexes,
        dimension_count=dimension_count,
        dim1_index=dim1_index,
    )
    detached = detached_headers is not None
    if direction == "column":
        data_rows = matrix if detached else matrix[1:]
        categories = [row[dim1_index - 1] if len(row) >= dim1_index else None for row in data_rows]
        headers = detached_headers or matrix[0]
        series_names = [str(headers[index - 1]) if len(headers) >= index else f"Series {index}" for index in selected]
        first_values = [row[selected[0] - 1] if len(row) >= selected[0] else None for row in data_rows]
    else:
        category_row = matrix[dim1_index - 1]
        categories = category_row if detached else category_row[1:]
        headers = detached_headers or [row[0] if row else None for row in matrix]
        series_names = [str(headers[index - 1]) if len(headers) >= index else f"Series {index}" for index in selected]
        first_row = matrix[selected[0] - 1]
        first_values = first_row if detached else first_row[1:]
    nonempty_categories = [value for value in categories if value not in (None, "")]
    return {
        "categories": nonempty_categories,
        "series_names": series_names,
        "values": [number for value in first_values if (number := _numeric(value)) is not None],
        "dim2_indexes": selected,
    }


def _header_values(matrices: list[list[list[Any]]], direction: str) -> list[Any]:
    matrix = _combine(matrices, direction)
    if direction == "column":
        if len(matrix) != 1:
            raise ValueError("Column-direction --header-range must contain one row")
        return matrix[0]
    if any(len(row) != 1 for row in matrix):
        raise ValueError("Row-direction --header-range must contain one column")
    return [row[0] for row in matrix]


def _boolean_argument(value: str) -> bool:
    normalized = str(value).strip().lower()
    if normalized == "true":
        return True
    if normalized == "false":
        return False
    raise argparse.ArgumentTypeError("expected true or false")


def _comma_separated_values(value: str) -> list[str]:
    values = [item.strip() for item in str(value).split(",")]
    if not values or any(not item for item in values):
        raise argparse.ArgumentTypeError("expected a comma-separated list")
    return values


def _series_y_axes_argument(value: str) -> list[str]:
    values = _comma_separated_values(value)
    invalid = [item for item in values if item not in {"left", "right"}]
    if invalid:
        raise argparse.ArgumentTypeError(
            f"unsupported series Y axis {invalid[0]!r}; expected left or right"
        )
    return values


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Recommend chart width and height before +chart-create-basic")
    parser.add_argument("target", help="Spreadsheet URL or spreadsheet token")
    worksheet = parser.add_mutually_exclusive_group()
    worksheet.add_argument("--worksheet-id")
    worksheet.add_argument("--worksheet-name")
    parser.add_argument("--chart-type", choices=sorted(SUPPORTED_CHART_TYPES), required=True)
    parser.add_argument("--data-range", required=True)
    parser.add_argument("--header-range")
    parser.add_argument("--data-direction", choices=("column", "row"), default="column")
    parser.add_argument("--dim1-index", type=int, default=1)
    parser.add_argument("--dim2-indexes")
    parser.add_argument("--series-types", type=_comma_separated_values)
    parser.add_argument("--series-y-axes", type=_series_y_axes_argument)
    parser.add_argument("--data-labels", default="none")
    parser.add_argument("--aggregate-categories", type=_boolean_argument, default=True)
    parser.add_argument("--legend-position", default="bottom")
    parser.add_argument("--title", default="")
    parser.add_argument("--timeout", type=int, default=60)
    return parser.parse_args(argv)


def main() -> None:
    args = parse_args()
    locator = _locator(args.target)
    try:
        sheets: list[dict[str, Any]] = []
        if _needs_workbook_metadata(
            [args.data_range, args.header_range],
            worksheet_id=args.worksheet_id,
            worksheet_name=args.worksheet_name,
        ):
            workbook = envelope_data(
                _run_read(
                    "+workbook-info",
                    stage="workbook metadata",
                    **locator,
                    timeout=args.timeout,
                )
            )
            sheets = extract_sheets(workbook)
        matrices = _read_ranges(
            locator,
            sheets,
            args.data_range,
            worksheet_id=args.worksheet_id,
            worksheet_name=args.worksheet_name,
            timeout=args.timeout,
        )
        headers = None
        if args.header_range:
            header_matrices = _read_ranges(
                locator,
                sheets,
                args.header_range,
                worksheet_id=args.worksheet_id,
                worksheet_name=args.worksheet_name,
                timeout=args.timeout,
                stage_prefix="header range",
            )
            headers = _header_values(header_matrices, args.data_direction)
        profile = profile_matrix(
            _combine(matrices, args.data_direction),
            direction=args.data_direction,
            dim1_index=args.dim1_index,
            dim2_indexes=args.dim2_indexes,
            detached_headers=headers,
        )
        result = recommend_chart_size(
            chart_type=args.chart_type,
            categories=profile["categories"],
            series_names=profile["series_names"],
            data_labels=args.data_labels,
            legend_position=args.legend_position,
            title=args.title,
            values=profile["values"],
            aggregate_categories=args.aggregate_categories,
            series_types=args.series_types,
            series_y_axes=args.series_y_axes,
        )
        result["data_profile"] = {
            "dim2_indexes": profile["dim2_indexes"],
            **result.pop("evidence"),
        }
    except (LarkCliError, KeyError, TypeError, ValueError) as exc:
        emit_error(ACTION, str(exc))
        raise SystemExit(1) from exc
    emit_success(ACTION, result)


if __name__ == "__main__":
    main()