"""
.osmoimg File Format — Registered file format for OsmoImage containers.

The .osmoimg format is a self-contained, compressed, optionally encrypted
model image that can be run by anyone with Model Osmosis installed.

File Structure:
    [8 bytes]  Magic number: b'OSMOIMG\\x00'
    [4 bytes]  Format version (uint32 LE): 1, 2, or 3
    [4 bytes]  Header length (uint32 LE)
    [N bytes]  Header JSON (runtime metadata, compressed with zlib)
    [...]      Payload bytes (zip for v1, codec-defined for v2)

The magic bytes let the OS and tools instantly identify .osmoimg files
without reading the full contents.
"""

from __future__ import annotations

import json
import os
import platform
import struct
import sys
import zlib
import zipfile
import tempfile
from pathlib import Path
from typing import Optional

# ─── Format Constants ──────────────────────────────────────────────────────

OSMOIMG_MAGIC = b"OSMOIMG\x00"       # 8-byte magic number
OSMOIMG_VERSION = 3                   # Max format version supported by this runtime
OSMOIMG_EXTENSION = ".osmoimg"        # File extension
OSMOIMG_MIME_TYPE = "application/x-osmoimg"
OSMOIMG_DESCRIPTION = "Model Osmosis Image"
OSMOIMG_PROG_ID = "ModelOsmosis.OsmoImage"


# ─── Header Read/Write ────────────────────────────────────────────────────

def write_osmoimg_header(
    file_handle,
    metadata: dict,
    version: int = OSMOIMG_VERSION,
) -> int:
    """Write the .osmoimg header (magic + version + metadata).

    Args:
        file_handle: Open file in binary write mode.
        metadata: Runtime metadata dict.

    Returns:
        Number of bytes written.
    """
    header_json = json.dumps(metadata, sort_keys=True, separators=(",", ":")).encode("utf-8")
    header_compressed = zlib.compress(header_json, level=9)

    # Magic (8) + Version (4) + Header length (4) + Header data
    file_handle.write(OSMOIMG_MAGIC)
    file_handle.write(struct.pack("<I", version))
    file_handle.write(struct.pack("<I", len(header_compressed)))
    file_handle.write(header_compressed)

    return 8 + 4 + 4 + len(header_compressed)


def read_osmoimg_header(file_handle) -> dict:
    """Read and validate the .osmoimg header.

    Args:
        file_handle: Open file in binary read mode.

    Returns:
        Runtime metadata dict.

    Raises:
        ValueError: If file is not a valid .osmoimg.
    """
    magic = file_handle.read(8)
    if magic != OSMOIMG_MAGIC:
        raise ValueError(
            f"Not a valid .osmoimg file (expected magic {OSMOIMG_MAGIC!r}, "
            f"got {magic!r})"
        )

    version = struct.unpack("<I", file_handle.read(4))[0]
    if version > OSMOIMG_VERSION:
        raise ValueError(
            f"Unsupported .osmoimg version {version} "
            f"(max supported: {OSMOIMG_VERSION})"
        )

    header_len = struct.unpack("<I", file_handle.read(4))[0]
    header_compressed = file_handle.read(header_len)
    header_json = zlib.decompress(header_compressed)

    metadata = json.loads(header_json)
    if isinstance(metadata, dict):
        metadata.setdefault("__osmoimg_version", version)
        metadata.setdefault("__osmoimg_header_len", header_len)
    return metadata


def is_osmoimg(path: str | Path) -> bool:
    """Check if a file is a valid .osmoimg by reading magic bytes.

    Args:
        path: Path to file.

    Returns:
        True if the file has valid .osmoimg magic bytes.
    """
    try:
        with open(path, "rb") as f:
            magic = f.read(8)
            return magic == OSMOIMG_MAGIC
    except (OSError, IOError):
        return False


# ─── OS File Association Registration ──────────────────────────────────────

def register_file_format(verbose: bool = True, icon_path: Optional[str] = None) -> bool:
    """Register the .osmoimg file format with the operating system.

    On Windows: creates registry entries for file association.
    On macOS/Linux: creates XDG MIME type and desktop entries.

    Args:
        verbose: Print status messages.
        icon_path: Optional path to a custom icon (.ico) for file association.

    Returns:
        True if registration was successful.
    """
    system = platform.system()

    if system == "Windows":
        return _register_windows(verbose, icon_path=icon_path)
    elif system == "Linux":
        return _register_linux(verbose)
    elif system == "Darwin":
        return _register_macos(verbose)
    else:
        if verbose:
            print(f"  ⚠ File registration not supported on {system}")
        return False


def unregister_file_format(verbose: bool = True) -> bool:
    """Remove the .osmoimg file format registration from the OS.

    Args:
        verbose: Print status messages.

    Returns:
        True if unregistration was successful.
    """
    system = platform.system()

    if system == "Windows":
        return _unregister_windows(verbose)
    elif system == "Linux":
        return _unregister_linux(verbose)
    else:
        if verbose:
            print(f"  ⚠ File unregistration not supported on {system}")
        return False


