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, Iterable, 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,
)


_ARCHIVE_READ_CHUNK_SIZE = 1024 * 1024
_WINDOWS_DRIVE_RE = re.compile(r"^[A-Za-z]:")
ARCHIVE_UPLOAD_MAX_ENTRIES = 10_000
# 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",
)


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:
    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)
        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)
    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,
    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,
) -> list[UploadedStoredFile]:
    inspection = inspect_archive(
        archive_data,
        filename=filename,
        password=password,
        max_entries=max_entries,
        max_expanded_bytes=max_expanded_bytes,
        max_expansion_ratio=max_expansion_ratio,
    )
    if inspection.requires_password and not inspection.password_verified:
        raise ArchivePasswordError("Archive password is required")
    selected_files = _selected_file_paths(inspection.entries, selected_paths)
    if not selected_files:
        raise FileStorageError("Select at least one archive file to import")
    actual_total_limit = min(
        max_expanded_bytes,
        inspection.compressed_size_bytes * max_expansion_ratio,
    )
    if inspection.archive_format == "zip":
        members = _read_selected_zip_members(
            archive_data,
            selected_files=selected_files,
            password=password,
            max_file_bytes=max_file_bytes,
            max_total_bytes=actual_total_limit,
        )
    else:
        members = _read_selected_tar_members(
            archive_data,
            selected_files=selected_files,
            max_file_bytes=max_file_bytes,
            max_total_bytes=actual_total_limit,
        )
    return _store_archive_members(
        session,
        members=members,
        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,
    )


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,
    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,
        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],
) -> list[ArchiveEntry]:
    try:
        with _open_tar(archive_data) as archive:
            entries: list[ArchiveEntry] = []
            for member in archive.getmembers():
                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"
                    )
                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(
            "Archive expansion ratio exceeds "
            f"{max_expansion_ratio}:1"
        )


def _with_derived_directories(
    entries: list[ArchiveEntry],
) -> 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"
                )
            by_path.setdefault(
                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,
) -> list[tuple[str, bytes]]:
    result: list[tuple[str, bytes]] = []
    total = 0
    try:
        with pyzipper.AESZipFile(_archive_source(archive_data)) as archive:
            for info in archive.infolist():
                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,
                        )
                except (RuntimeError, ValueError, zipfile.BadZipFile) as exc:
                    if info.flag_bits & 0x1:
                        raise ArchivePasswordError(
                            "Archive password is incorrect"
                        ) from exc
                    raise
                result.append((path, data))
    except (FileStorageError, ArchivePasswordError):
        raise
    except (OSError, ValueError, zipfile.BadZipFile) as exc:
        raise FileStorageError("ZIP extraction failed") from exc
    return result


def _read_selected_tar_members(
    archive_data: bytes | str | PathLike[str],
    *,
    selected_files: set[str],
    max_file_bytes: int,
    max_total_bytes: int,
) -> list[tuple[str, bytes]]:
    result: list[tuple[str, bytes]] = []
    total = 0
    try:
        with _open_tar(archive_data) as archive:
            for member in archive.getmembers():
                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,
                    )
                result.append((path, data))
    except FileStorageError:
        raise
    except (OSError, tarfile.TarError) as exc:
        raise FileStorageError("TAR extraction failed") from exc
    return result


def _read_member(
    source: BinaryIO,
    *,
    path: str,
    max_file_bytes: int,
    max_total_bytes: int,
    current_total: int,
) -> 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)
    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,
) -> list[UploadedStoredFile]:
    uploaded: list[UploadedStoredFile] = []
    base_folder = normalize_folder(folder)
    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,
            )
        )
    return uploaded


def _safe_member_path(value: str) -> str:
    raw = str(value or "").replace("\\", "/").strip()
    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:
        return normalize_logical_path(raw.rstrip("/"))
    except ValueError as exc:
        raise FileStorageError(f"Unsafe archive member path {value!r}") from exc


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