"""pyaxml: Rust-backed Android AXML/ARSC parser with Python bindings."""

from __future__ import annotations

try:
    from lxml import etree as ET

    _LXML = True
except ImportError:
    import xml.etree.ElementTree as ET

    _LXML = False

from pyaxml._pyaxml import AXML as _AXML, ARSC as _ARSC
from pyaxml.proto import axml_pb2
from pyaxml import compute


# ── XML element helpers ───────────────────────────────────────────────────────


def _to_element(xml_str: str):
    """Parse an XML string into an Element, falling back to stdlib if lxml rejects it."""
    if _LXML:
        try:
            return ET.fromstring(xml_str.encode())
        except ET.XMLSyntaxError:
            import xml.etree.ElementTree as _stdlib_ET

            return _stdlib_ET.fromstring(xml_str)
    return ET.fromstring(xml_str)


def _from_element(element) -> str:
    """Serialize an Element to an XML string, accepting both lxml and stdlib types."""
    if _LXML:
        try:
            from lxml import etree as _lxml

            if isinstance(element, _lxml._Element):
                return _lxml.tostring(element, encoding="unicode")
        except (ImportError, AttributeError):
            pass
    import xml.etree.ElementTree as _stdlib_ET

    return _stdlib_ET.tostring(element, encoding="unicode")


# ── String pool encoding helpers (Android binary format) ─────────────────────


def _decode_string_block(data: bytes, is_utf8: bool) -> str:
    """Decode a StringBlock.data value from the proto representation.

    The Rust backend stores raw string content bytes (no length prefix, no null
    terminator) in `StringBlock.data`.  For UTF-8 pools the bytes are plain
    UTF-8; for UTF-16 pools they are raw UTF-16-LE.
    """
    if not data:
        return ""
    try:
        if is_utf8:
            return data.decode("utf-8", errors="replace")
        else:
            return data.decode("utf-16-le", errors="replace")
    except (UnicodeDecodeError, ValueError):
        return ""


def _encode_string_block(s: str, is_utf8: bool) -> bytes:
    """Encode a Python str into the proto StringBlock.data format (raw bytes, no prefix)."""
    if is_utf8:
        return s.encode("utf-8")
    else:
        return s.encode("utf-16-le")


def _switch_string_bytes(
    data: bytes, old_str: str, new_str: str, is_utf8: bool
) -> bytes:
    """Re-encode data as new_str if it currently decodes to old_str; otherwise unchanged."""
    if _decode_string_block(data, is_utf8) == old_str:
        return _encode_string_block(new_str, is_utf8)
    return data


def _decode_locale_tag(config) -> str:
    """Extract a locale tag string from an ARSCResTableConfig proto message."""
    try:
        raw = config.locale.to_bytes(4, "little")
        lang = raw[0:2].rstrip(b"\x00").decode("ascii", errors="replace")
        region = raw[2:4].rstrip(b"\x00").decode("ascii", errors="replace")
        if not lang:
            return "default"
        return f"{lang}-{region}" if region else lang
    except Exception:
        return "default"


# ── StringBlocks compatibility class ─────────────────────────────────────────


class StringBlocks:
    """Compatibility class for string pool manipulation via protobuf."""

    def __init__(self, proto=None):
        self._proto = proto

    @property
    def proto(self):
        """The underlying StringBlocks protobuf message."""
        return self._proto

    @proto.setter
    def proto(self, value):
        self._proto = value

    def switch(self, old_str: str, new_str: str) -> None:
        """Replace all occurrences of old_str with new_str in the string pool."""
        is_utf8 = bool(self._proto.hnd.flag & 256)
        for block in self._proto.stringblocks:
            block.data = _switch_string_bytes(block.data, old_str, new_str, is_utf8)

    def decode_str(self, idx: int) -> str:
        """Return the string at pool index idx."""
        is_utf8 = bool(self._proto.hnd.flag & 256)
        return _decode_string_block(self._proto.stringblocks[idx].data, is_utf8)