# ─── Windows Registry ─────────────────────────────────────────────────────

def _register_windows(verbose: bool, icon_path: Optional[str] = None) -> bool:
    """Register .osmoimg on Windows via the registry."""
    try:
        import winreg

        osmosis_path = _find_osmosis_exe()
        icon_ref = _resolve_windows_icon_ref(osmosis_path=osmosis_path, icon_path=icon_path)

        # Register the extension → ProgID mapping
        # HKEY_CURRENT_USER\Software\Classes\.osmoimg
        with winreg.CreateKey(
            winreg.HKEY_CURRENT_USER,
            rf"Software\Classes\{OSMOIMG_EXTENSION}",
        ) as key:
            winreg.SetValue(key, "", winreg.REG_SZ, OSMOIMG_PROG_ID)
            winreg.SetValueEx(key, "Content Type", 0, winreg.REG_SZ, OSMOIMG_MIME_TYPE)

        # Register the ProgID with description and commands
        # HKEY_CURRENT_USER\Software\Classes\ModelOsmosis.OsmoImage
        prog_key_path = rf"Software\Classes\{OSMOIMG_PROG_ID}"

        with winreg.CreateKey(winreg.HKEY_CURRENT_USER, prog_key_path) as key:
            winreg.SetValue(key, "", winreg.REG_SZ, OSMOIMG_DESCRIPTION)

        # DefaultIcon
        with winreg.CreateKey(
            winreg.HKEY_CURRENT_USER,
            rf"{prog_key_path}\DefaultIcon",
        ) as key:
            winreg.SetValue(key, "", winreg.REG_SZ, icon_ref)

        # Shell\open\command — run with osmosis
        with winreg.CreateKey(
            winreg.HKEY_CURRENT_USER,
            rf"{prog_key_path}\shell\open\command",
        ) as key:
            winreg.SetValue(
                key, "", winreg.REG_SZ,
                f'"{osmosis_path}" run "%1"',
            )

        # Shell\inspect\command
        with winreg.CreateKey(
            winreg.HKEY_CURRENT_USER,
            rf"{prog_key_path}\shell\inspect",
        ) as key:
            winreg.SetValue(key, "", winreg.REG_SZ, "Inspect with Model Osmosis")

        with winreg.CreateKey(
            winreg.HKEY_CURRENT_USER,
            rf"{prog_key_path}\shell\inspect\command",
        ) as key:
            winreg.SetValue(
                key, "", winreg.REG_SZ,
                f'"{osmosis_path}" inspect "%1"',
            )

        # Shell\export\command
        with winreg.CreateKey(
            winreg.HKEY_CURRENT_USER,
            rf"{prog_key_path}\shell\export",
        ) as key:
            winreg.SetValue(key, "", winreg.REG_SZ, "Export with Model Osmosis")

        with winreg.CreateKey(
            winreg.HKEY_CURRENT_USER,
            rf"{prog_key_path}\shell\export\command",
        ) as key:
            winreg.SetValue(
                key, "", winreg.REG_SZ,
                f'"{osmosis_path}" export "%1"',
            )

        # Notify shell of change
        try:
            import ctypes
            SHCNE_ASSOCCHANGED = 0x08000000
            SHCNF_IDLIST = 0x0000
            ctypes.windll.shell32.SHChangeNotify(
                SHCNE_ASSOCCHANGED, SHCNF_IDLIST, None, None
            )
        except Exception:
            pass

        if verbose:
            print(f"  ✓ Registered .osmoimg file format on Windows")
            print(f"    Extension: {OSMOIMG_EXTENSION}")
            print(f"    ProgID:    {OSMOIMG_PROG_ID}")
            print(f"    Icon:      {icon_ref}")
            print(f"    Actions:   Open, Inspect, Export")
        return True

    except Exception as e:
        if verbose:
            print(f"  ✗ Failed to register on Windows: {e}")
        return False


def _unregister_windows(verbose: bool) -> bool:
    """Remove .osmoimg registration from Windows registry."""
    try:
        import winreg

        def _delete_key_tree(hive, path):
            try:
                with winreg.OpenKey(hive, path, 0, winreg.KEY_ALL_ACCESS) as key:
                    while True:
                        try:
                            subkey = winreg.EnumKey(key, 0)
                            _delete_key_tree(hive, rf"{path}\{subkey}")
                        except OSError:
                            break
                winreg.DeleteKey(hive, path)
            except FileNotFoundError:
                pass

        _delete_key_tree(
            winreg.HKEY_CURRENT_USER,
            rf"Software\Classes\{OSMOIMG_EXTENSION}",
        )
        _delete_key_tree(
            winreg.HKEY_CURRENT_USER,
            rf"Software\Classes\{OSMOIMG_PROG_ID}",
        )

        try:
            import ctypes
            ctypes.windll.shell32.SHChangeNotify(0x08000000, 0, None, None)
        except Exception:
            pass

        if verbose:
            print(f"  ✓ Unregistered .osmoimg file format")
        return True

    except Exception as e:
        if verbose:
            print(f"  ✗ Failed to unregister: {e}")
        return False


