#!/usr/bin/env python3
"""Convert XXTLanControl assist-task color results into Lua snippets."""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path
from typing import Any


def load_json(path: str) -> Any:
    if path == "-":
        return json.load(sys.stdin)
    with open(path, "r", encoding="utf-8") as handle:
        return json.load(handle)


def select_step(payload: Any, step_id: str | None) -> Any:
    if not step_id:
        return payload
    if not isinstance(payload, dict):
        raise SystemExit("--step requires a manifest object")
    for step in payload.get("steps", []):
        if isinstance(step, dict) and str(step.get("stepId", "")) == step_id:
            return step
    raise SystemExit(f"step not found: {step_id}")


def unwrap_result(payload: Any) -> tuple[dict[str, Any], dict[str, Any]]:
    if not isinstance(payload, dict):
        raise SystemExit("input must be a JSON object")
    meta = payload
    data = payload
    if isinstance(data.get("task"), dict):
        meta = data["task"]
        data = data["task"]
    if isinstance(data.get("result"), dict):
        data = data["result"]
    return meta, data


def normalize_color(value: Any) -> str:
    if isinstance(value, int):
        if value <= 0xFFFFFF:
            return f"0x{value:06X}"
        return f"0x{value:X}"
    text = str(value or "").strip()
    if not text:
        raise ValueError("missing color")
    text = re.sub(r"^(?:0x|#)", "", text, flags=re.IGNORECASE)
    text = re.sub(r"[^0-9a-fA-F]", "", text).upper()
    if len(text) == 3:
        text = "".join(ch * 2 for ch in text)
    if not text:
        raise ValueError("invalid color")
    if len(text) > 8:
        text = text[-8:]
    return f"0x{text}"


def as_int(value: Any, field: str) -> int:
    try:
        return int(round(float(value)))
    except (TypeError, ValueError):
        raise ValueError(f"invalid {field}: {value!r}") from None


def normalize_point(item: Any) -> dict[str, Any] | None:
    if isinstance(item, (list, tuple)) and len(item) >= 3:
        point = {"x": item[0], "y": item[1], "color": item[2]}
        if len(item) >= 4:
            point["similarity"] = item[3]
    elif isinstance(item, dict):
        point = item
    else:
        return None

    try:
        out: dict[str, Any] = {
            "x": as_int(point.get("x"), "x"),
            "y": as_int(point.get("y"), "y"),
            "color": normalize_color(point.get("color", point.get("hex", point.get("rgb")))),
        }
        sim = point.get("similarity", point.get("sim", point.get("csim")))
        if sim is not None and str(sim).strip() != "":
            out["similarity"] = as_int(sim, "similarity")
        return out
    except ValueError:
        return None


def collect_points(result: dict[str, Any], include_records: bool) -> list[dict[str, Any]]:
    raw_points: list[Any] = []
    for key in ("points", "statePoints", "colorPoints", "colors", "group"):
        value = result.get(key)
        if isinstance(value, list):
            raw_points.extend(value)
    if include_records and isinstance(result.get("records"), list):
        raw_points.extend(result["records"])

    points: list[dict[str, Any]] = []
    seen: set[tuple[int, int, str]] = set()
    for item in raw_points:
        point = normalize_point(item)
        if not point:
            continue
        key = (point["x"], point["y"], point["color"])
        if key in seen:
            continue
        seen.add(key)
        points.append(point)
    return points


def parse_tap(text: str | None) -> dict[str, int] | None:
    if not text:
        return None
    parts = re.split(r"[,:\s]+", text.strip())
    if len(parts) < 2:
        raise SystemExit("--tap must be formatted as x,y")
    return {"x": as_int(parts[0], "tap x"), "y": as_int(parts[1], "tap y")}