# ── StringBlocksProxy proxy for axml.stringblocks ──────────────────────────


class StringBlocksProxy:
    """Proxy returned by AXML.stringblocks; allows .proto get/set to mutate the AXML."""

    def __init__(self, owner: "AXML"):
        self._owner = owner

    @property
    def proto(self):
        """The StringBlocks sub-message from the AXML proto."""
        return self._owner.to_proto().stringblocks

    @proto.setter
    def proto(self, new_sb) -> None:
        """Apply new string pool without recalculating chunk sizes (corrupted output).

        Call ``axml.compute()`` before ``pack()`` to produce a correctly-sized binary.
        """
        new_strings = [bytes(b.data) for b in new_sb.stringblocks]
        self._owner._inner.apply_stringblocks_no_compute(new_strings)


# ── AXML ──────────────────────────────────────────────────────────────────────


class AXML:
    """Android Binary XML document wrapper."""

    def __init__(self):
        self._inner = _AXML()

    @classmethod
    def from_axml(cls, data: bytes) -> "AXML":
        """Parse binary AXML data.

        Args:
            data: Binary AXML bytes (e.g., from AndroidManifest.xml).

        Returns:
            Parsed AXML document.

        Raises:
            ValueError: If parsing fails.
        """
        obj = cls.__new__(cls)
        obj._inner = _AXML.from_axml(data)
        return obj

    def __iter__(self):
        """Support tuple-unpacking: ``axml, remainder = AXML.from_axml(data)``."""
        return iter((self, b""))

    def pack(self) -> bytes:
        """Serialize back to binary AXML format."""
        return self._inner.pack()

    def to_xml(self):
        """Decode to an XML Element tree (lxml or stdlib ET).

        Returns:
            Root Element of the decoded XML.

        Raises:
            ValueError: If decoding fails.
        """
        return _to_element(self._inner.to_xml())

    def from_xml(self, element) -> None:
        """Recompile the AXML from an XML Element or XML string.

        Args:
            element: lxml.etree.Element, xml.etree.ElementTree.Element, or XML string.

        Raises:
            ValueError: If recompilation fails.
        """
        if isinstance(element, str):
            self._inner.from_xml(element)
        else:
            self._inner.from_xml(_from_element(element))

    def compute(self) -> None:
        """Recalculate all chunk-size header fields (string pool + outer total size).

        Without calling this, ``pack()`` uses whatever sizes are currently stored
        (original parsed values, values set manually, or stale values from
        ``stringblocks.proto = …``), useful for producing intentionally malformed
        files.  Call ``compute()`` to restore correct values.
        """
        self._inner.compute()

    @property
    def total_size(self) -> int:
        """Stored AXML total-file-size header field (bytes 4..8).

        ``pack()`` writes this value verbatim.  Set it to an arbitrary value to
        produce a binary with a wrong size field without changing any content.
        Call ``compute()`` to reset it to the correct value.
        """
        return self._inner.get_total_size()

    @total_size.setter
    def total_size(self, value: int) -> None:
        self._inner.set_total_size(value)

    def string_count(self) -> int:
        """Return the number of strings in the string pool."""
        return self._inner.string_count()

    def get_string(self, idx: int) -> str:
        """Get a string by index from the string pool.

        Args:
            idx: String pool index.

        Returns:
            The decoded string.

        Raises:
            ValueError: If index is out of bounds.
        """
        return self._inner.get_string(idx)

    @property
    def stringblocks(self) -> StringBlocksProxy:
        """Proxy for accessing and modifying the string pool via proto."""
        return StringBlocksProxy(self)

    def to_proto_bytes(self) -> bytes:
        """Serialize to binary proto bytes (proto encoding, not Android binary XML)."""
        return bytes(self._inner.to_proto_bytes())

    @classmethod
    def from_proto_bytes(cls, data: bytes) -> "AXML":
        """Deserialize from binary proto bytes produced by :meth:`to_proto_bytes`.

        Args:
            data: Binary proto bytes.

        Returns:
            Reconstructed AXML document.
        """
        obj = cls.__new__(cls)
        obj._inner = _AXML.from_proto_bytes(data)
        return obj

    def to_proto(self) -> axml_pb2.AXML:
        """Serialize to a google.protobuf AXML message object."""
        msg = axml_pb2.AXML()
        msg.ParseFromString(self._inner.to_proto_bytes())
        return msg

    @classmethod
    def from_proto(cls, proto: axml_pb2.AXML) -> "AXML":
        """Deserialize from a google.protobuf AXML message object.

        Args:
            proto: axml_pb2.AXML message.

        Returns:
            Reconstructed AXML document.
        """
        obj = cls.__new__(cls)
        obj._inner = _AXML.from_proto_bytes(proto.SerializeToString())
        return obj

    def export_proto_file(self, path: str) -> None:
        """Write the binary proto representation to *path*.

        The file can be reloaded with :meth:`from_proto_file`.
        """
        self._inner.export_proto_file(path)

    @classmethod
    def from_proto_file(cls, path: str) -> "AXML":
        """Load an AXML from a binary proto file written by :meth:`export_proto_file`.

        Args:
            path: Path to the ``.pb`` file.

        Returns:
            Reconstructed AXML document.
        """
        obj = cls.__new__(cls)
        obj._inner = _AXML.from_proto_file(path)
        return obj

    def export_stringblocks_proto_file(self, path: str) -> None:
        """Write the string pool as a binary proto file to *path*."""
        self._inner.export_stringblocks_proto_file(path)

    def import_stringblocks_proto_file(self, path: str) -> None:
        """Replace the string pool from a binary proto file.

        The pool is marked dirty; call :meth:`compute` before :meth:`pack`
        to produce a correctly-sized binary.
        """
        self._inner.import_stringblocks_proto_file(path)


