"""Fixed Origin adapter for origin.plot@v2.

This process accepts exactly one argument: a Node-created job directory. It never
installs packages, evaluates user code, downloads data, or resolves paths from the
request. terminal.json is its only terminal-state contract.
"""
from __future__ import annotations

import csv
import hashlib
import json
import math
import os
import sys
from datetime import datetime, timezone
from importlib.metadata import PackageNotFoundError, version
from pathlib import Path
from typing import Any

PLOT_CONFIG = {
    "line": ("line", "l"),
    "scatter": ("scatter", "s"),
    "line_scatter": ("linesymb", "y"),
    "column": ("column", "c"),
    "bar": ("bar", 215),
    "grouped_column": ("column", "c"),
    "y_error": ("ERRBAR", "y"),
    "contour": ("TriContour", 243),
    "surface_3d": ("glCMAP", 103),
    "ternary": ("ternary", 245),
}
XYZ_PLOT_TYPES = {"contour", "surface_3d", "ternary", "heatmap"}
FORMATS = {"opju", "png", "svg", "pdf"}
LINE_STYLES = {
    "solid": 1,
    "dash": 2,
    "dot": 3,
    "dash_dot": 4,
    "dash_dot_dot": 5,
}
SYMBOLS = {
    "square": 0,
    "circle": 1,
    "triangle_up": 2,
    "diamond": 3,
    "cross": 9,
    "plus": 10,
}
LEGEND_POSITIONS = {
    "top_left": (700, 500),
    "top_right": (6800, 500),
    "bottom_left": (700, 7200),
    "bottom_right": (6800, 7200),
}
ADAPTER_VERSION = "0.5.0"


def _probe() -> int:
    health = "ready"
    detail = "Origin COM 与托管 Python 运行时可用"
    software_version = None
    try:
        if sys.platform != "win32":
            raise RuntimeError("Origin adapter requires Windows")
        import winreg
        import originpro

        with winreg.OpenKey(winreg.HKEY_CLASSES_ROOT, r"Origin.ApplicationSI\CLSID"):
            pass
        originpro_version = version("originpro")
        detail = f"Origin COM 与托管 Python 运行时可用（originpro {originpro_version}）"
        del originpro
    except (FileNotFoundError, ImportError, OSError, PackageNotFoundError, RuntimeError) as exc:
        health = "unavailable"
        detail = str(exc)
    print(json.dumps({
        "adapter_version": ADAPTER_VERSION,
        "software": "OriginPro",
        "software_version": software_version,
        "health": health,
        "detail": detail,
    }, ensure_ascii=False))
    return 0


def _validate_semantics(request: dict[str, Any]) -> None:
    input_keys = {item["key"] for item in request["inputs"]}
    plot = request["operation"]["plot"]
    plot_type = plot["type"]
    series = plot["series"]
    used_inputs = {item["input"] for item in series}
    if used_inputs != input_keys:
        raise ValueError("INPUT_BINDINGS_MUST_BE_USED_EXACTLY")
    if plot_type == "grouped_column" and len(series) < 2:
        raise ValueError("GROUPED_COLUMN_REQUIRES_MULTIPLE_SERIES")
    if plot_type in XYZ_PLOT_TYPES and len(series) != 1:
        raise ValueError("XYZ_PLOT_REQUIRES_ONE_SERIES")
    required_roles = (
        ("x", "y", "z") if plot_type in XYZ_PLOT_TYPES
        else ("x", "y", "y_error") if plot_type == "y_error"
        else ("x", "y")
    )
    for item in series:
        if any(role not in item for role in required_roles):
            raise ValueError("SERIES_REQUIRED_ROLE_MISSING")
    for axis_name in ("x_axis", "y_axis", "z_axis"):
        axis = plot.get(axis_name) or {}
        minimum = axis.get("minimum")
        maximum = axis.get("maximum")
        if minimum is not None and maximum is not None and minimum >= maximum:
            raise ValueError(f"{axis_name.upper()}_LIMITS_INVALID")
        if axis.get("scale") in {"log10", "ln", "log2"} and (
            minimum is not None and minimum <= 0 or maximum is not None and maximum <= 0
        ):
            raise ValueError(f"{axis_name.upper()}_LOG_LIMIT_INVALID")


def _atomic_json(path: Path, value: Any) -> None:
    temporary = path.with_name(path.name + ".tmp-" + os.urandom(8).hex())
    try:
        with temporary.open("w", encoding="utf-8", newline="\n") as handle:
            json.dump(value, handle, ensure_ascii=False, indent=2)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary, path)
    finally:
        temporary.unlink(missing_ok=True)