# ─── Linux XDG ─────────────────────────────────────────────────────────────

def _register_linux(verbose: bool) -> bool:
    """Register .osmoimg on Linux via XDG MIME types."""
    try:
        osmosis_path = _find_osmosis_exe()
        xdg_data = Path.home() / ".local" / "share"

        # Create MIME type XML
        mime_dir = xdg_data / "mime" / "packages"
        mime_dir.mkdir(parents=True, exist_ok=True)

        mime_xml = f"""<?xml version="1.0" encoding="UTF-8"?>
<mime-info xmlns="http://www.freedesktop.org/standards/shared-mime-info">
  <mime-type type="{OSMOIMG_MIME_TYPE}">
    <comment>{OSMOIMG_DESCRIPTION}</comment>
    <glob pattern="*{OSMOIMG_EXTENSION}"/>
    <magic priority="50">
      <match type="string" value="{OSMOIMG_MAGIC.decode('latin-1')}" offset="0"/>
    </magic>
  </mime-type>
</mime-info>"""

        (mime_dir / "osmoimg.xml").write_text(mime_xml)

        # Create .desktop file
        apps_dir = xdg_data / "applications"
        apps_dir.mkdir(parents=True, exist_ok=True)

        desktop_entry = f"""[Desktop Entry]
Type=Application
Name=Model Osmosis
Comment=Run OsmoImage model containers
Exec={osmosis_path} run %f
MimeType={OSMOIMG_MIME_TYPE}
Terminal=true
Categories=Science;Development;
"""
        (apps_dir / "osmosis.desktop").write_text(desktop_entry)

        # Update MIME database
        os.system("update-mime-database " + str(xdg_data / "mime") + " 2>/dev/null")
        os.system("update-desktop-database " + str(apps_dir) + " 2>/dev/null")

        if verbose:
            print(f"  ✓ Registered .osmoimg file format on Linux")
        return True

    except Exception as e:
        if verbose:
            print(f"  ✗ Failed to register on Linux: {e}")
        return False


def _unregister_linux(verbose: bool) -> bool:
    """Remove .osmoimg registration on Linux."""
    try:
        xdg_data = Path.home() / ".local" / "share"
        mime_file = xdg_data / "mime" / "packages" / "osmoimg.xml"
        desktop_file = xdg_data / "applications" / "osmosis.desktop"

        if mime_file.exists():
            mime_file.unlink()
        if desktop_file.exists():
            desktop_file.unlink()

        os.system("update-mime-database " + str(xdg_data / "mime") + " 2>/dev/null")

        if verbose:
            print(f"  ✓ Unregistered .osmoimg file format")
        return True

    except Exception as e:
        if verbose:
            print(f"  ✗ Failed to unregister: {e}")
        return False


# ─── macOS ─────────────────────────────────────────────────────────────────

def _register_macos(verbose: bool) -> bool:
    """Register .osmoimg on macOS (via launch services)."""
    if verbose:
        print(f"  ℹ macOS: .osmoimg files will be associated when opening one")
        print(f"    Run: osmosis run model.osmoimg")
    return True


# ─── Utilities ─────────────────────────────────────────────────────────────

def _find_osmosis_exe() -> str:
    """Find the installed osmosis executable path."""
    import shutil

    # Check if osmosis is on PATH
    osmosis = shutil.which("osmosis")
    if osmosis:
        return osmosis

    # Fallback: construct from Python executable
    python_dir = Path(sys.executable).parent
    if platform.system() == "Windows":
        scripts_dir = python_dir / "Scripts"
        exe = scripts_dir / "osmosis.exe"
        if exe.exists():
            return str(exe)
    else:
        bin_dir = python_dir
        exe = bin_dir / "osmosis"
        if exe.exists():
            return str(exe)

    # Final fallback
    return "osmosis"


def _resolve_windows_icon_ref(osmosis_path: str, icon_path: Optional[str]) -> str:
    """Resolve icon reference for Windows ProgID DefaultIcon."""
    # Explicit CLI path wins.
    if icon_path:
        icon_file = Path(icon_path).expanduser()
        if icon_file.exists():
            return f"{icon_file.resolve()},0"

    # Optional env override.
    env_icon = os.environ.get("OSMOIMG_ICON_PATH", "").strip()
    if env_icon:
        env_file = Path(env_icon).expanduser()
        if env_file.exists():
            return f"{env_file.resolve()},0"

    # Friendly defaults in current working directory.
    cwd = Path.cwd()
    for candidate in ("osmoimg icon.ico", "osmoimg_icon.ico", "osmoimg.ico"):
        icon_file = cwd / candidate
        if icon_file.exists():
            return f"{icon_file.resolve()},0"

    # Fallback: icon from executable.
    return f"{osmosis_path},0"
