#!/usr/bin/env python3
"""
Archive handling unit:
- Extract zip/7z
- Handle nested archives (only extraction & discovery, no log-type heuristics)
- Batch: treat multiple top-level archives as multiple devices
"""

from __future__ import annotations

import os
import shutil
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Tuple

from fs_utils import ensure_dir, find_prefixed_ancestor, run


ARCHIVE_SUFFIXES: Tuple[str, ...] = (".zip", ".7z")


@dataclass(frozen=True)
class ArchiveProcessor:
    max_depth: int = 4

    def _build_nested_extract_dir(self, archive: Path, extract_root: Path) -> Path:
        nested_root = ensure_dir(extract_root / "_nested_extracts")
        return nested_root / f"{archive.stem}_extracted"

    def extract_archive(self, archive: Path, dir_out: Path) -> None:
        ensure_dir(dir_out)
        suf = archive.suffix.lower()
        if suf == ".zip":
            with zipfile.ZipFile(archive, "r") as zf:
                zf.extractall(dir_out)
            return
        if suf == ".7z":
            if shutil.which("7z") is None:
                raise RuntimeError("未检测到 7z，可安装 p7zip 或提供 zip/已解压目录")
            run(["7z", "x", str(archive), f"-o{str(dir_out)}"])
            return
        raise RuntimeError(f"不支持的压缩格式: {archive}")

    def list_archives(self, root: Path) -> List[Path]:
        out: List[Path] = []
        for dirpath, _, filenames in os.walk(root):
            p = Path(dirpath)
            for fn in filenames:
                fpath = p / fn
                if fpath.suffix.lower() in ARCHIVE_SUFFIXES:
                    out.append(fpath)
        out.sort(key=lambda p: p.stat().st_mtime, reverse=True)
        return out

    def ensure_extracted_root(self, input_path: Path, work_root: Path) -> Path:
        """Return a directory that contains extracted contents to scan.

        Archive inputs under input* are extracted under the run work root. This
        keeps source inputs immutable while preserving source-adjacent run layout.
        """
        ensure_dir(work_root)
        if input_path.is_dir():
            return input_path
        if input_path.suffix.lower() in ARCHIVE_SUFFIXES:
            if find_prefixed_ancestor(input_path, "input") is not None:
                base = ensure_dir(work_root / "_extracts") / f"{input_path.stem}_extracted"
            else:
                base = input_path.parent / f"{input_path.stem}_extracted"
            self.extract_archive(input_path, base)
            return base
        # Treat single file input as "under its parent".
        return input_path.parent

    def extract_next_nested(self, root: Path, extract_root: Optional[Path] = None) -> Optional[Path]:
        """Extract the newest archive under root and return the extracted dir."""
        zips = self.list_archives(root)
        if not zips:
            return None
        next_zip = zips[0]
        if extract_root and find_prefixed_ancestor(root, "input") is not None:
            next_root = self._build_nested_extract_dir(next_zip, extract_root)
        else:
            next_root = next_zip.parent / (next_zip.stem + "_extracted")
        self.extract_archive(next_zip, next_root)
        return next_root

    def prepare_device_dirs(self, root_path: Path, device_root: Optional[Path] = None) -> List[Path]:
        """If root_path contains >=2 archives at top-level, treat them as devices."""
        archives = sorted(
            [p for p in root_path.iterdir() if p.is_file() and p.suffix.lower() in ARCHIVE_SUFFIXES]
        )
        if len(archives) < 2:
            return []

        device_dirs: List[Path] = []
        for idx, archive in enumerate(archives, start=1):
            if device_root is not None:
                device_dir = ensure_dir(device_root) / f"device{idx}_{archive.name}"
            else:
                device_dir = archive.parent / archive.stem
            ensure_dir(device_dir)
            marker = device_dir / ".source_archive"
            try:
                previous = marker.read_text(encoding="utf-8").strip() if marker.exists() else ""
            except Exception:
                previous = ""

            if previous != archive.name or not any(
                child.name != ".source_archive" for child in device_dir.iterdir()
            ):
                print(f"解压设备压缩包: {archive} -> {device_dir}")
                self.extract_archive(archive, device_dir)
                marker.write_text(archive.name + "\n", encoding="utf-8")
            else:
                print(f"复用已解压设备目录: {device_dir} <- {archive.name}")

            device_dirs.append(device_dir)
        return device_dirs