from __future__ import annotations

import mimetypes
import re
import stat
import tarfile
import zipfile
from dataclasses import dataclass
from io import BytesIO
from os import PathLike
from pathlib import Path, PurePosixPath
from typing import Any, BinaryIO, Callable, Iterable, Iterator, Literal

import pyzipper
from sqlalchemy.orm import Session

from govoplan_files.backend.db.models import FileAsset
from govoplan_files.backend.storage.backends import (
    StorageBackendError,
    get_storage_backend,
)
from govoplan_files.backend.storage.common import (
    FileConflictResolution,
    FileStorageError,
    UploadedStoredFile,
)
from govoplan_files.backend.storage.files import (
    create_file_asset,
    current_versions_and_blobs,
)
from govoplan_files.backend.storage.paths import (
    filename_from_path,
    normalize_folder,
    normalize_logical_path,
)
from govoplan_files.backend.storage.native_zip import (
    native_zip_library,
    read_native_zip_members,
)


_ARCHIVE_READ_CHUNK_SIZE = 1024 * 1024
_WINDOWS_DRIVE_RE = re.compile(r"^[A-Za-z]:")
ARCHIVE_UPLOAD_MAX_ENTRIES = 10_000
ARCHIVE_MAX_PATH_BYTES = 4096
ARCHIVE_MAX_PATH_DEPTH = 128
# Kept for callers and documentation using the former ZIP-specific name.
ZIP_UPLOAD_MAX_FILES = ARCHIVE_UPLOAD_MAX_ENTRIES
SUPPORTED_ARCHIVE_SUFFIXES = (
    ".tar.bz2",
    ".tar.gz",
    ".tar.xz",
    ".tbz2",
    ".tgz",
    ".txz",
    ".tar",
    ".zip",
)

ArchiveProgress = Callable[[str, int, int, int, int], None]


class ArchivePasswordError(FileStorageError):
    pass


@dataclass(frozen=True, slots=True)
class ArchiveEntry:
    path: str
    kind: Literal["file", "directory"]
    size_bytes: int
    compressed_size_bytes: int | None = None
    encrypted: bool = False


@dataclass(frozen=True, slots=True)
class ArchiveInspection:
    archive_format: str
    entries: tuple[ArchiveEntry, ...]
    file_count: int
    directory_count: int
    expanded_size_bytes: int
    compressed_size_bytes: int
    requires_password: bool
    password_verified: bool


def create_zip_file(
    session: Session, assets: Iterable[FileAsset], output_path: str | Path
) -> None:
    backend = get_storage_backend()
    asset_list = list(assets)
    version_blobs = current_versions_and_blobs(session, asset_list)
    with zipfile.ZipFile(
        output_path, mode="w", compression=zipfile.ZIP_DEFLATED
    ) as archive:
        for asset in asset_list:
            _version, blob = version_blobs[asset.id]
            info = zipfile.ZipInfo(asset.display_path)
            info.compress_type = zipfile.ZIP_DEFLATED
            with archive.open(info, "w") as member:
                try:
                    for chunk in backend.iter_bytes(blob.storage_key):
                        if chunk:
                            member.write(chunk)
                except StorageBackendError as exc:
                    raise FileStorageError(str(exc)) from exc


def archive_format_for_filename(filename: str) -> str:
    lowered = filename.strip().casefold()
    if lowered.endswith(".zip"):
        return "zip"
    if lowered.endswith((".tar.gz", ".tgz")):
        return "tar.gz"
    if lowered.endswith((".tar.bz2", ".tbz2")):
        return "tar.bz2"
    if lowered.endswith((".tar.xz", ".txz")):
        return "tar.xz"
    if lowered.endswith(".tar"):
        return "tar"
    raise FileStorageError(
        "Unsupported archive format. Use ZIP, TAR, TAR.GZ, TAR.BZ2, or TAR.XZ."
    )


def is_supported_archive_filename(filename: str) -> bool:
    try:
        archive_format_for_filename(filename)
    except FileStorageError:
        return False
    return True


