"""
Bundle export/import helpers for:
- Agent package (personality, memory, adapters, related settings)
- Full migration package (agent package + chats and optional encryption key)
"""
from __future__ import annotations

import hashlib
import ipaddress
import json
import os
import shutil
import tempfile
import zipfile
from dataclasses import dataclass
from datetime import datetime, timezone
from urllib.parse import urlparse

from core import get_config, reload_config, save_config
from persistence import MUGI_DIR, secure_atomic_write


BUNDLE_SCHEMA_VERSION = 1
AGENT_BUNDLE_TYPE = "agent_bundle_v1"
MIGRATION_BUNDLE_TYPE = "migration_bundle_v1"
SUPPORTED_BUNDLE_TYPES = {AGENT_BUNDLE_TYPE, MIGRATION_BUNDLE_TYPE}

MANIFEST_FILENAME = "manifest.json"
_DRIVE_WORKSPACE_FILES = {
    "workspace/memory/drives.json": b'{\n  "active_drives": [],\n  "resolved_drives": [],\n  "expired_drives": [],\n  "stats": {\n    "total_discovered": 0,\n    "total_resolved": 0,\n    "total_expired": 0,\n    "longest_active_days": 0\n  }\n}\n',
    "workspace/memory/growth-events.md": b"# Growth Events\n",
    "workspace/memory/tool_errors.json": b"[]\n",
}


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _file_sha256(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def _bytes_sha256(data: bytes) -> str:
    h = hashlib.sha256()
    h.update(data)
    return h.hexdigest()


def _safe_rel(rel_path: str) -> str:
    rel = rel_path.replace("\\", "/").lstrip("/")
    if not rel or rel == ".":
        raise ValueError("Invalid relative path")
    if rel.startswith("../") or "/../" in rel or rel.endswith("/.."):
        raise ValueError("Path traversal detected")
    return rel


def _zip_write_file(zipf: zipfile.ZipFile, src_path: str, dst_rel: str, records: list[dict]) -> None:
    safe_rel = _safe_rel(dst_rel)
    zipf.write(src_path, arcname=safe_rel)
    st = os.stat(src_path)
    records.append(
        {
            "path": safe_rel,
            "sha256": _file_sha256(src_path),
            "size": int(st.st_size),
        }
    )


def _zip_write_bytes(zipf: zipfile.ZipFile, data: bytes, dst_rel: str, records: list[dict]) -> None:
    safe_rel = _safe_rel(dst_rel)
    zipf.writestr(safe_rel, data)
    records.append(
        {
            "path": safe_rel,
            "sha256": _bytes_sha256(data),
            "size": len(data),
        }
    )


def _iter_files(root: str) -> list[str]:
    files: list[str] = []
    if not os.path.isdir(root):
        return files
    for base, _, names in os.walk(root):
        for name in names:
            full = os.path.join(base, name)
            if os.path.isfile(full):
                files.append(full)
    files.sort()
    return files


@dataclass
class _Section:
    source: str
    bundle_prefix: str
    kind: str  # "file" | "dir"
    include_in_agent: bool
    include_in_migration: bool


def _sections() -> list[_Section]:
    workspace = os.path.join(MUGI_DIR, "workspace")
    return [
        _Section(os.path.join(MUGI_DIR, "mugi_config.toml"), "config/mugi_config.toml", "file", True, True),
        _Section(os.path.join(MUGI_DIR, "app_prefs.json"), "prefs/app_prefs.json", "file", True, True),
        _Section(workspace, "workspace", "dir", True, True),
        _Section(os.path.join(MUGI_DIR, "mlx", "adapters"), "mlx/adapters", "dir", True, True),
        _Section(os.path.join(MUGI_DIR, "evolution", "evolution_meta.json"), "evolution/evolution_meta.json", "file", True, True),
        _Section(os.path.join(MUGI_DIR, "evolution", "trajectories"), "evolution/trajectories", "dir", True, True),
        _Section(os.path.join(MUGI_DIR, "plugins"), "plugins", "dir", True, True),
        _Section(os.path.join(MUGI_DIR, "supertonic_voices"), "voices/supertonic_voices", "dir", True, True),
        _Section(os.path.join(MUGI_DIR, "chats.db"), "chat/chats.db", "file", False, True),
        _Section(os.path.join(MUGI_DIR, "chats_db.json"), "chat/chats_db.json", "file", False, True),
    ]


def _bundle_sections(bundle_type: str) -> list[_Section]:
    out: list[_Section] = []
    for s in _sections():
        if bundle_type == AGENT_BUNDLE_TYPE and s.include_in_agent:
            out.append(s)
        elif bundle_type == MIGRATION_BUNDLE_TYPE and s.include_in_migration:
            out.append(s)
    return out


def _build_manifest(bundle_type: str, file_records: list[dict], encryption_key_mode: str) -> dict:
    return {
        "schema_version": BUNDLE_SCHEMA_VERSION,
        "bundle_type": bundle_type,
        "created_at": _now_iso(),
        "app_version": os.environ.get("MUGI_APP_VERSION", "unknown"),
        "encryption_key_mode": encryption_key_mode,
        "included_sections": sorted({r["path"].split("/", 1)[0] for r in file_records}),
        "files": file_records,
    }


def create_bundle(bundle_type: str, destination_zip_path: str, include_encryption_key: bool = False) -> dict:
    if bundle_type not in SUPPORTED_BUNDLE_TYPES:
        raise ValueError("unsupported bundle type")
    dest = os.path.abspath(os.path.expanduser(destination_zip_path or ""))
    if not dest:
        raise ValueError("destination path required")
    parent = os.path.dirname(dest)
    if not parent or not os.path.isdir(parent):
        raise ValueError("destination directory does not exist")

    os.makedirs(os.path.dirname(dest), exist_ok=True)
    file_records: list[dict] = []

    key_mode = "not_applicable"
    if bundle_type == MIGRATION_BUNDLE_TYPE:
        key_mode = "embedded" if include_encryption_key else "separate"

    with tempfile.TemporaryDirectory(prefix="mugi_bundle_export_") as _tmp:
        with zipfile.ZipFile(dest, "w", compression=zipfile.ZIP_DEFLATED) as zipf:
            for section in _bundle_sections(bundle_type):
                if section.kind == "file":
                    if os.path.isfile(section.source):
                        _zip_write_file(zipf, section.source, section.bundle_prefix, file_records)
                else:
                    for src in _iter_files(section.source):
                        rel = os.path.relpath(src, section.source).replace("\\", "/")
                        bundle_rel = f"{section.bundle_prefix}/{rel}"
                        _zip_write_file(zipf, src, bundle_rel, file_records)

            # Ensure drive-engine files are present in agent/migration bundles even
            # if they haven't been created yet in workspace memory.
            added_paths = {r.get("path", "") for r in file_records}
            for rel, payload in _DRIVE_WORKSPACE_FILES.items():
                if rel not in added_paths:
                    _zip_write_bytes(zipf, payload, rel, file_records)

            if bundle_type == MIGRATION_BUNDLE_TYPE and include_encryption_key:
                from chat_storage import export_encryption_key

                key_data = export_encryption_key().encode("utf-8")
                _zip_write_bytes(zipf, key_data, "chat/encryption_key.txt", file_records)

            manifest = _build_manifest(bundle_type, file_records, encryption_key_mode=key_mode)
            manifest_bytes = json.dumps(manifest, indent=2).encode("utf-8")
            zipf.writestr(MANIFEST_FILENAME, manifest_bytes)

    return {
        "ok": True,
        "path": dest,
        "bundle_type": bundle_type,
        "file_count": len(file_records),
        "encryption_key_mode": key_mode,
    }


def _is_symlink_zip_entry(info: zipfile.ZipInfo) -> bool:
    mode = (info.external_attr >> 16) & 0xFFFF
    return (mode & 0o170000) == 0o120000


def _extract_bundle_safely(zip_path: str, extract_dir: str) -> None:
    with zipfile.ZipFile(zip_path, "r") as zipf:
        for info in zipf.infolist():
            name = info.filename.replace("\\", "/")
            if not name:
                continue
            if name.startswith("/") or name.startswith("../") or "/../" in name:
                raise ValueError("zip contains unsafe path")
            if _is_symlink_zip_entry(info):
                raise ValueError("zip contains symlink entry")
            target = os.path.realpath(os.path.join(extract_dir, name))
            extract_root = os.path.realpath(extract_dir)
            if not (target == extract_root or target.startswith(extract_root + os.sep)):
                raise ValueError("zip path escapes extraction dir")
            if info.is_dir():
                os.makedirs(target, exist_ok=True)
                continue
            os.makedirs(os.path.dirname(target), exist_ok=True)
            with zipf.open(info, "r") as src, open(target, "wb") as dst:
                shutil.copyfileobj(src, dst)


def _load_manifest(extract_dir: str) -> dict:
    manifest_path = os.path.join(extract_dir, MANIFEST_FILENAME)
    if not os.path.isfile(manifest_path):
        raise ValueError("manifest.json missing")
    with open(manifest_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    if not isinstance(data, dict):
        raise ValueError("manifest.json invalid")
    return data


def _validate_manifest(bundle_type: str, manifest: dict, extract_dir: str) -> list[dict]:
    if int(manifest.get("schema_version", -1)) != BUNDLE_SCHEMA_VERSION:
        raise ValueError("unsupported manifest schema_version")
    if manifest.get("bundle_type") != bundle_type:
        raise ValueError("bundle type mismatch")
    files = manifest.get("files")
    if not isinstance(files, list):
        raise ValueError("manifest files list missing")
    validated: list[dict] = []
    for entry in files:
        if not isinstance(entry, dict):
            raise ValueError("manifest file entry invalid")
        rel = _safe_rel(str(entry.get("path") or ""))
        expected = str(entry.get("sha256") or "").strip().lower()
        if len(expected) != 64:
            raise ValueError("manifest checksum invalid")
        abs_path = os.path.realpath(os.path.join(extract_dir, rel))
        if not os.path.isfile(abs_path):
            raise ValueError(f"missing file in bundle: {rel}")
        actual = _file_sha256(abs_path)
        if actual != expected:
            raise ValueError(f"checksum mismatch: {rel}")
        validated.append({"rel": rel, "abs": abs_path})
    return validated


def _target_path_for_bundle_rel(rel: str) -> str | None:
    if rel == "config/mugi_config.toml":
        return os.path.join(MUGI_DIR, "mugi_config.toml")
    if rel == "prefs/app_prefs.json":
        return os.path.join(MUGI_DIR, "app_prefs.json")
    if rel.startswith("workspace/"):
        rest = rel[len("workspace/") :]
        return os.path.join(MUGI_DIR, "workspace", rest)
    if rel.startswith("mlx/adapters/"):
        rest = rel[len("mlx/adapters/") :]
        return os.path.join(MUGI_DIR, "mlx", "adapters", rest)
    if rel == "evolution/evolution_meta.json":
        return os.path.join(MUGI_DIR, "evolution", "evolution_meta.json")
    if rel.startswith("evolution/trajectories/"):
        rest = rel[len("evolution/trajectories/") :]
        return os.path.join(MUGI_DIR, "evolution", "trajectories", rest)
    if rel.startswith("plugins/"):
        rest = rel[len("plugins/") :]
        return os.path.join(MUGI_DIR, "plugins", rest)
    if rel.startswith("voices/supertonic_voices/"):
        rest = rel[len("voices/supertonic_voices/") :]
        return os.path.join(MUGI_DIR, "supertonic_voices", rest)
    if rel == "chat/chats.db":
        return os.path.join(MUGI_DIR, "chats.db")
    if rel == "chat/chats_db.json":
        return os.path.join(MUGI_DIR, "chats_db.json")
    if rel == "chat/encryption_key.txt":
        return None
    return None


def _detect_existing(validated_files: list[dict]) -> list[str]:
    existing: set[str] = set()
    for item in validated_files:
        rel = item["rel"]
        target = _target_path_for_bundle_rel(rel)
        if target and os.path.exists(target):
            existing.add(target)
    return sorted(existing)


def _replace_with_backup(src: str, dst: str, backup_root: str, applied: list[dict]) -> None:
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    backup_path = None
    existed = os.path.exists(dst)
    if existed:
        backup_path = os.path.join(backup_root, f"backup_{len(applied)}")
        if os.path.isdir(dst):
            shutil.copytree(dst, backup_path)
        else:
            os.makedirs(os.path.dirname(backup_path), exist_ok=True)
            shutil.copy2(dst, backup_path)
    if os.path.isdir(dst):
        shutil.rmtree(dst, ignore_errors=True)
    elif os.path.isfile(dst):
        os.remove(dst)

    if os.path.isdir(src):
        shutil.copytree(src, dst)
    else:
        shutil.copy2(src, dst)
    applied.append({"dst": dst, "backup": backup_path, "existed": existed})


def _rollback(applied: list[dict]) -> None:
    for item in reversed(applied):
        dst = item["dst"]
        backup = item.get("backup")
        existed = bool(item.get("existed"))
        try:
            if os.path.isdir(dst):
                shutil.rmtree(dst, ignore_errors=True)
            elif os.path.isfile(dst):
                os.remove(dst)
            if existed and backup and os.path.exists(backup):
                if os.path.isdir(backup):
                    shutil.copytree(backup, dst)
                else:
                    os.makedirs(os.path.dirname(dst), exist_ok=True)
                    shutil.copy2(backup, dst)
        except Exception:
            pass


def _is_private_host(host: str) -> bool:
    h = (host or "").strip().lower()
    if not h:
        return False
    if h in {"localhost", "127.0.0.1"}:
        return False
    if h.endswith(".local"):
        return True
    try:
        ip = ipaddress.ip_address(h)
        return bool(ip.is_private or ip.is_loopback or ip.is_link_local)
    except ValueError:
        return False


def _sanitize_app_prefs_and_collect_warnings() -> list[str]:
    warnings: list[str] = []
    prefs_path = os.path.join(MUGI_DIR, "app_prefs.json")
    if not os.path.isfile(prefs_path):
        return warnings
    try:
        with open(prefs_path, "r", encoding="utf-8") as f:
            prefs = json.load(f)
        if not isinstance(prefs, dict):
            return warnings
    except Exception:
        return warnings

    changed = False

    mlx_model = str(prefs.get("mlx_model") or "").strip()
    if mlx_model and os.path.isabs(mlx_model) and not os.path.isdir(mlx_model):
        prefs["mlx_model"] = ""
        warnings.append("Reset invalid absolute MLX model path from imported preferences.")
        changed = True

    default_backend_url = "http://127.0.0.1:1234/v1"

    for key in ("backend_base_url", "last_remote_backend_url"):
        val = str(prefs.get(key) or "").strip()
        if not val:
            continue
        parsed = urlparse(val if "://" in val else f"http://{val}")
        if _is_private_host(parsed.hostname or ""):
            prefs[key] = default_backend_url if key == "backend_base_url" else ""
            warnings.append(f"Reset LAN-specific `{key}` from imported preferences.")
            changed = True

    selected_model = str(prefs.get("selected_llmster_model") or "").strip()
    if selected_model:
        try:
            import app_llmster

            available = app_llmster.llmster_list_models()
            keys = {str(m.get("key") or "").strip() for m in available if isinstance(m, dict)}
            keys.discard("")
            if keys and selected_model not in keys:
                prefs["selected_llmster_model"] = ""
                warnings.append("Cleared selected llmster model because it is not available on this machine.")
                changed = True
        except Exception:
            pass

    if changed:
        secure_atomic_write(prefs_path, json.dumps(prefs, indent=2))
        try:
            from utils.app_prefs_store import invalidate_app_prefs_cache
            invalidate_app_prefs_cache()
        except Exception:
            pass
    return warnings


def _sanitize_config_allowed_dirs_and_collect_warnings() -> list[str]:
    warnings: list[str] = []
    try:
        config = get_config(use_cache=False)
        allowed = ((config.get("paths") or {}).get("allowed_dirs")) or []
        if isinstance(allowed, list) and allowed:
            config.setdefault("paths", {})["allowed_dirs"] = []
            save_config(config)
            warnings.append("Cleared imported allowed directories; reconfigure local folder access in Settings.")
    except Exception:
        pass
    return warnings


def import_bundle(
    expected_bundle_type: str,
    bundle_zip_path: str,
    confirm_overwrite: bool = False,
) -> dict:
    if expected_bundle_type not in SUPPORTED_BUNDLE_TYPES:
        raise ValueError("unsupported bundle type")

    src_zip = os.path.abspath(os.path.expanduser(bundle_zip_path or ""))
    if not src_zip or not os.path.isfile(src_zip):
        raise ValueError("bundle file not found")

    with tempfile.TemporaryDirectory(prefix="mugi_bundle_import_") as extract_dir:
        _extract_bundle_safely(src_zip, extract_dir)
        manifest = _load_manifest(extract_dir)
        validated = _validate_manifest(expected_bundle_type, manifest, extract_dir)

        existing_paths = _detect_existing(validated)
        if existing_paths and not confirm_overwrite:
            return {
                "ok": False,
                "requires_overwrite_confirmation": True,
                "existing_paths": existing_paths,
                "message": "Import will overwrite existing data. Export current data first if you want to keep it.",
            }

        # Build section-level apply units so directory imports are clean replaces.
        section_map = {
            "config/mugi_config.toml": (os.path.join(extract_dir, "config", "mugi_config.toml"), os.path.join(MUGI_DIR, "mugi_config.toml")),
            "prefs/app_prefs.json": (os.path.join(extract_dir, "prefs", "app_prefs.json"), os.path.join(MUGI_DIR, "app_prefs.json")),
            "workspace": (os.path.join(extract_dir, "workspace"), os.path.join(MUGI_DIR, "workspace")),
            "mlx/adapters": (os.path.join(extract_dir, "mlx", "adapters"), os.path.join(MUGI_DIR, "mlx", "adapters")),
            "evolution/evolution_meta.json": (os.path.join(extract_dir, "evolution", "evolution_meta.json"), os.path.join(MUGI_DIR, "evolution", "evolution_meta.json")),
            "evolution/trajectories": (os.path.join(extract_dir, "evolution", "trajectories"), os.path.join(MUGI_DIR, "evolution", "trajectories")),
            "plugins": (os.path.join(extract_dir, "plugins"), os.path.join(MUGI_DIR, "plugins")),
            "voices/supertonic_voices": (os.path.join(extract_dir, "voices", "supertonic_voices"), os.path.join(MUGI_DIR, "supertonic_voices")),
            "chat/chats.db": (os.path.join(extract_dir, "chat", "chats.db"), os.path.join(MUGI_DIR, "chats.db")),
            "chat/chats_db.json": (os.path.join(extract_dir, "chat", "chats_db.json"), os.path.join(MUGI_DIR, "chats_db.json")),
        }
        apply_order = [
            "config/mugi_config.toml",
            "prefs/app_prefs.json",
            "workspace",
            "mlx/adapters",
            "evolution/evolution_meta.json",
            "evolution/trajectories",
            "plugins",
            "voices/supertonic_voices",
            "chat/chats.db",
            "chat/chats_db.json",
        ]

        has_mlx = any(item["rel"].startswith("mlx/adapters/") for item in validated)
        if has_mlx:
            try:
                import app_mlx

                app_mlx.stop_mlx_server()
            except Exception:
                pass

        applied: list[dict] = []
        warnings: list[str] = []
        with tempfile.TemporaryDirectory(prefix="mugi_bundle_backup_") as backup_root:
            try:
                for key in apply_order:
                    src, dst = section_map[key]
                    if not os.path.exists(src):
                        continue
                    _replace_with_backup(src, dst, backup_root, applied)
            except Exception:
                _rollback(applied)
                raise

        key_path = os.path.join(extract_dir, "chat", "encryption_key.txt")
        if os.path.isfile(key_path):
            try:
                from chat_storage import import_encryption_key

                with open(key_path, "r", encoding="utf-8") as f:
                    key_data = f.read().strip()
                if key_data:
                    import_encryption_key(key_data)
                    warnings.append("Imported chat encryption key from bundle.")
            except Exception:
                warnings.append("Could not import bundled encryption key; import it manually if chat data is unreadable.")
        elif expected_bundle_type == MIGRATION_BUNDLE_TYPE:
            warnings.append("Migration bundle did not include encryption key; import it separately if needed.")

        warnings.extend(_sanitize_app_prefs_and_collect_warnings())
        warnings.extend(_sanitize_config_allowed_dirs_and_collect_warnings())
        warnings.append("Keychain credentials are not transferred (Discord token, IMAP password, and similar secrets). Re-enter them in Settings.")
        warnings.append("Imported plugins may require re-approval on first use.")
        try:
            from utils.app_prefs_store import invalidate_app_prefs_cache
            invalidate_app_prefs_cache()
        except Exception:
            pass

        reload_config()

        return {
            "ok": True,
            "bundle_type": expected_bundle_type,
            "warnings": warnings,
            "existing_paths": existing_paths,
            "manifest": {
                "bundle_type": manifest.get("bundle_type"),
                "schema_version": manifest.get("schema_version"),
                "created_at": manifest.get("created_at"),
                "file_count": len(validated),
            },
        }