def _read_rows(path: Path, sheet: str | None) -> tuple[list[str], list[list[Any]]]:
    suffix = path.suffix.lower()
    if suffix == ".csv":
        with path.open("r", encoding="utf-8-sig", newline="") as handle:
            rows = list(csv.reader(handle))
        if len(rows) < 2:
            raise ValueError("CSV_INPUT_EMPTY")
        return [str(item) for item in rows[0]], rows[1:]
    if suffix == ".json":
        value = json.loads(path.read_text(encoding="utf-8"))
        if isinstance(value, list) and value and all(isinstance(item, dict) for item in value):
            headers = list(value[0])
            return headers, [[item.get(name) for name in headers] for item in value]
        if isinstance(value, dict) and value and all(isinstance(item, list) for item in value.values()):
            headers = list(value)
            length = max(len(value[name]) for name in headers)
            return headers, [[value[name][index] if index < len(value[name]) else None for name in headers] for index in range(length)]
        raise ValueError("JSON_INPUT_SHAPE_UNSUPPORTED")
    if suffix == ".xlsx":
        from openpyxl import load_workbook

        workbook = load_workbook(path, read_only=True, data_only=True)
        try:
            worksheet = workbook[sheet] if sheet else workbook.active
            rows = list(worksheet.iter_rows(values_only=True))
        finally:
            workbook.close()
        if len(rows) < 2:
            raise ValueError("XLSX_INPUT_EMPTY")
        return [str(item or "") for item in rows[0]], [list(row) for row in rows[1:]]
    raise ValueError("INPUT_TYPE_UNSUPPORTED")


def _column_index(headers: list[str], value: Any, field: str) -> int:
    if not isinstance(value, str) or value not in headers:
        raise ValueError(f"{field.upper()}_COLUMN_NOT_FOUND")
    return headers.index(value)


def _manifest(path: Path, media_type: str) -> dict[str, Any]:
    return {
        "artifact_id": {
            "project.opju": "project",
            "figure.png": "figure_png",
            "figure.svg": "figure_svg",
            "figure.pdf": "figure_pdf",
            "plot-spec.json": "plot_spec",
            "provenance.json": "provenance",
        }[path.name],
        "filename": path.name,
        "media_type": media_type,
        "size_bytes": path.stat().st_size,
        "sha256": _file_sha256(path),
    }


def _file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _axis_title(axis: Any, fallback: str) -> str:
    if not isinstance(axis, dict):
        return fallback
    title = str(axis.get("title") or fallback)
    unit = str(axis.get("unit") or "")
    return f"{title} ({unit})" if unit else title


def _hex_color(value: str) -> tuple[int, int, int]:
    return tuple(int(value[index:index + 2], 16) for index in (1, 3, 5))


def _apply_canvas(graph: Any, canvas: Any) -> None:
    if not isinstance(canvas, dict):
        return
    width = float(canvas["width_mm"])
    height = float(canvas["height_mm"])
    graph.set_float("width", width / 25.4 * graph.get_float("resx"))
    graph.set_float("height", height / 25.4 * graph.get_float("resy"))


def _apply_axis(layer: Any, name: str, spec: Any, fallback: str) -> None:
    settings = spec if isinstance(spec, dict) else {}
    axis = layer.axis(name)
    axis.title = _axis_title(settings, fallback)
    if "scale" in settings:
        axis.scale = settings["scale"]
    if any(key in settings for key in ("minimum", "maximum", "major_step")):
        axis.set_limits(
            settings.get("minimum"), settings.get("maximum"), settings.get("major_step")
        )
    if "tick_label_angle" in settings:
        layer.set_float(f"{name}.label.rotate", float(settings["tick_label_angle"]))
    if "tick_label_font_size" in settings:
        layer.set_float(f"{name}.label.pt", float(settings["tick_label_font_size"]))
    if "title_font_size" in settings:
        title_object = {"x": "xb", "y": "yl", "z": "zf"}[name]
        layer.label(title_object).set_int("fsize", round(settings["title_font_size"]))
    if "grid" in settings:
        grid_value = {"none": 0, "major": 1, "major_minor": 3}[settings["grid"]]
        layer.set_int(f"{name}.grid.show", grid_value)