def inspect_archive(
    archive_data: bytes | str | PathLike[str],
    *,
    filename: str,
    password: str | None = None,
    max_entries: int = ARCHIVE_UPLOAD_MAX_ENTRIES,
    max_expanded_bytes: int = 2 * 1024 * 1024 * 1024,
    max_expansion_ratio: int = 100,
) -> ArchiveInspection:
    from govoplan_files.backend.storage.archive_workers import inspect_archive_isolated
    return inspect_archive_isolated(
        archive_data, filename=filename, password=password, max_entries=max_entries,
        max_expanded_bytes=max_expanded_bytes, max_expansion_ratio=max_expansion_ratio,
    )


def _inspect_archive_content(
    archive_data: bytes | str | PathLike[str],
    *,
    filename: str,
    password: str | None = None,
    max_entries: int = ARCHIVE_UPLOAD_MAX_ENTRIES,
    max_expanded_bytes: int = 2 * 1024 * 1024 * 1024,
    max_expansion_ratio: int = 100,
) -> ArchiveInspection:
    archive_format = archive_format_for_filename(filename)
    compressed_size = _archive_size(archive_data)
    if archive_format == "zip":
        entries, requires_password, password_verified = _inspect_zip(
            archive_data,
            password=password,
        )
    else:
        entries = _inspect_tar(
            archive_data,
            compressed_size=compressed_size,
            max_entries=max_entries,
            max_expanded_bytes=max_expanded_bytes,
            max_expansion_ratio=max_expansion_ratio,
        )
        requires_password = False
        password_verified = True
    _validate_archive_limits(
        entries,
        compressed_size=compressed_size,
        max_entries=max_entries,
        max_expanded_bytes=max_expanded_bytes,
        max_expansion_ratio=max_expansion_ratio,
    )
    complete_entries = _with_derived_directories(entries, max_entries=max_entries)
    return ArchiveInspection(
        archive_format=archive_format,
        entries=tuple(complete_entries),
        file_count=sum(entry.kind == "file" for entry in complete_entries),
        directory_count=sum(entry.kind == "directory" for entry in complete_entries),
        expanded_size_bytes=sum(
            entry.size_bytes for entry in entries if entry.kind == "file"
        ),
        compressed_size_bytes=compressed_size,
        requires_password=requires_password,
        password_verified=password_verified,
    )


def extract_archive_upload(
    session: Session,
    *,
    tenant_id: str,
    owner_type: str,
    owner_id: str,
    user_id: str,
    archive_data: bytes | str | PathLike[str],
    filename: str,
    folder: str | None,
    campaign_id: str | None,
    selected_paths: Iterable[str] | None = None,
    password: str | None = None,
    conflict_strategy: str = "reject",
    conflict_resolutions: Iterable[FileConflictResolution] | None = None,
    metadata: dict[str, Any] | None = None,
    is_admin: bool = False,
    encryption_vault_id: str | None = None,
    max_entries: int = ARCHIVE_UPLOAD_MAX_ENTRIES,
    max_file_bytes: int = 50 * 1024 * 1024,
    max_expanded_bytes: int = 2 * 1024 * 1024 * 1024,
    max_expansion_ratio: int = 100,
    progress: ArchiveProgress | None = None,
) -> list[UploadedStoredFile]:
    from govoplan_files.backend.storage.archive_workers import extract_archive_isolated
    return extract_archive_isolated(
        session, archive_data=archive_data, filename=filename, password=password,
        selected_paths=selected_paths, max_entries=max_entries, max_file_bytes=max_file_bytes,
        max_expanded_bytes=max_expanded_bytes, max_expansion_ratio=max_expansion_ratio,
        progress=progress, store_options={
            "tenant_id": tenant_id, "owner_type": owner_type, "owner_id": owner_id,
            "user_id": user_id, "folder": folder, "campaign_id": campaign_id,
            "conflict_strategy": conflict_strategy, "conflict_resolutions": conflict_resolutions,
            "metadata": metadata, "is_admin": is_admin, "encryption_vault_id": encryption_vault_id,
        },
    )