# ── ARSC ──────────────────────────────────────────────────────────────────────


class ARSC:
    """Android Resource Table container."""

    def __init__(self):
        self._inner = _ARSC.__new__(_ARSC)

    @classmethod
    def from_axml(cls, data: bytes) -> "ARSC":
        """Parse binary ARSC data (resources.arsc).

        Args:
            data: Binary ARSC bytes.

        Returns:
            Parsed ARSC resource table.

        Raises:
            ValueError: If parsing fails.
        """
        obj = cls.__new__(cls)
        obj._inner = _ARSC.from_axml(data)
        return obj

    def __iter__(self):
        """Support tuple-unpacking: ``arsc, remainder = ARSC.from_axml(data)``."""
        return iter((self, b""))

    def list_packages(self, language: str | None = None) -> str:
        """List resource entries grouped by locale.

        Args:
            language: Optional locale filter (e.g., "en", "fr-FR", "default").
                      Pass None to include all locales.

        Returns:
            XML string with <resources lang="..."> sections.
        """
        return self._inner.list_packages(language)

    def package_count(self) -> int:
        """Return the number of packages in the resource table."""
        return self._inner.package_count()

    def compute(self, recursive: bool = True) -> None:
        """Recalculate all chunk-size header fields and store them.

        ``recursive=True`` (default): recomputes string pools and all package
        chunk sizes, then updates the global total-size and package-count fields.
        ``recursive=False``: recomputes only the global string pool and the
        top-level size fields; package string pools keep their current raw bytes.
        """
        self._inner.compute(recursive)

    @property
    def total_size(self) -> int:
        """Stored ARSC total-file-size header field (bytes 4..8).

        ``pack()`` writes this value verbatim.  Modify to produce a binary with
        an intentionally wrong file size.  Call ``compute()`` to restore the
        correct value.
        """
        return self._inner.get_total_size()

    @total_size.setter
    def total_size(self, value: int) -> None:
        self._inner.set_total_size(value)

    @property
    def header_package_count(self) -> int:
        """Stored package-count header field (bytes 8..12).

        ``pack()`` writes this value verbatim.  Modify to produce a binary where
        the declared package count differs from actual packages.
        Call ``compute()`` to restore the correct value.
        """
        return self._inner.get_header_package_count()

    @header_package_count.setter
    def header_package_count(self, value: int) -> None:
        self._inner.set_header_package_count(value)

    def get_pkg_chunk_size(self, idx: int) -> int:
        """Return the stored chunk-size header field for package at *idx*."""
        return self._inner.get_pkg_chunk_size(idx)

    def set_pkg_chunk_size(self, idx: int, size: int) -> None:
        """Set the chunk-size header field for package at *idx*.

        ``pack()`` writes this value verbatim for that package's chunk header.
        Call ``compute()`` with ``recursive=True`` to restore the correct value.
        """
        self._inner.set_pkg_chunk_size(idx, size)

    @property
    def proto(self) -> axml_pb2.ARSC:
        """The full ARSC protobuf message."""
        return self.to_proto()

    def get_packages(self) -> list:
        """Return a list of AXMLResTablePackage proto messages."""
        return list(self.to_proto().restablespackage)

    def get_id_public(
        self,
        pkg,
        type_name: str,
        name: str,
        language: str | None = None,
    ) -> "tuple[int, int] | None":
        """Find the resource ID and value data for (type_name, name).

        Args:
            pkg: An AXMLResTablePackage proto (as returned by get_packages()).
            type_name: Resource type string (e.g. "string", "xml", "layout").
            name: Resource entry name (e.g. "network_security_config").
            language: Optional locale filter (e.g. "en", "fr-FR", "default").
                      When None, all locales are searched and the first match
                      is returned.

        Returns:
            ``(resource_id, value_data)`` where ``value_data`` is the entry's
            Res_value data field (global string-pool index for TYPE_STRING
            resources).

        Raises:
            KeyError: If the type or resource is not found.
        """
        key_flag = bool(pkg.key_sp_string.hnd.flag & 256)
        type_flag = bool(pkg.type_sp_string.hnd.flag & 256)
        key_strs = [
            _decode_string_block(b.data, key_flag)
            for b in pkg.key_sp_string.stringblocks
        ]
        type_strs = [
            _decode_string_block(b.data, type_flag)
            for b in pkg.type_sp_string.stringblocks
        ]
        try:
            type_idx = type_strs.index(type_name)
        except ValueError:
            raise KeyError(f"Type '{type_name}' not found in package")

        for restype in pkg.restypes:
            typetype = restype.typetype
            if not typetype:
                continue
            if typetype.id - 1 != type_idx:
                continue
            type_id = typetype.id
            if language is not None:
                entry_lang = _decode_locale_tag(typetype.config)
                if entry_lang != language:
                    continue
            for entry_idx, table in enumerate(typetype.tables):
                if not table.present:
                    continue
                key_idx = table.index
                if key_idx < len(key_strs) and key_strs[key_idx] == name:
                    res_id = (pkg.id << 24) | (type_id << 16) | entry_idx
                    value_data = table.key.data if table.key else 0
                    return res_id, value_data

        suffix = f" for language='{language}'" if language is not None else ""
        raise KeyError(f"Resource '{type_name}/{name}' not found{suffix}")

    def pack(self) -> bytes:
        """Serialize back to binary `resources.arsc` format."""
        return bytes(self._inner.pack())

    def add_id_public(self, pkg, type_name: str, name: str, path: str) -> int:
        """Add a new resource entry (type_name / name → path file string).

        Finds or creates `type_name` in the first package's type string pool,
        adds `name` to the key string pool, adds `path` to the global string
        pool, and appends TypeSpec + TypeType chunks for the default locale.

        Args:
            pkg: Package proto (used only for API compatibility; the first
                 package in the ARSC is always modified).
            type_name: Resource type (e.g. ``"xml"``, ``"layout"``).
            name: Resource entry name (e.g. ``"network_security_config"``).
            path: File path value (e.g. ``"res/xml/network_security_config.xml"``).

        Returns:
            The new resource ID (0x7fTTEEEE).
        """
        return self._inner.add_resource(type_name, name, path)

    def to_proto_bytes(self) -> bytes:
        """Serialize to binary proto bytes (proto encoding, not Android binary ARSC format)."""
        return bytes(self._inner.to_proto_bytes())

    @classmethod
    def from_proto_bytes(cls, data: bytes) -> "ARSC":
        """Deserialize from binary proto bytes produced by :meth:`to_proto_bytes`.

        Args:
            data: Binary proto bytes.

        Returns:
            Reconstructed ARSC resource table.
        """
        obj = cls.__new__(cls)
        obj._inner = _ARSC.from_proto_bytes(data)
        return obj

    def to_proto(self) -> axml_pb2.ARSC:
        """Serialize to a google.protobuf ARSC message object."""
        msg = axml_pb2.ARSC()
        msg.ParseFromString(self._inner.to_proto_bytes())
        return msg

    @classmethod
    def from_proto(cls, proto: axml_pb2.ARSC) -> "ARSC":
        """Deserialize from a google.protobuf ARSC message object.

        Args:
            proto: axml_pb2.ARSC message.

        Returns:
            Reconstructed ARSC resource table.
        """
        obj = cls.__new__(cls)
        obj._inner = _ARSC.from_proto_bytes(proto.SerializeToString())
        return obj

    def export_proto_file(self, path: str) -> None:
        """Write the binary proto representation to *path*."""
        self._inner.export_proto_file(path)

    @classmethod
    def from_proto_file(cls, path: str) -> "ARSC":
        """Load an ARSC from a binary proto file written by :meth:`export_proto_file`.

        Args:
            path: Path to the ``.pb`` file.

        Returns:
            Reconstructed ARSC resource table.
        """
        obj = cls.__new__(cls)
        obj._inner = _ARSC.from_proto_file(path)
        return obj

    def export_pkg_proto_file(self, idx: int, path: str) -> None:
        """Write package *idx* as a binary proto file to *path*."""
        self._inner.export_pkg_proto_file(idx, path)

    def import_pkg_proto_file(self, idx: int, path: str) -> None:
        """Replace package *idx* from a binary proto file.

        Args:
            idx: Zero-based package index.
            path: Path to the ``.pb`` file.

        Raises:
            ValueError: If *idx* is out of range or the file cannot be decoded.
        """
        self._inner.import_pkg_proto_file(idx, path)


# ── AXMLGuess auto-detect AXML vs ARSC ─────────────────────────────────────


class AXMLGuess:
    """Auto-detect binary format (AXML or ARSC) and parse accordingly."""

    @classmethod
    def from_axml(cls, data: bytes) -> tuple:
        """Parse data as AXML, falling back to ARSC on failure.

        Args:
            data: Binary bytes (AXML or ARSC).

        Returns:
            Tuple of (AXML or ARSC instance, empty remainder bytes).
        """
        try:
            return AXML.from_axml(data), b""
        except Exception:
            return ARSC.from_axml(data), b""


__all__ = ["AXML", "ARSC", "AXMLGuess", "StringBlocks", "compute"]

try:
    from importlib.metadata import version as _pkg_version, PackageNotFoundError

    __version__ = _pkg_version("pyaxml")
except PackageNotFoundError:
    import re as _re, pathlib as _pl

    _m = _re.search(
        r'^version\s*=\s*"([^"]+)"',
        (_pl.Path(__file__).parents[2] / "Cargo.toml").read_text(),
        _re.M,
    )
    __version__ = _m.group(1) if _m else "unknown"
