#!/usr/bin/env python3
from __future__ import annotations

import argparse
import hashlib
import json
import re
import urllib.request
from collections import defaultdict
from pathlib import Path
from typing import Any

from PIL import Image


VISUAL_RESOURCE_PROPS = (
    ("layout", "icon-image", "icon"),
    ("paint", "fill-pattern", "fill-pattern"),
    ("paint", "line-pattern", "line-pattern"),
)

VECTOR_PROPS = {
    "fill": ("fill-color", "fill-opacity", "fill-outline-color"),
    "line": ("line-color", "line-width", "line-dasharray"),
    "circle": ("circle-color", "circle-radius", "circle-stroke-color", "circle-stroke-width"),
    "symbol": ("text-color", "text-halo-color", "text-halo-width"),
}


def canonical(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))


def sid(prefix: str, value: Any) -> str:
    return prefix + "-" + hashlib.sha256(canonical(value).encode("utf-8")).hexdigest()[:16]


def load_json(path_or_url: str) -> dict[str, Any]:
    if re.match(r"^https?://", path_or_url):
        with urllib.request.urlopen(path_or_url, timeout=30) as r:
            return json.loads(r.read().decode("utf-8"))
    return json.loads(Path(path_or_url).read_text(encoding="utf-8"))


def resolve_sprite_base(style: dict[str, Any], override: str | None) -> str | None:
    return override or style.get("sprite")


def load_sprite(base: str | None, output: Path) -> tuple[dict[str, Any], Image.Image | None]:
    if not base:
        return {}, None
    output.mkdir(parents=True, exist_ok=True)
    if re.match(r"^https?://", base):
        with urllib.request.urlopen(base + ".json", timeout=30) as r:
            meta_bytes = r.read()
        with urllib.request.urlopen(base + ".png", timeout=60) as r:
            png_bytes = r.read()
        (output / "sprite.json").write_bytes(meta_bytes)
        (output / "sprite.png").write_bytes(png_bytes)
    else:
        base_path = Path(base)
        (output / "sprite.json").write_bytes(base_path.with_suffix(".json").read_bytes())
        (output / "sprite.png").write_bytes(base_path.with_suffix(".png").read_bytes())
    meta = json.loads((output / "sprite.json").read_text(encoding="utf-8"))
    image = Image.open(output / "sprite.png").convert("RGBA")
    return meta, image


def walk_gets(value: Any, out: set[str]) -> None:
    if isinstance(value, list):
        if len(value) >= 2 and value[0] == "get" and isinstance(value[1], str):
            out.add(value[1])
        for item in value[1:]:
            walk_gets(item, out)
    elif isinstance(value, dict):
        for item in value.values():
            walk_gets(item, out)


def collect_sprite_literals(value: Any, sprite_keys: set[str], out: set[str]) -> None:
    if isinstance(value, str):
        if value in sprite_keys:
            out.add(value)
    elif isinstance(value, list):
        for item in value[1:]:
            collect_sprite_literals(item, sprite_keys, out)
    elif isinstance(value, dict):
        for item in value.values():
            collect_sprite_literals(item, sprite_keys, out)


def finite_positive_domain(expr: Any, field: str) -> set[Any] | None:
    if not isinstance(expr, list) or not expr:
        return None
    head = expr[0]
    if head == "==" and len(expr) >= 3:
        left, right = expr[1], expr[2]
        if isinstance(left, list) and len(left) >= 2 and left[0] == "get" and left[1] == field:
            if isinstance(right, (str, int, float)):
                return {right}
        if isinstance(right, list) and len(right) >= 2 and right[0] == "get" and right[1] == field:
            if isinstance(left, (str, int, float)):
                return {left}
        return None
    if head == "match" and len(expr) >= 5:
        input_expr = expr[1]
        if not (isinstance(input_expr, list) and len(input_expr) >= 2 and input_expr[0] == "get" and input_expr[1] == field):
            return None
        arms = expr[2:]
        fallback = arms[-1]
        if fallback is not False:
            return None
        allowed: set[Any] = set()
        pairs = arms[:-1]
        for i in range(0, len(pairs), 2):
            if i + 1 >= len(pairs):
                break
            labels, result = pairs[i], pairs[i + 1]
            if result is not True:
                continue
            if isinstance(labels, list):
                for x in labels:
                    if isinstance(x, (str, int, float)):
                        allowed.add(x)
            elif isinstance(labels, (str, int, float)):
                allowed.add(labels)
        return allowed or None
    if head == "all":
        parts = [finite_positive_domain(x, field) for x in expr[1:]]
        parts = [x for x in parts if x is not None]
        if not parts:
            return None
        out = set(parts[0])
        for part in parts[1:]:
            out &= part
        return out or None
    if head == "any":
        parts = [finite_positive_domain(x, field) for x in expr[1:]]
        if not parts or any(x is None for x in parts):
            return None
        out: set[Any] = set()
        for part in parts:
            out |= part or set()
        return out or None
    # Negation and inequality provide exclusions, not a finite positive domain.
    if head in {"!", "!="}:
        return None
    return None