def extract_zip_upload(
    session: Session,
    *,
    tenant_id: str,
    owner_type: str,
    owner_id: str,
    user_id: str,
    zip_data: bytes | str | PathLike[str],
    folder: str | None,
    campaign_id: str | None,
    conflict_strategy: str = "reject",
    conflict_resolutions: Iterable[FileConflictResolution] | None = None,
    metadata: dict[str, Any] | None = None,
    is_admin: bool = False,
    encryption_vault_id: str | None = None,
    max_files: int = ZIP_UPLOAD_MAX_FILES,
    max_file_bytes: int = 50 * 1024 * 1024,
    max_total_bytes: int = 250 * 1024 * 1024,
) -> list[UploadedStoredFile]:
    """Backward-compatible wrapper for the original immediate ZIP endpoint."""

    return extract_archive_upload(
        session,
        tenant_id=tenant_id,
        owner_type=owner_type,
        owner_id=owner_id,
        user_id=user_id,
        archive_data=zip_data,
        filename="archive.zip",
        folder=folder,
        campaign_id=campaign_id,
        conflict_strategy=conflict_strategy,
        conflict_resolutions=conflict_resolutions,
        metadata=metadata,
        is_admin=is_admin,
        encryption_vault_id=encryption_vault_id,
        max_entries=max_files,
        max_file_bytes=max_file_bytes,
        max_expanded_bytes=max_total_bytes,
    )


def _inspect_zip(
    archive_data: bytes | str | PathLike[str],
    *,
    password: str | None,
) -> tuple[list[ArchiveEntry], bool, bool]:
    try:
        with pyzipper.AESZipFile(_archive_source(archive_data)) as archive:
            infos = archive.infolist()
            entries = [_zip_entry(info) for info in infos]
            encrypted_info = next(
                (
                    info
                    for info in infos
                    if not info.is_dir() and bool(info.flag_bits & 0x1)
                ),
                None,
            )
            if encrypted_info is None:
                return entries, False, True
            if not password:
                return entries, True, False
            try:
                with archive.open(
                    encrypted_info, pwd=password.encode("utf-8")
                ) as member:
                    member.read(1)
            except (RuntimeError, ValueError, zipfile.BadZipFile) as exc:
                raise ArchivePasswordError("Archive password is incorrect") from exc
            return entries, True, True
    except ArchivePasswordError:
        raise
    except (OSError, ValueError, zipfile.BadZipFile) as exc:
        raise FileStorageError("Invalid ZIP upload") from exc


def _inspect_tar(
    archive_data: bytes | str | PathLike[str],
    *,
    compressed_size: int,
    max_entries: int,
    max_expanded_bytes: int,
    max_expansion_ratio: int,
) -> list[ArchiveEntry]:
    try:
        with _open_tar(archive_data) as archive:
            entries: list[ArchiveEntry] = []
            expanded_size = 0
            for member in archive:
                # Check each header before advancing across its payload. In a
                # compressed TAR getmembers() would inflate everything first.
                if len(entries) >= max_entries:
                    raise FileStorageError(
                        f"Archive contains too many entries (limit {max_entries})"
                    )
                path = _safe_member_path(member.name)
                if member.isdir():
                    entries.append(
                        ArchiveEntry(path=path, kind="directory", size_bytes=0)
                    )
                    continue
                if not member.isfile():
                    raise FileStorageError(
                        f"Archive member {member.name!r} is not a regular file or directory"
                    )
                if member.size < 0:
                    raise FileStorageError("Archive member has invalid size metadata")
                expanded_size += member.size
                if expanded_size > max_expanded_bytes:
                    raise FileStorageError(
                        "Archive is too large after extraction "
                        f"(limit {max_expanded_bytes} bytes)"
                    )
                if expanded_size and (
                    compressed_size <= 0
                    or expanded_size > compressed_size * max_expansion_ratio
                ):
                    raise FileStorageError(
                        f"Archive expansion ratio exceeds {max_expansion_ratio}:1"
                    )
                entries.append(
                    ArchiveEntry(
                        path=path,
                        kind="file",
                        size_bytes=max(0, int(member.size)),
                    )
                )
            return entries
    except FileStorageError:
        raise
    except (OSError, tarfile.TarError) as exc:
        raise FileStorageError("Invalid TAR upload") from exc