def _apply_series_style(origin_plot: Any, style: Any) -> None:
    if not isinstance(style, dict):
        return
    if "color" in style:
        origin_plot.color = _hex_color(style["color"])
    if "line_width" in style:
        origin_plot.set_float("line.width", float(style["line_width"]))
    if "line_style" in style:
        origin_plot.set_int("line.type", LINE_STYLES[style["line_style"]])
    if "symbol" in style:
        origin_plot.symbol_kind = SYMBOLS[style["symbol"]]
    if "symbol_size" in style:
        origin_plot.symbol_size = float(style["symbol_size"])
    if "transparency" in style:
        origin_plot.transparency = int(style["transparency"])


def _apply_legend(layer: Any, legend: Any) -> None:
    if not isinstance(legend, dict):
        return
    label = layer.label("Legend")
    label.show = legend.get("enabled", True)
    if "font_size" in legend:
        label.set_int("fsize", round(legend["font_size"]))
    if "position" in legend:
        left, top = LEGEND_POSITIONS[legend["position"]]
        label.set_int("left", left)
        label.set_int("top", top)


def _input_file(job_dir: Path, key: str) -> Path:
    directory = job_dir / "input" / key
    files = [path for path in directory.iterdir() if path.is_file() and not path.name.startswith(".")]
    if len(files) != 1:
        raise ValueError(f"INPUT_FILE_COUNT_INVALID:{key}")
    return files[0]


def _resolve_series(
    input_data: dict[str, tuple[list[str], list[list[Any]]]],
    series_specs: list[dict[str, Any]],
) -> tuple[list[dict[str, Any]], dict[tuple[str, int], str]]:
    resolved: list[dict[str, Any]] = []
    labels: dict[tuple[str, int], str] = {}
    for series in series_specs:
        input_key = series["input"]
        headers, _ = input_data[input_key]
        role_indexes = {
            role: _column_index(headers, series[role], role)
            for role in ("x", "y", "z", "y_error") if role in series
        }
        y_index = role_indexes["y"]
        label = series.get("label")
        label_key = (input_key, y_index)
        effective_label = label or series["y"]
        if label_key in labels and labels[label_key] != effective_label:
            raise ValueError("SERIES_LABEL_CONFLICT")
        labels[label_key] = effective_label
        resolved.append({"input": input_key, "label": label, **role_indexes})
    return resolved, labels


def _number(value: Any, role: str) -> float:
    if isinstance(value, bool):
        raise ValueError(f"{role.upper()}_VALUE_NOT_NUMERIC")  # noqa: TRY004
    try:
        number = float(value)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{role.upper()}_VALUE_NOT_NUMERIC") from exc
    if not math.isfinite(number):
        raise ValueError(f"{role.upper()}_VALUE_NOT_FINITE")
    return number


def _is_evenly_spaced(values: list[float]) -> bool:
    if len(values) <= 2:
        return True
    step = values[1] - values[0]
    tolerance = max(abs(step) * 1e-9, 1e-12)
    return all(
        math.isclose(current - previous, step, rel_tol=1e-9, abs_tol=tolerance)
        for previous, current in zip(values[1:-1], values[2:], strict=True)
    )


def _heatmap_matrix(
    rows: list[list[Any]], resolved: dict[str, Any]
) -> tuple[list[list[float]], tuple[float, float, float, float]]:
    points: dict[tuple[float, float], float] = {}
    for row in rows:
        try:
            x = _number(row[resolved["x"]], "x")
            y = _number(row[resolved["y"]], "y")
            z = _number(row[resolved["z"]], "z")
        except IndexError as exc:
            raise ValueError("HEATMAP_ROW_INCOMPLETE") from exc
        if (x, y) in points:
            raise ValueError("HEATMAP_COORDINATES_DUPLICATED")
        points[(x, y)] = z
    x_values = sorted({item[0] for item in points})
    y_values = sorted({item[1] for item in points})
    if len(x_values) < 2 or len(y_values) < 2:
        raise ValueError("HEATMAP_GRID_TOO_SMALL")
    if len(points) != len(x_values) * len(y_values):
        raise ValueError("HEATMAP_GRID_INCOMPLETE")
    if not _is_evenly_spaced(x_values) or not _is_evenly_spaced(y_values):
        raise ValueError("HEATMAP_GRID_NOT_REGULAR")
    matrix = [[points[(x, y)] for x in x_values] for y in y_values]
    return matrix, (x_values[0], x_values[-1], y_values[0], y_values[-1])