def result_tap(result: dict[str, Any], override: str | None) -> dict[str, int] | None:
    explicit = parse_tap(override)
    if explicit:
        return explicit
    for key in ("actionPoint", "tap", "action"):
        value = result.get(key)
        if isinstance(value, dict) and "x" in value and "y" in value:
            return {"x": as_int(value["x"], "tap x"), "y": as_int(value["y"], "tap y")}
    return None


def lua_string(text: str) -> str:
    return "'" + text.replace("\\", "\\\\").replace("'", "\\'") + "'"


def lua_name(value: str | None, meta: dict[str, Any]) -> str:
    raw = value or str(meta.get("stepId") or meta.get("title") or meta.get("id") or "screen")
    name = re.sub(r"[^0-9A-Za-z_]+", "_", raw.strip()).strip("_")
    return name or "screen"


def point_line(point: dict[str, Any], indent: str) -> str:
    values = [str(point["x"]), str(point["y"]), point["color"]]
    if "similarity" in point:
        values.append(str(point["similarity"]))
    return f"{indent}{{{', '.join(values)}}},"


def render_points(points: list[dict[str, Any]]) -> str:
    lines = ["{"]
    lines.extend(point_line(point, "    ") for point in points)
    lines.append("}")
    return "\n".join(lines)


def render_is_colors(points: list[dict[str, Any]], similarity: int) -> str:
    lines = ["screen.keep()", "local ok = screen.is_colors({"]
    lines.extend(point_line(point, "    ") for point in points)
    lines.extend([f"}}, {similarity})", "screen.unkeep()"])
    return "\n".join(lines)


def render_tap(tap: dict[str, int] | None) -> str:
    if not tap:
        return "-- TODO: add touch.tap(x, y)"
    return f"touch.tap({tap['x']}, {tap['y']})\nsys.msleep(300)"


def render_xxtdo_screen(points: list[dict[str, Any]], name: str, similarity: int, tap: dict[str, int] | None) -> str:
    lines = [
        "{",
        f"    name = {lua_string(name)},",
        f"    csim = {similarity},",
    ]
    lines.extend(point_line(point, "    ") for point in points)
    lines.append("    run = function()")
    if tap:
        lines.append(f"        touch.tap({tap['x']}, {tap['y']})")
        lines.append("        sys.msleep(300)")
    else:
        lines.append("        -- TODO: add next action")
    lines.extend(["    end,", "}"])
    return "\n".join(lines)


def write_output(text: str, output: str | None) -> None:
    if output:
        Path(output).write_text(text + "\n", encoding="utf-8")
    else:
        print(text)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", help="JSON file path, or - for stdin")
    parser.add_argument("--step", help="Select a manifest steps[].stepId")
    parser.add_argument("--format", choices=("points", "is-colors", "xxtdo-screen", "tap", "json"), default="points")
    parser.add_argument("--name", help="Screen name for xxtdo-screen output")
    parser.add_argument("--similarity", type=int, default=90)
    parser.add_argument("--tap", help="Override tap point as x,y")
    parser.add_argument("--no-records", action="store_true", help="Do not merge result.records into points")
    parser.add_argument("--output", help="Write output to file instead of stdout")
    args = parser.parse_args()

    selected = select_step(load_json(args.input), args.step)
    meta, result = unwrap_result(selected)
    points = collect_points(result, include_records=not args.no_records)
    tap = result_tap(result, args.tap)

    if args.format in {"points", "is-colors", "xxtdo-screen"} and not points:
        raise SystemExit("no color points found in result")

    if args.format == "points":
        output = render_points(points)
    elif args.format == "is-colors":
        output = render_is_colors(points, args.similarity)
    elif args.format == "xxtdo-screen":
        output = render_xxtdo_screen(points, lua_name(args.name, meta), args.similarity, tap)
    elif args.format == "tap":
        output = render_tap(tap)
    else:
        output = json.dumps({"points": points, "tap": tap}, ensure_ascii=False, indent=2)

    write_output(output, args.output)


if __name__ == "__main__":
    main()