def _zip_entry(info: zipfile.ZipInfo) -> ArchiveEntry:
    path = _safe_member_path(info.filename)
    unix_mode = (info.external_attr >> 16) & 0xFFFF
    file_type = stat.S_IFMT(unix_mode)
    if file_type and not (stat.S_ISREG(unix_mode) or stat.S_ISDIR(unix_mode)):
        raise FileStorageError(
            f"Archive member {info.filename!r} is not a regular file or directory"
        )
    if info.file_size < 0 or info.compress_size < 0:
        raise FileStorageError(
            f"Archive member {info.filename!r} has invalid size metadata"
        )
    return ArchiveEntry(
        path=path,
        kind="directory" if info.is_dir() else "file",
        size_bytes=0 if info.is_dir() else int(info.file_size),
        compressed_size_bytes=(None if info.is_dir() else int(info.compress_size)),
        encrypted=bool(info.flag_bits & 0x1),
    )


def _validate_archive_limits(
    entries: list[ArchiveEntry],
    *,
    compressed_size: int,
    max_entries: int,
    max_expanded_bytes: int,
    max_expansion_ratio: int,
) -> None:
    if len(entries) > max_entries:
        raise FileStorageError(
            f"Archive contains too many entries (limit {max_entries})"
        )
    seen: dict[str, str] = {}
    expanded_size = 0
    for entry in entries:
        previous_kind = seen.get(entry.path)
        if previous_kind is not None:
            raise FileStorageError(f"Archive contains duplicate path {entry.path!r}")
        seen[entry.path] = entry.kind
        if entry.kind == "file":
            expanded_size += entry.size_bytes
            if expanded_size > max_expanded_bytes:
                raise FileStorageError(
                    "Archive is too large after extraction "
                    f"(limit {max_expanded_bytes} bytes)"
                )
    if expanded_size and (
        compressed_size <= 0 or expanded_size > compressed_size * max_expansion_ratio
    ):
        raise FileStorageError(
            f"Archive expansion ratio exceeds {max_expansion_ratio}:1"
        )


def _with_derived_directories(
    entries: list[ArchiveEntry],
    *,
    max_entries: int,
) -> list[ArchiveEntry]:
    by_path = {entry.path: entry for entry in entries}
    for entry in entries:
        parent = PurePosixPath(entry.path).parent
        while str(parent) not in {"", "."}:
            path = str(parent)
            existing = by_path.get(path)
            if existing and existing.kind != "directory":
                raise FileStorageError(
                    f"Archive path {path!r} is both a file and a directory"
                )
            if path not in by_path:
                if len(by_path) >= max_entries:
                    raise FileStorageError(
                        "Archive contains too many entries including parent "
                        f"directories (limit {max_entries})"
                    )
                by_path[path] = ArchiveEntry(path=path, kind="directory", size_bytes=0)
            parent = parent.parent
    return sorted(
        by_path.values(),
        key=lambda entry: (
            tuple(entry.path.casefold().split("/")),
            entry.kind != "directory",
        ),
    )


def _selected_file_paths(
    entries: tuple[ArchiveEntry, ...],
    selected_paths: Iterable[str] | None,
) -> set[str]:
    file_paths = {entry.path for entry in entries if entry.kind == "file"}
    if selected_paths is None:
        return file_paths
    known = {entry.path: entry for entry in entries}
    normalized = {_safe_member_path(path) for path in selected_paths}
    unknown = normalized - known.keys()
    if unknown:
        raise FileStorageError(
            f"Archive selection contains unknown path {sorted(unknown)[0]!r}"
        )
    selected_files: set[str] = set()
    for path in normalized:
        entry = known[path]
        if entry.kind == "file":
            selected_files.add(path)
            continue
        prefix = f"{path}/"
        selected_files.update(
            file_path for file_path in file_paths if file_path.startswith(prefix)
        )
    return selected_files


