"""Voice export/import bundles (ZIP). See architecture/voice-bundle.md.

Export streams a ZIP (no temp files). Import treats the archive as hostile:
only fixed flat member names, capped sizes, hashes verified, nothing executed,
no path from the archive is ever used.
"""
from __future__ import annotations

import hashlib
import json
import re
import shutil
import tempfile
import zipfile
from pathlib import Path
from typing import Iterator

from core.provenance import CONSENT_EXTS, MAX_CONSENT_BYTES, validate_provenance
from core.voice import (
    _AUDIO_EXTS, _PICTURE_EXTS, _find_voice_audio, _load_meta,
    _save_meta, _validate_voice_meta, _voice_problems,
)

FORMAT_VERSION = 1
MAX_BUNDLE_BYTES = 150 * 1024 * 1024      # compressed upload
MAX_TOTAL_BYTES = 200 * 1024 * 1024       # uncompressed, all members
MAX_AUDIO_BYTES = 150 * 1024 * 1024
MAX_PICTURE_BYTES = 10 * 1024 * 1024
MAX_TEXT_BYTES = 256 * 1024
MAX_RATIO = 200                           # per-member compression ratio (zip-bomb)
ID_RE = re.compile(r"^[A-Za-z0-9_\-][A-Za-z0-9_\-.]{0,99}$")
_PICTURE_SET = {e.lstrip(".") for e in _PICTURE_EXTS}
_AUDIO_SET = {e.lstrip(".") for e in _AUDIO_EXTS}
_CONSENT_SET = {e.lstrip(".") for e in CONSENT_EXTS}


class BundleError(ValueError):
    """Invalid bundle (-> HTTP 400)."""


def _limit_for(name: str) -> int | None:
    if name in ("manifest.json", "meta.json", "reference.txt"):
        return MAX_TEXT_BYTES
    stem, _, ext = name.rpartition(".")
    if stem == "audio" and ext in _AUDIO_SET:
        return MAX_AUDIO_BYTES
    if stem == "picture" and ext in _PICTURE_SET:
        return MAX_PICTURE_BYTES
    if stem == "consent" and ext in _CONSENT_SET:
        return MAX_CONSENT_BYTES
    return None


def _app_version() -> str:
    try:
        return (Path(__file__).resolve().parents[1] / "VERSION").read_text().strip()
    except OSError:
        return "unknown"


# ── export ────────────────────────────────────────────────────────────────────

class _Sink:
    """Non-seekable write target; the generator drains it between chunks."""
    def __init__(self): self.buf = bytearray(); self.pos = 0
    def write(self, b): self.buf += b; self.pos += len(b); return len(b)
    def tell(self): return self.pos
    def flush(self):
        pass  # in-memory buffer, nothing to flush; zipfile just calls it
    def drain(self):
        out = bytes(self.buf); self.buf.clear(); return out


def export_members(wav: Path) -> tuple[dict, list[tuple[str, Path | bytes]]]:
    """Return (meta_for_bundle, [(member_name, path-or-bytes)]) for a voice."""
    meta = _load_meta(wav)
    meta["audio_problems"] = list(meta.get("audio_problems") or _voice_problems(wav))
    members: list[tuple[str, Path | bytes]] = [(f"audio{wav.suffix.lower()}", wav)]
    ref = wav.with_suffix(".reference.txt")
    if ref.is_file():
        members.append(("reference.txt", ref))
    for ext in _PICTURE_EXTS:
        pic = wav.with_suffix(ext)
        if pic.is_file():
            members.append((f"picture{ext}", pic))
            break
    consent = wav.with_suffix(".consent")
    rec = (meta.get("provenance") or {}).get("recorded_consent") or ""
    if consent.is_file() and Path(rec).suffix.lower() in CONSENT_EXTS:
        members.append((f"consent{Path(rec).suffix.lower()}", consent))
    members.append(("meta.json", (json.dumps(meta, indent=2, ensure_ascii=False) + "\n").encode()))
    return meta, members


def stream_bundle(voice_id: str, members: list[tuple[str, Path | bytes]]) -> Iterator[bytes]:
    """Yield a ZIP; manifest.json (with sha256 per file) is written last."""
    sink, hashes = _Sink(), {}
    with zipfile.ZipFile(sink, "w", zipfile.ZIP_DEFLATED) as zf:
        for name, src in members:
            h = hashlib.sha256()
            with zf.open(name, "w", force_zip64=True) as out:
                if isinstance(src, bytes):
                    h.update(src); out.write(src)
                else:
                    with src.open("rb") as f:
                        while chunk := f.read(1024 * 1024):
                            h.update(chunk); out.write(chunk)
                            if sink.buf:
                                yield sink.drain()
            hashes[name] = h.hexdigest()
            yield sink.drain()
        manifest = {"format": FORMAT_VERSION, "app_version": _app_version(),
                    "voice_id": voice_id, "files": hashes}
        zf.writestr("manifest.json", json.dumps(manifest, indent=2))
    yield sink.drain()


# ── import ────────────────────────────────────────────────────────────────────

def _read_member(zf: zipfile.ZipFile, info: zipfile.ZipInfo, limit: int, dest: Path | None = None) -> tuple[bytes, str]:
    """Read at most `limit` bytes (declared size is not trusted); returns (data|b'', sha256)."""
    h, total, data = hashlib.sha256(), 0, bytearray()
    out = dest.open("wb") if dest else None
    try:
        with zf.open(info) as f:
            while chunk := f.read(1024 * 1024):
                total += len(chunk)
                if total > limit:
                    raise BundleError(f"Bundle member too large: {info.filename}")
                h.update(chunk)
                (out.write(chunk) if out else data.extend(chunk))
    finally:
        if out: out.close()
    return bytes(data), h.hexdigest()