def _validate_artifact(path: Path, extension: str) -> None:
    if not path.is_file() or path.stat().st_size == 0:
        raise RuntimeError(f"{extension.upper()}_EXPORT_EMPTY")
    head = path.read_bytes()[:1024]
    if extension == "png" and not head.startswith(b"\x89PNG\r\n\x1a\n"):
        raise RuntimeError("PNG_EXPORT_INVALID")
    if extension == "pdf" and not head.startswith(b"%PDF-"):
        raise RuntimeError("PDF_EXPORT_INVALID")
    if extension == "svg" and b"<svg" not in head.lower():
        raise RuntimeError("SVG_EXPORT_INVALID")
    if extension == "opju" and len(head) < 64:
        raise RuntimeError("OPJU_EXPORT_INVALID")


def run(job_dir: Path) -> list[dict[str, Any]]:
    job_dir = job_dir.resolve(strict=True)
    request_record = json.loads((job_dir / "request" / "request.json").read_text(encoding="utf-8"))
    request = request_record["request"]
    _validate_semantics(request)
    input_specs = request["inputs"]
    input_files = {item["key"]: _input_file(job_dir, item["key"]) for item in input_specs}
    input_data = {
        item["key"]: _read_rows(
            input_files[item["key"]],
            (item.get("selector") or {}).get("sheet"),
        )
        for item in input_specs
    }
    plot_spec = request["operation"]["plot"]
    plot_type = plot_spec["type"]
    if plot_type not in PLOT_CONFIG and plot_type != "heatmap":
        raise ValueError("PLOT_TYPE_NOT_IMPLEMENTED")
    series_specs = plot_spec["series"]
    resolved_series, labels = _resolve_series(input_data, series_specs)

    import originpro as op

    output = job_dir / "output"
    output.mkdir(exist_ok=True)
    op.set_show(False)
    try:
        op.new()
        worksheets: dict[str, Any] = {}
        for input_spec in input_specs:
            input_key = input_spec["key"]
            headers, rows = input_data[input_key]
            worksheet = op.new_sheet("w", lname=input_key)
            worksheets[input_key] = worksheet
            for index, header in enumerate(headers):
                column_label = labels.get((input_key, index), header)
                worksheet.from_list(
                    index,
                    [row[index] if index < len(row) else None for row in rows],
                    lname=column_label,
                )
        if plot_type == "heatmap":
            import numpy as np

            resolved = resolved_series[0]
            _, rows = input_data[resolved["input"]]
            matrix, xy_map = _heatmap_matrix(rows, resolved)
            matrix_sheet = op.new_sheet("m")
            matrix_sheet.from_np(np.array(matrix, dtype=float))
            matrix_sheet.xymap = xy_map
            matrix_sheet.set_label(
                0, _axis_title(plot_spec.get("z_axis"), str(series_specs[0]["z"]))
            )
            graph = op.new_graph(template="heatmap")
            layer = graph[0]
            origin_plots = [layer.add_mplot(matrix_sheet, 0, type=105)]
        else:
            template, origin_plot_type = PLOT_CONFIG[plot_type]
            graph = op.new_graph(template=template)
            layer = graph[0]
            origin_plots = []
            for resolved in resolved_series:
                arguments = {
                    "coly": resolved["y"],
                    "colx": resolved["x"],
                    "type": origin_plot_type,
                }
                if plot_type in XYZ_PLOT_TYPES:
                    arguments["colz"] = resolved["z"]
                if plot_type == "y_error":
                    arguments["colyerr"] = resolved["y_error"]
                origin_plots.append(layer.add_plot(worksheets[resolved["input"]], **arguments))
            if plot_type == "grouped_column":
                layer.group()
        layer.rescale()
        _apply_canvas(graph, plot_spec.get("canvas"))
        _apply_axis(
            layer, "x", plot_spec.get("x_axis"), str(series_specs[0].get("x") or "X")
        )
        _apply_axis(layer, "y", plot_spec.get("y_axis"), "Y")
        if plot_type == "surface_3d":
            _apply_axis(
                layer, "z", plot_spec.get("z_axis"), str(series_specs[0].get("z") or "Z")
            )
        for origin_plot, series_spec in zip(origin_plots, series_specs, strict=True):
            _apply_series_style(origin_plot, series_spec.get("style"))
        _apply_legend(layer, plot_spec.get("legend"))
        if plot_spec.get("title"):
            title = layer.add_label(str(plot_spec["title"]))
            title_font_size = (plot_spec.get("title_style") or {}).get("font_size", 18)
            title.set_int("fsize", round(title_font_size))
            title.set_int("left", 2200)
            title.set_int("top", 120)
        requested_outputs = request["outputs"]
        formats = [item["format"] for item in requested_outputs]
        if any(item not in FORMATS for item in formats):
            raise ValueError("OUTPUT_FORMAT_UNSUPPORTED")
        artifacts: list[dict[str, Any]] = []
        if "opju" in formats:
            project = output / "project.opju"
            op.save(str(project))
            _validate_artifact(project, "opju")
            artifacts.append(_manifest(project, "application/x-origin-project"))
        media = {"png": "image/png", "svg": "image/svg+xml", "pdf": "application/pdf"}
        png_output = next((item for item in requested_outputs if item["format"] == "png"), None)
        dpi = (png_output.get("options") or {}).get("dpi", 300) if png_output else 300
        if not isinstance(dpi, int) or isinstance(dpi, bool) or not 72 <= dpi <= 1200:
            raise ValueError("OUTPUT_DPI_INVALID")
        canvas_width_mm = (plot_spec.get("canvas") or {}).get("width_mm", 160)
        pixel_width = round(dpi * canvas_width_mm / 25.4)
        for extension in ("png", "svg", "pdf"):
            if extension in formats:
                target = output / f"figure.{extension}"
                exported = Path(graph.save_fig(
                    str(target),
                    type=extension,
                    width=pixel_width if extension == "png" else 0,
                    ratio=100 if extension in {"svg", "pdf"} else 0,
                )).resolve()
                if exported != target.resolve() or not target.is_file():
                    raise RuntimeError(f"{extension.upper()}_EXPORT_FAILED")
                _validate_artifact(target, extension)
                artifacts.append(_manifest(target, media[extension]))
        try:
            originpro_version = version("originpro")
        except PackageNotFoundError:
            originpro_version = "embedded"
        provenance = {
            "adapter_version": ADAPTER_VERSION,
            "originpro_version": originpro_version,
            "request_digest": request_record["request_digest"],
            "inputs": [
                {
                    "key": item["key"],
                    "filename": input_files[item["key"]].name,
                    "sha256": _file_sha256(input_files[item["key"]]),
                }
                for item in input_specs
            ],
            "outputs": requested_outputs,
            "requested_dpi": dpi if png_output else None,
            "png_pixel_width": pixel_width if png_output else None,
            "canvas_mm": plot_spec.get("canvas"),
        }
        plot_spec_path = output / "plot-spec.json"
        provenance_path = output / "provenance.json"
        _atomic_json(plot_spec_path, request)
        _atomic_json(provenance_path, provenance)
        artifacts.append(_manifest(plot_spec_path, "application/json"))
        artifacts.append(_manifest(provenance_path, "application/json"))
        return artifacts
    finally:
        if op.oext:
            op.exit()