def eval_simple_resource(expr: Any, values: dict[str, Any]) -> Any:
    if not isinstance(expr, list):
        return expr
    if not expr:
        return expr
    head = expr[0]
    if head == "get" and len(expr) >= 2:
        return values.get(expr[1])
    if head == "to-string" and len(expr) >= 2:
        value = eval_simple_resource(expr[1], values)
        return None if value is None else str(value)
    if head == "concat":
        parts = []
        for item in expr[1:]:
            value = eval_simple_resource(item, values)
            if value is None:
                return None
            parts.append(str(value))
        return "".join(parts)
    return None

def concat_prefix(expr: Any) -> str | None:
    if not isinstance(expr, list) or not expr or expr[0] != "concat":
        return None
    parts = expr[1:]
    prefix = ""
    saw_dynamic = False
    for item in parts:
        if isinstance(item, str) and not saw_dynamic:
            prefix += item
        else:
            saw_dynamic = True
    return prefix if prefix and saw_dynamic else None


def direct_get_field(expr: Any) -> str | None:
    if isinstance(expr, list) and len(expr) >= 2 and expr[0] == "get" and isinstance(expr[1], str):
        return expr[1]
    return None


def resource_candidates(expr: Any, sprite_keys: set[str], filter_expr: Any) -> tuple[list[str], list[str]]:
    literals: set[str] = set()
    collect_sprite_literals(expr, sprite_keys, literals)
    gets: set[str] = set()
    walk_gets(expr, gets)

    field = direct_get_field(expr)
    if field:
        domain = finite_positive_domain(filter_expr, field)
        if domain:
            for value in domain:
                if isinstance(value, str) and value in sprite_keys:
                    literals.add(value)

    # For simple generated keys (e.g. concat(prefix, to-string(get(field)))),
    # only expand when the Style filter itself supplies a finite positive domain.
    if len(gets) == 1:
        dynamic_field = next(iter(gets))
        domain = finite_positive_domain(filter_expr, dynamic_field)
        if domain:
            for value in domain:
                resolved = eval_simple_resource(expr, {dynamic_field: value})
                if isinstance(resolved, str) and resolved in sprite_keys:
                    literals.add(resolved)

    return sorted(literals), sorted(gets)


def sprite_visual_hash(image: Image.Image, meta: dict[str, Any]) -> str:
    crop = image.crop((
        int(meta["x"]), int(meta["y"]),
        int(meta["x"]) + int(meta["width"]),
        int(meta["y"]) + int(meta["height"]),
    ))
    return hashlib.sha256(crop.tobytes()).hexdigest()