def _read_selected_zip_members(
    archive_data: bytes | str | PathLike[str],
    *,
    selected_files: set[str],
    password: str | None,
    max_file_bytes: int,
    max_total_bytes: int,
    progress: ArchiveProgress | None = None,
    total_bytes: int = 0,
) -> Iterator[tuple[str, bytes]]:
    total = 0
    completed = 0
    try:
        with pyzipper.AESZipFile(_archive_source(archive_data)) as archive:
            infos = archive.infolist()
            # Accelerate only classic ZIPCrypto using established public C
            # decoding APIs; AES and other formats retain their existing path.
            native_compatible = (
                bool(password)
                and "\x00" not in password
                and any(info.flag_bits & 1 for info in infos)
                and all(
                    not getattr(info, "wz_aes_version", None)
                    and info.compress_type in {zipfile.ZIP_STORED, zipfile.ZIP_DEFLATED}
                    and (info.filename.isascii() or bool(info.flag_bits & 0x800))
                    for info in infos
                )
            )
            native = native_zip_library() if native_compatible else None
            if native is not None:
                yield from read_native_zip_members(
                    native,
                    archive_data,
                    infos_by_path={
                        _safe_member_path(info.filename): info for info in infos
                    },
                    selected_files=selected_files,
                    normalize_name=_safe_member_path,
                    password=password,
                    max_file_bytes=max_file_bytes,
                    max_total_bytes=max_total_bytes,
                    progress=progress,
                    total_bytes=total_bytes,
                )
                return
            for info in infos:
                if info.is_dir():
                    continue
                path = _safe_member_path(info.filename)
                if path not in selected_files:
                    continue
                pwd = password.encode("utf-8") if password else None
                try:
                    with archive.open(info, pwd=pwd) as source:
                        data, total = _read_member(
                            source,
                            path=path,
                            max_file_bytes=max_file_bytes,
                            max_total_bytes=max_total_bytes,
                            current_total=total,
                            on_bytes=(
                                lambda count: progress(
                                    "extracting",
                                    completed,
                                    len(selected_files),
                                    count,
                                    total_bytes,
                                )
                            )
                            if progress
                            else None,
                        )
                except (RuntimeError, ValueError, zipfile.BadZipFile) as exc:
                    if info.flag_bits & 0x1:
                        raise ArchivePasswordError(
                            "Archive password is incorrect"
                        ) from exc
                    raise
                completed += 1
                if progress:
                    progress(
                        "extracting", completed, len(selected_files), total, total_bytes
                    )
                yield path, data
                del data
    except (FileStorageError, ArchivePasswordError):
        raise
    except (OSError, ValueError, zipfile.BadZipFile) as exc:
        raise FileStorageError("ZIP extraction failed") from exc


def _read_selected_tar_members(
    archive_data: bytes | str | PathLike[str],
    *,
    selected_files: set[str],
    max_file_bytes: int,
    max_total_bytes: int,
    progress: ArchiveProgress | None = None,
    total_bytes: int = 0,
) -> Iterator[tuple[str, bytes]]:
    total = 0
    completed = 0
    try:
        with _open_tar(archive_data) as archive:
            # Iterate as headers are read instead of inflating the entire TAR
            # with getmembers(), then seeking backwards to inflate it again.
            for member in archive:
                if not member.isfile():
                    continue
                path = _safe_member_path(member.name)
                if path not in selected_files:
                    continue
                source = archive.extractfile(member)
                if source is None:
                    raise FileStorageError(f"Archive member {path!r} could not be read")
                with source:
                    data, total = _read_member(
                        source,
                        path=path,
                        max_file_bytes=max_file_bytes,
                        max_total_bytes=max_total_bytes,
                        current_total=total,
                        on_bytes=(
                            lambda count: progress(
                                "extracting",
                                completed,
                                len(selected_files),
                                count,
                                total_bytes,
                            )
                        )
                        if progress
                        else None,
                    )
                completed += 1
                if progress:
                    progress(
                        "extracting", completed, len(selected_files), total, total_bytes
                    )
                yield path, data
                del data
    except FileStorageError:
        raise
    except (OSError, tarfile.TarError) as exc:
        raise FileStorageError("TAR extraction failed") from exc


def _read_member(
    source: BinaryIO,
    *,
    path: str,
    max_file_bytes: int,
    max_total_bytes: int,
    current_total: int,
    on_bytes: Callable[[int], None] | None = None,
) -> tuple[bytes, int]:
    parts: list[bytes] = []
    actual_size = 0
    while True:
        read_size = min(
            _ARCHIVE_READ_CHUNK_SIZE,
            max_file_bytes + 1 - actual_size,
        )
        chunk = source.read(read_size)
        if not chunk:
            break
        actual_size += len(chunk)
        if actual_size > max_file_bytes:
            raise FileStorageError(f"Archive member {path!r} exceeds per-file limit")
        if current_total + actual_size > max_total_bytes:
            raise FileStorageError("Archive is too large after extraction")
        parts.append(chunk)
        if on_bytes:
            on_bytes(current_total + actual_size)
    return b"".join(parts), current_total + actual_size