def import_bundle(zip_path: Path, voices_dir: Path, scan_dir: Path, decode_check=None) -> dict:
    """Validate and install a bundle. Returns {voice_id, status, wav}.

    status: "imported" | "renamed" (name taken -> suffix) | "unchanged" (identical already present).
    `decode_check(path)` must raise if the audio is not decodable.
    """
    if not zipfile.is_zipfile(zip_path):
        raise BundleError("Not a ZIP file")
    voices_dir.mkdir(parents=True, exist_ok=True)
    stage = Path(tempfile.mkdtemp(prefix=".import-", dir=voices_dir))
    try:
        with zipfile.ZipFile(zip_path) as zf:
            infos = zf.infolist()
            names = [i.filename for i in infos]
            if len(set(names)) != len(names):
                raise BundleError("Duplicate members in bundle")
            total = 0
            for i in infos:
                if i.is_dir() or _limit_for(i.filename) is None:  # also rejects any path / traversal / unknown name
                    raise BundleError("Unexpected member in bundle")
                if i.file_size > _limit_for(i.filename) or (i.compress_size and i.file_size / i.compress_size > MAX_RATIO and i.file_size > 1024 * 1024):
                    raise BundleError("Bundle member too large or suspiciously compressed")
                total += i.file_size
            if total > MAX_TOTAL_BYTES:
                raise BundleError("Bundle is too large when unpacked")
            by_name = {i.filename: i for i in infos}
            if "manifest.json" not in by_name:
                raise BundleError("manifest.json missing")
            raw, _ = _read_member(zf, by_name["manifest.json"], MAX_TEXT_BYTES)
            try:
                manifest = json.loads(raw)
            except ValueError:
                raise BundleError("manifest.json is not valid JSON")
            if not isinstance(manifest, dict) or manifest.get("format") != FORMAT_VERSION:
                raise BundleError("Unsupported bundle format version")
            files = manifest.get("files")
            if not isinstance(files, dict) or set(files) != set(by_name) - {"manifest.json"}:
                raise BundleError("Manifest does not match bundle contents")
            voice_id = manifest.get("voice_id")
            if not isinstance(voice_id, str) or not ID_RE.match(voice_id) or ".." in voice_id:
                raise BundleError("Invalid voice id in manifest")
            audio = [n for n in files if n.startswith("audio.")]
            if len(audio) != 1:
                raise BundleError("Bundle needs exactly one audio file")
            staged: dict[str, Path] = {}
            for n, expected in files.items():
                p = stage / n  # n is a validated fixed name, never a path from the archive
                _, digest = _read_member(zf, by_name[n], _limit_for(n), p)
                if digest != expected:
                    raise BundleError(f"Hash mismatch: {n}")
                staged[n] = p
        # decode + metadata re-validation (nothing is installed yet)
        if decode_check:
            try:
                decode_check(staged[audio[0]])
            except Exception:
                raise BundleError("Audio file cannot be decoded")
        try:
            meta = json.loads(staged["meta.json"].read_bytes()) if "meta.json" in staged else {}
            consent = next((n for n in staged if n.startswith("consent.")), None)
            if "provenance" in meta:
                prev = {"recorded_consent": "x"} if consent else None
                meta["provenance"] = validate_provenance(meta["provenance"], prev)
                meta["provenance"].pop("recorded_consent", None)
                if consent:
                    meta["provenance"]["recorded_consent"] = "consent" + Path(consent).suffix
            _validate_voice_meta(meta)
        except (ValueError, TypeError, AttributeError) as exc:
            raise BundleError(f"Invalid metadata: {exc}")
        # name handling: never overwrite; identical audio already present -> idempotent no-op
        new_hash = files[audio[0]]
        suffix = Path(audio[0]).suffix
        vid, status, n = voice_id, "imported", 1
        while (existing := _find_voice_audio(vid, scan_dir)) is not None or (voices_dir / f"{vid}{suffix}").exists():
            if existing is not None and hashlib.sha256(existing.read_bytes()).hexdigest() == new_hash:
                return {"voice_id": vid, "status": "unchanged", "wav": str(existing)}
            n += 1
            vid, status = f"{voice_id}_{n}", "renamed"
            if n > 99:
                raise BundleError("Too many voices with this name")
        if consent:
            meta.get("provenance", {}).update(recorded_consent=f"{vid}{Path(consent).suffix}")
        dest = voices_dir / f"{vid}{suffix}"
        shutil.move(str(staged[audio[0]]), dest)
        if "reference.txt" in staged:
            shutil.move(str(staged["reference.txt"]), dest.with_suffix(".reference.txt"))
        for n_, p in staged.items():
            if n_.startswith("picture."):
                shutil.move(str(p), dest.with_suffix(Path(n_).suffix))
            elif consent and n_ == consent:
                shutil.move(str(p), dest.with_suffix(".consent"))
        meta.pop("needs_tts_restart", None)
        _save_meta(dest, meta)
        return {"voice_id": vid, "status": status, "wav": str(dest)}
    finally:
        shutil.rmtree(stage, ignore_errors=True)