def extract(style_path: str, sprite_override: str | None, output: Path, catalog_name: str) -> dict[str, Any]:
    output.mkdir(parents=True, exist_ok=True)
    style = load_json(style_path)
    style_bytes = canonical(style).encode("utf-8")
    sprite_meta, sprite_image = load_sprite(resolve_sprite_base(style, sprite_override), output / "_sprite")
    sprite_keys = set(sprite_meta)

    groups: dict[str, dict[str, Any]] = {}
    objects: list[dict[str, Any]] = []
    unresolved: list[dict[str, Any]] = []

    for index, layer in enumerate(style.get("layers", [])):
        source_layer = layer.get("source-layer")
        group_seed = {
            "source": layer.get("source"),
            "source_layer": source_layer,
        }
        group_id = sid("GRP", group_seed)
        groups.setdefault(group_id, {
            "group_id": group_id,
            "provenance": group_seed,
            "style_layer_indexes": [],
        })["style_layer_indexes"].append(index)

        found_resource = False
        for section, key, role in VISUAL_RESOURCE_PROPS:
            expr = (layer.get(section) or {}).get(key)
            if expr is None:
                continue
            found_resource = True
            candidates, dynamic_fields = resource_candidates(expr, sprite_keys, layer.get("filter"))
            base = {
                "group_id": group_id,
                "geometry_role": layer.get("type"),
                "component_role": role,
                "style_index": index,
                "style_layer_id": layer.get("id"),
                "filter": layer.get("filter"),
                "resource_expression": expr,
                "dynamic_fields": dynamic_fields,
                "provenance": {
                    "source": layer.get("source"),
                    "source_layer": source_layer,
                    "style_layer_id": layer.get("id"),
                },
            }
            for resource_key in candidates:
                meta = sprite_meta[resource_key]
                obj = dict(base)
                obj["resource_key"] = resource_key
                obj["resource_meta"] = meta
                obj["visual_hash"] = sprite_visual_hash(sprite_image, meta) if sprite_image else None
                signature = {
                    "geometry_role": obj["geometry_role"],
                    "component_role": role,
                    "visual_hash": obj["visual_hash"],
                    "style_index": index,
                }
                obj["visual_object_id"] = sid("VO", signature)
                objects.append(obj)
            if not candidates:
                item = dict(base)
                item["reason"] = "dynamic_or_unresolved_resource"
                item["visual_object_id"] = sid("VO", {
                    "geometry_role": item["geometry_role"],
                    "component_role": role,
                    "expression": expr,
                    "style_index": index,
                })
                unresolved.append(item)

        layer_type = layer.get("type")
        paint = layer.get("paint") or {}
        vector_payload = {k: paint[k] for k in VECTOR_PROPS.get(layer_type, ()) if k in paint}
        text_expr = (layer.get("layout") or {}).get("text-field")
        if vector_payload or text_expr is not None or not found_resource:
            payload = {
                "group_id": group_id,
                "geometry_role": layer_type,
                "component_role": "vector-or-text",
                "style_index": index,
                "style_layer_id": layer.get("id"),
                "filter": layer.get("filter"),
                "visual_properties": vector_payload,
                "text_expression": text_expr,
                "provenance": {
                    "source": layer.get("source"),
                    "source_layer": source_layer,
                    "style_layer_id": layer.get("id"),
                },
            }
            payload["visual_object_id"] = sid("VO", {
                "geometry_role": layer_type,
                "component_role": "vector-or-text",
                "visual_properties": vector_payload,
                "text_expression": text_expr,
                "style_index": index,
            })
            objects.append(payload)

    catalog = {
        "schema": "style-visual-object/v1",
        "catalog_name": catalog_name,
        "principle": "Visual Object extraction is provider-neutral; provenance names are not semantic classification.",
        "style": {
            "name": style.get("name"),
            "version": style.get("version"),
            "sha256": hashlib.sha256(style_bytes).hexdigest(),
            "source": style_path,
            "sprite": style.get("sprite"),
            "glyphs": style.get("glyphs"),
            "layer_count": len(style.get("layers", [])),
        },
        "sprite": {
            "resource_count": len(sprite_meta),
        },
        "groups": list(groups.values()),
        "visual_objects": objects,
        "unresolved_visual_objects": unresolved,
        "summary": {
            "group_count": len(groups),
            "visual_object_count": len(objects),
            "unresolved_count": len(unresolved),
            "sprite_backed_count": sum(1 for x in objects if x.get("resource_key")),
        },
    }
    (output / "anonymous-catalog.json").write_text(json.dumps(catalog, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return catalog


def freeze_standard(anonymous_path: Path, output: Path) -> dict[str, Any]:
    catalog = json.loads(anonymous_path.read_text(encoding="utf-8"))
    standard = json.loads(json.dumps(catalog))
    standard["schema"] = "style-visual-object-standard/v1"
    standard["catalog_status"] = "FROZEN"
    standard["standard_rule"] = "Layer name is attached only after anonymous Visual Object generation."
    for obj in standard["visual_objects"]:
        obj["standard_layer_name"] = obj.get("provenance", {}).get("source_layer")
    for obj in standard["unresolved_visual_objects"]:
        obj["standard_layer_name"] = obj.get("provenance", {}).get("source_layer")
    output.mkdir(parents=True, exist_ok=True)
    (output / "standard-catalog.json").write_text(json.dumps(standard, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return standard


def main() -> None:
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("extract")
    p.add_argument("--style", required=True)
    p.add_argument("--sprite")
    p.add_argument("--output", required=True)
    p.add_argument("--name", required=True)
    f = sub.add_parser("freeze-standard")
    f.add_argument("--anonymous", required=True)
    f.add_argument("--output", required=True)
    args = ap.parse_args()
    if args.cmd == "extract":
        result = extract(args.style, args.sprite, Path(args.output), args.name)
    else:
        result = freeze_standard(Path(args.anonymous), Path(args.output))
    print(json.dumps(result["summary"], ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()