"""Declarative structure decoding for firmware memory.

A schema is plain data (JSON-friendly), so the same description can come from
the repo, a Ghidra export, or a plugin:

    {
      "name": "gsp_ctrl",
      "fields": [
        {"name": "magic",   "offset": 0,  "fmt": "u32"},
        {"name": "state",   "offset": 8,  "fmt": "u32", "enum": {"0": "IDLE", "1": "RUN"}},
        {"name": "buf",     "offset": 64, "fmt": "bytes", "size": 256},
        {"name": "name",    "offset": 64, "fmt": "cstr",  "size": 32}
      ]
    }

Supported formats: u8/u16/u32/u64, s8/s16/s32/s64, f32/f64, bool, bytes,
cstr, ptr (u64 by default; override with ``size``).
"""
from __future__ import annotations

import struct as _struct
from dataclasses import dataclass, field as _field

from .errors import StructError

_INT_FORMATS = {
    "u8": ("B", 1),
    "u16": ("H", 2),
    "u32": ("I", 4),
    "u64": ("Q", 8),
    "s8": ("b", 1),
    "s16": ("h", 2),
    "s32": ("i", 4),
    "s64": ("q", 8),
    "bool": ("B", 1),
}
_FLOAT_FORMATS = {"f32": ("f", 4), "f64": ("d", 8)}


@dataclass
class Field:
    name: str
    offset: int
    fmt: str
    size: int = 0
    enum: dict | None = None
    note: str | None = None
    count: int = 1

    @property
    def nbytes(self) -> int:
        if self.fmt in _INT_FORMATS:
            return _INT_FORMATS[self.fmt][1]
        if self.fmt in _FLOAT_FORMATS:
            return _FLOAT_FORMATS[self.fmt][1]
        if self.fmt in ("bytes", "cstr"):
            return self.size
        if self.fmt == "ptr":
            return self.size or 8
        if self.fmt == "u64x":
            return 8
        raise StructError(f"field {self.name!r}: unknown format {self.fmt!r}")


@dataclass
class Schema:
    name: str
    fields: list[Field] = _field(default_factory=list)
    size: int = 0

    def __post_init__(self) -> None:
        computed = max((f.offset + f.nbytes for f in self.fields), default=0)
        if not self.size:
            self.size = computed
        elif self.size < computed:
            raise StructError(
                f"schema {self.name!r}: declared size {self.size:#x} < fields end {computed:#x}"
            )
        self.validate()

    def validate(self) -> None:
        seen: set[str] = set()
        for field in self.fields:
            if field.name in seen:
                raise StructError(f"schema {self.name!r}: duplicate field {field.name!r}")
            seen.add(field.name)
            if field.offset < 0:
                raise StructError(f"schema {self.name!r}: negative offset for {field.name!r}")
            if field.nbytes <= 0:
                raise StructError(f"schema {self.name!r}: zero-size field {field.name!r}")

    def decode(self, data: bytes, base: int = 0) -> dict:
        if len(data) < self.size:
            raise StructError(
                f"schema {self.name!r}: need {self.size} bytes, got {len(data)}"
            )
        out: dict = {}
        for field in self.fields:
            out[field.name] = _decode_field(field, data, base)
        return out


def _decode_field(field: Field, data: bytes, base: int):
    start = field.offset
    if field.fmt in _INT_FORMATS:
        code, width = _INT_FORMATS[field.fmt]
        value = _struct.unpack_from("<" + code, data, start)[0]
        if field.fmt == "bool":
            value = bool(value)
        return value
    if field.fmt in _FLOAT_FORMATS:
        code, _width = _FLOAT_FORMATS[field.fmt]
        return _struct.unpack_from("<" + code, data, start)[0]
    if field.fmt == "ptr":
        return int.from_bytes(data[start : start + field.nbytes], "little")
    if field.fmt == "bytes":
        return bytes(data[start : start + field.size])
    if field.fmt == "cstr":
        raw = data[start : start + field.size]
        return raw.split(b"\x00", 1)[0].decode("utf-8", "replace")
    raise StructError(f"unknown format {field.fmt!r}")


def decode(schema: Schema, data: bytes, base: int = 0) -> dict:
    return schema.decode(data, base=base)


def schema_from_json(payload: dict) -> Schema:
    if not isinstance(payload, dict) or "fields" not in payload:
        raise StructError("schema JSON must be an object with a 'fields' list")
    fields = []
    for item in payload["fields"]:
        try:
            fields.append(
                Field(
                    name=item["name"],
                    offset=int(item["offset"]),
                    fmt=item.get("fmt", "bytes"),
                    size=int(item.get("size", 0)),
                    enum=item.get("enum"),
                    note=item.get("note"),
                    count=int(item.get("count", 1)),
                )
            )
        except KeyError as exc:
            raise StructError(f"schema field missing key: {exc}") from exc
    return Schema(name=payload.get("name", "schema"), fields=fields, size=int(payload.get("size", 0)))


def schema_to_json(schema: Schema) -> dict:
    return {
        "name": schema.name,
        "size": schema.size,
        "fields": [
            {
                k: v
                for k, v in {
                    "name": f.name,
                    "offset": f.offset,
                    "fmt": f.fmt,
                    "size": f.size or None,
                    "enum": f.enum,
                    "note": f.note,
                }.items()
                if v is not None
            }
            for f in schema.fields
        ],
    }


def render(schema: Schema, values: dict, indent: str = "  ") -> str:
    lines = [f"{schema.name} ({{base}}+{schema.size:#x} bytes)"]
    width = max((len(f.name) for f in schema.fields), default=0)
    for field in schema.fields:
        value = values.get(field.name)
        lines.append(f"{indent}{field.name:<{width}} @+0x{field.offset:03x} = {_format_value(field, value)}")
        if field.note:
            lines.append(f"{indent}{'':<{width}}   ; {field.note}")
    return "\n".join(lines)


def _format_value(field: Field, value) -> str:
    if isinstance(value, bool):
        return str(value)
    if isinstance(value, float):
        return f"{value:g}"
    if isinstance(value, int):
        text = f"{value:#x}" if not field.enum else f"{value} ({field.enum.get(str(value), '?')})"
        if field.enum and field.fmt not in ("bool",):
            text = f"{value:#x} ({field.enum.get(str(value), field.enum.get(value, '?'))})"
        return text
    if isinstance(value, (bytes, bytearray)):
        head = bytes(value[:16]).hex(" ")
        suffix = " ..." if len(value) > 16 else ""
        return f"[{len(value)}B] {head}{suffix}"
    return repr(value)