from __future__ import annotations

import argparse
import hashlib
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path

import defusedxml.ElementTree as ET

from dissect.target.tools.utils.cli import catch_sigpipe

TEMPLATE = """\
# This file is generated by target-build-magic
# Do not edit manually, for overrides see 'overrides.py'
#
# Generated on:         {timestamp}
# Source SHA1 hash:     {hash}
# Number of types:      {count}
# References:           http://www.freedesktop.org/standards/shared-mime-info
#                       https://gitlab.freedesktop.org/xdg/shared-mime-info/-/blob/master/data/freedesktop.org.xml.in

TYPES = {types}
"""

IGNORED_TYPES = (
    # Conflicts with ELF
    "application/x-iso9660-appimage",
    "application/vnd.appimage",
    "application/x-core",
    # Conflicts with application/zip
    "application/x-karbon",
    "application/x-kchart",
    "application/x-kformula",
    "application/x-kivio",
    "application/x-kontour",
    "application/x-kpresenter",
    "application/x-krita",
    "application/x-kspread",
    "application/x-kword",
    # Conflicts with TIFF
    "image/x-3ds",
    # Conflicts with MZ
    "application/x-msdownload",
    # Conflicts with TAR
    "text/x-devicetree-source",
    # Conflicts with DEB
    "application/x-archive",
    # Conflicts with XML
    "application/x-quicktime-media-link",
    # Conflicts with OGG
    "application/annodex",
    # Conflicts with ZLIB
    "text/x-todo-txt",
    # Conflicts with PNG
    "image/apng",
    # Too broad definition
    "video/mp2t",
)

IGNORED_SUB_TYPES = (
    "application/zip",
    "application/xml",
    "application/json",
    "application/annodex",
)


def parse_freedesktop_xml(path: Path) -> str:
    """Attempt to parse a FreeDesktop shared-mime-info document.

    Does not yet parse nested magic values.
    """

    ns = {"mime": "http://www.freedesktop.org/standards/shared-mime-info"}
    types = []

    content = path.read_bytes()
    hash = hashlib.sha1(content).hexdigest()

    magics_seen = set()

    for mime_type in ET.fromstring(content.decode()).findall("mime:mime-type", ns):
        mime_type_name = mime_type.get("type")

        if (sub_class := mime_type.find("mime:sub-class-of", ns)) is not None and sub_class.get(
            "type"
        ) in IGNORED_SUB_TYPES:
            continue

        if mime_type_name in IGNORED_TYPES:
            continue

        comment = mime_type.find("mime:comment", ns).text
        magics = []

        for magic_match in mime_type.findall("mime:magic/mime:match", ns):
            value = magic_match.get("value")
            type = magic_match.get("type")

            if not value:
                raise ValueError(f"Missing value in {magic_match!r}")

            if not type:
                raise ValueError(f"Missing type in {magic_match!r}")

            try:
                value = convert_value(type, value)
            except Exception as e:
                print(f"Failed to convert {type!r} with value {value!r} for type {mime_type_name!r}: {e!s}")
                continue

            # Calculate offset
            if ":" in (offset := magic_match.get("offset", "0")):
                offset = offset.split(":")[0]
            try:
                offset = int(offset)
            except ValueError:
                pass

            if (offset, value) in magics_seen:
                print(f"Duplicate magic value for type {mime_type_name!r}: offset={offset!r} value={value!r}")
                continue

            if len(value) < 2:
                print(
                    f"Magic value too short to be usable for type {mime_type_name!r}: offset={offset!r} value={value!r}"
                )
                continue

            magics_seen.add((offset, value))
            magics.append({"offset": offset, "value": value})

        patterns = [g.get("pattern") for g in mime_type.findall("mime:glob", ns)]

        if not patterns and not magics:
            print(f"No usable definitions for {mime_type_name!r}")
            continue

        types.append(
            {
                "type": mime_type_name,
                "name": comment,
                "pattern": patterns,
                "magic": magics,
            }
        )

    print(f"Parsed {len(types)} types")

    return TEMPLATE.format(timestamp=datetime.now(tz=timezone.utc), count=len(types), hash=hash, types=types)


def convert_value(type: str, value: str) -> bytes:
    """Convert a value to bytes."""

    # Regular string conversion to bytes
    # e.g. "PK\\003\\004"       -> b"PK\x03\x04"
    # e.g. "\\x00\\x01"         -> b"\x00\x01"
    # e.g. "\x89PNG\r\n\x1A\n"  -> b"\x89PNG\r\n\x1a\n"
    if type == "string":
        if "\\x" in value or "\\0" in value:
            return value.encode("latin1").decode("unicode-escape").encode("latin1")
        return value.encode("latin1")

    # Integer conversion
    # e.g. "0x9501"     -> b"\x95\x01"
    # e.g. "0xa1b2c3d4" -> b"\xa1\xb2\xc3\xd4"
    if type in ("byte", "little16", "little32", "big16", "big32", "host16", "host32"):
        return bytes.fromhex(value[2:] if value[:2] == "0x" else value)

    raise ValueError(f"Unknown type {type!r}")


@catch_sigpipe
def main() -> int:
    help_formatter = argparse.ArgumentDefaultsHelpFormatter
    parser = argparse.ArgumentParser(
        description="target-build-magic",
        fromfile_prefix_chars="@",
        formatter_class=help_formatter,
    )
    parser.add_argument("path", metavar="PATH", type=Path, help="path to freedesktop.org.xml.in file")
    parser.add_argument(
        "--output", "-o", type=Path, help="optional alternative output path for generated types document"
    )
    parser.add_argument("--format", action="store_true", default=False, help="format output using ruff")
    args = parser.parse_args()

    if not args.path.is_file():
        print(f"Provided file does not exist: {args.path}")
        return 1

    output = (args.output or Path(__file__).parents[1].joinpath("helpers/magic/mimetypes/freedesktop.py")).resolve()

    if not output.parent.is_dir():
        print(f"Output path parent directory {output.parent} does not exist")
        return 1

    try:
        content = parse_freedesktop_xml(args.path)
    except Exception as e:
        print(f"Exception while generating output: {e!s}")
        return 1

    try:
        output.write_text(content)
    except Exception as e:
        print(f"Exception while writing output to {output!s}: {e!s}")
        return 1

    if args.format:
        subprocess.run(["ruff", "format", output], capture_output=True)
        subprocess.run(["ruff", "check", "--fix", "--unsafe-fixes", output], capture_output=True)

    return 0


if __name__ == "__main__":
    sys.exit(main())