def _store_archive_members(
    session: Session,
    *,
    members: Iterable[tuple[str, bytes]],
    tenant_id: str,
    owner_type: str,
    owner_id: str,
    user_id: str,
    folder: str | None,
    campaign_id: str | None,
    conflict_strategy: str,
    conflict_resolutions: Iterable[FileConflictResolution] | None,
    metadata: dict[str, Any] | None,
    is_admin: bool,
    encryption_vault_id: str | None,
    progress: ArchiveProgress | None = None,
    total_files: int = 0,
    total_bytes: int = 0,
) -> list[UploadedStoredFile]:
    uploaded: list[UploadedStoredFile] = []
    base_folder = normalize_folder(folder)
    stored_bytes = 0
    for inner_path, data in members:
        target_path = f"{base_folder}/{inner_path}" if base_folder else inner_path
        uploaded.append(
            create_file_asset(
                session,
                tenant_id=tenant_id,
                owner_type=owner_type,
                owner_id=owner_id,
                user_id=user_id,
                filename=filename_from_path(inner_path),
                data=data,
                display_path=target_path,
                content_type=mimetypes.guess_type(inner_path)[0]
                or "application/octet-stream",
                metadata=metadata,
                campaign_id=campaign_id,
                conflict_strategy=conflict_strategy,
                conflict_resolutions=conflict_resolutions,
                is_admin=is_admin,
                encryption_vault_id=encryption_vault_id,
            )
        )
        stored_bytes += len(data)
        if progress:
            progress("storing", len(uploaded), total_files, stored_bytes, total_bytes)
        del data
    return uploaded


def _safe_member_path(value: str) -> str:
    raw = str(value or "").replace("\\", "/").strip()
    try:
        oversized = (
            len(raw) > ARCHIVE_MAX_PATH_BYTES
            or len(raw.encode("utf-8")) > ARCHIVE_MAX_PATH_BYTES
        )
    except UnicodeEncodeError as exc:
        raise FileStorageError("Archive member path is not valid Unicode") from exc
    if oversized:
        raise FileStorageError(
            f"Archive member path exceeds {ARCHIVE_MAX_PATH_BYTES} UTF-8 bytes"
        )
    if not raw or "\x00" in raw or raw.startswith("/") or _WINDOWS_DRIVE_RE.match(raw):
        raise FileStorageError(f"Unsafe archive member path {value!r}")
    if any(part == ".." for part in raw.split("/")):
        raise FileStorageError(f"Unsafe archive member path {value!r}")
    try:
        normalized = normalize_logical_path(raw.rstrip("/"))
    except ValueError as exc:
        raise FileStorageError(f"Unsafe archive member path {value!r}") from exc
    if normalized.count("/") + 1 > ARCHIVE_MAX_PATH_DEPTH:
        raise FileStorageError(
            f"Archive member path exceeds {ARCHIVE_MAX_PATH_DEPTH} components"
        )
    return normalized


def _archive_size(archive_data: bytes | str | PathLike[str]) -> int:
    if isinstance(archive_data, bytes):
        return len(archive_data)
    try:
        return Path(archive_data).stat().st_size
    except OSError as exc:
        raise FileStorageError("Archive upload could not be read") from exc


def _archive_source(
    archive_data: bytes | str | PathLike[str],
) -> BytesIO | str | PathLike[str]:
    return BytesIO(archive_data) if isinstance(archive_data, bytes) else archive_data


def _open_tar(
    archive_data: bytes | str | PathLike[str],
) -> tarfile.TarFile:
    if isinstance(archive_data, bytes):
        return tarfile.open(fileobj=BytesIO(archive_data), mode="r:*")
    try:
        return tarfile.open(name=archive_data, mode="r:*")
    except OSError as exc:
        raise FileStorageError("Archive upload could not be read") from exc