"""Archive sniffing and hardened extraction (zip / 7z / rar)."""
from __future__ import annotations

import os
import shutil
import stat
import subprocess
import zipfile
from pathlib import Path

from .util import safe_relpath

MAX_UNPACKED = 32 * 1024 ** 3


class ArchiveError(Exception):
    pass


def sniff(path: Path) -> str:
    """Return 'zip', '7z', 'rar' based on magic bytes (extensions lie)."""
    with open(path, "rb") as fh:
        head = fh.read(8)
    if head[:4] in (b"PK\x03\x04", b"PK\x05\x06"):
        return "zip"
    if head[:6] == b"7z\xbc\xaf\x27\x1c":
        return "7z"
    if head[:4] == b"Rar!":
        return "rar"
    raise ArchiveError(f"{Path(path).name}: not a zip, 7z or rar archive")


def _validate_name(name: str) -> str:
    rel = safe_relpath(name)
    if rel is None:
        raise ArchiveError(f"unsafe path in archive: {name!r}")
    return rel


def _extract_zip(path: Path, dest: Path) -> None:
    try:
        zf = zipfile.ZipFile(path)
    except zipfile.BadZipFile as e:
        raise ArchiveError(f"corrupt zip: {e}") from e
    with zf:
        infos = zf.infolist()
        # validate everything before writing anything
        total = 0
        plan = []
        for info in infos:
            rel = _validate_name(info.filename)
            mode = info.external_attr >> 16
            if stat.S_ISLNK(mode):
                raise ArchiveError(f"symlink entries are not allowed: {info.filename!r}")
            total += info.file_size
            plan.append((info, rel))
        if total > MAX_UNPACKED:
            raise ArchiveError("archive unpacks to an unreasonable size")
        root = os.path.realpath(dest)
        for info, rel in plan:
            target = os.path.join(root, rel)
            if os.path.commonpath([root, os.path.realpath(target)]) != root:
                raise ArchiveError(f"unsafe path in archive: {info.filename!r}")
            if info.is_dir():
                os.makedirs(target, exist_ok=True)
                continue
            os.makedirs(os.path.dirname(target), exist_ok=True)
            with zf.open(info) as src, open(target, "wb") as out:
                shutil.copyfileobj(src, out)


def _tool(*names: str) -> str | None:
    for n in names:
        p = shutil.which(n)
        if p:
            return p
    return None


def _run(cmd: list[str]) -> str:
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, errors="replace", timeout=3600)
    except FileNotFoundError as e:
        raise ArchiveError(f"required tool missing: {cmd[0]}") from e
    if r.returncode != 0:
        raise ArchiveError((r.stderr or r.stdout).strip().splitlines()[-1:][0] if (r.stderr or r.stdout).strip()
                           else f"{cmd[0]} failed")
    return r.stdout


def _list_7z(tool: str, path: Path) -> list[str]:
    out = _run([tool, "l", "-slt", "-ba", "--", str(path)])
    return [line[7:] for line in out.splitlines() if line.startswith("Path = ")]


def _forbid_symlinks(dest: Path) -> None:
    for dirpath, dirnames, filenames in os.walk(dest):
        for n in dirnames + filenames:
            p = os.path.join(dirpath, n)
            if os.path.islink(p):
                raise ArchiveError(f"symlink found in extracted archive: {n!r}")


def _extract_7z(path: Path, dest: Path) -> None:
    tool = _tool("7z", "7zz", "7za", "7zr")
    if not tool:
        raise ArchiveError("7z is not installed")
    for n in _list_7z(tool, path):
        _validate_name(n)
    _run([tool, "x", "-y", f"-o{dest}", "--", str(path)])
    _forbid_symlinks(dest)


def _extract_rar(path: Path, dest: Path) -> None:
    unrar = _tool("unrar")
    if unrar:
        names = _run([unrar, "lb", "--", str(path)]).splitlines()
        for n in names:
            if n.strip():
                _validate_name(n)
        _run([unrar, "x", "-y", "-o+", "--", str(path), str(dest) + os.sep])
        _forbid_symlinks(dest)
    else:
        _extract_7z(path, dest)  # 7z can read rar when its codec is present


def extract(archive: Path, dest: Path) -> None:
    """Extract into an existing empty directory, rejecting traversal and symlinks."""
    archive, dest = Path(archive), Path(dest)
    dest.mkdir(parents=True, exist_ok=True)
    kind = sniff(archive)
    {"zip": _extract_zip, "7z": _extract_7z, "rar": _extract_rar}[kind](archive, dest)


def copy_tree(src: Path, dest: Path) -> None:
    """Import an already-extracted folder; symlinks are skipped."""
    for dirpath, dirnames, filenames in os.walk(src):
        rel = Path(dirpath).relative_to(src)
        (dest / rel).mkdir(parents=True, exist_ok=True)
        dirnames[:] = [d for d in dirnames if not os.path.islink(os.path.join(dirpath, d))]
        for f in filenames:
            s = os.path.join(dirpath, f)
            if not os.path.islink(s):
                shutil.copy2(s, dest / rel / f)