def main() -> int:
    if sys.argv[1:] == ["--probe"]:
        return _probe()
    if len(sys.argv) != 2:
        print("[ERR] Usage: worker.py <job-directory>", file=sys.stderr)
        return 2
    job_dir = Path(sys.argv[1])
    request_record: dict[str, Any] = {}
    try:
        request_record = json.loads((job_dir / "request" / "request.json").read_text(encoding="utf-8"))
        artifacts = run(job_dir)
        terminal = {
            "job_id": request_record["job_id"],
            "lease_id": request_record["lease_id"],
            "request_digest": request_record["request_digest"],
            "status": "succeeded",
            "error": {},
            "artifact_manifest": artifacts,
            "terminal_at": datetime.now(timezone.utc).isoformat(),
        }
        _atomic_json(job_dir / "artifacts.json", artifacts)
        _atomic_json(job_dir / "terminal.json", terminal)
        print("[OK] Origin job completed.")
        return 0
    except Exception as exception:
        terminal = {
            "job_id": request_record.get("job_id", ""),
            "lease_id": request_record.get("lease_id", ""),
            "request_digest": request_record.get("request_digest", ""),
            "status": "failed",
            "error": {"code": type(exception).__name__, "detail": str(exception)[:500]},
            "artifact_manifest": [],
            "terminal_at": datetime.now(timezone.utc).isoformat(),
        }
        _atomic_json(job_dir / "terminal.json", terminal)
        print(f"[ERR] {type(exception).__name__}: {exception}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())