import hashlib
import json
import os
import shutil
import tempfile
import time
import urllib.request
import zipfile
from contextlib import contextmanager
from datetime import datetime, timezone


BASE_DIR = os.path.dirname(os.path.dirname(__file__))
CONFIG_PATH = os.path.join(BASE_DIR, "config", "config.json")
UPDATES_DIR = os.path.join(BASE_DIR, ".updates")
DOWNLOAD_DIR = os.path.join(UPDATES_DIR, "downloads")
BACKUP_DIR = os.path.join(UPDATES_DIR, "backups")
STATE_PATH = os.path.join(UPDATES_DIR, "state.json")
LOCK_PATH = os.path.join(UPDATES_DIR, "update.lock")

PRESERVE_NAMES = {
    ".cache",
    ".updates",
    "__pycache__",
    "config",
    "data",
    "api_calls.log",
    "api_usage.json",
    "airlabs_usage.json",
    "close.txt",
    "current_overhead.json",
    "current_tracked.json",
    "farthest.txt",
    "flight_counter.json",
    "flightaware_usage.json",
    "tracked_flight.json",
    "wifi_connect.log",
}

DEFAULT_CONFIG = {
    "enabled": False,
    "manifest_url": "",
    "check_on_start": True,
    "auto_install": True,
    "timeout_seconds": 20,
    "keep_backups": 2,
}


class UpdateError(Exception):
    pass


@contextmanager
def _operation_lock():
    os.makedirs(UPDATES_DIR, exist_ok=True)
    try:
        fd = os.open(LOCK_PATH, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
    except FileExistsError:
        raise UpdateError("Another update operation is already running")
    try:
        os.write(fd, str(os.getpid()).encode("ascii"))
        os.close(fd)
        yield
    finally:
        try:
            os.remove(LOCK_PATH)
        except OSError:
            pass


def _now_iso():
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat()


def _load_json(path, default):
    try:
        with open(path, "r", encoding="utf-8-sig") as f:
            return json.load(f)
    except Exception:
        return default


def _write_json(path, data):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)


def update_config():
    cfg = _load_json(CONFIG_PATH, {})
    updates = dict(DEFAULT_CONFIG)
    updates.update(cfg.get("updates", {}) or {})
    updates["timeout_seconds"] = max(5, min(120, int(updates.get("timeout_seconds") or 20)))
    updates["keep_backups"] = max(1, min(10, int(updates.get("keep_backups") or 2)))
    return updates


def _state():
    state = _load_json(STATE_PATH, {})
    state.setdefault("installed_version", _read_installed_version())
    state.setdefault("last_status", "never checked")
    return state


def _save_state(**updates):
    state = _state()
    state.update(updates)
    state["updated_at"] = _now_iso()
    _write_json(STATE_PATH, state)
    return state


def _read_installed_version():
    version_file = os.path.join(BASE_DIR, "VERSION")
    try:
        with open(version_file, "r", encoding="utf-8") as f:
            version = f.read().strip()
            if version:
                return version
    except OSError:
        pass
    return "0.0.0"


def _version_tuple(value):
    parts = []
    for piece in str(value or "0").replace("-", ".").split("."):
        number = "".join(ch for ch in piece if ch.isdigit())
        parts.append(int(number or 0))
    return tuple(parts)


def is_newer(remote_version, installed_version=None):
    installed = installed_version or _state().get("installed_version") or _read_installed_version()
    return _version_tuple(remote_version) > _version_tuple(installed)


def _fetch_json(url, timeout):
    with urllib.request.urlopen(url, timeout=timeout) as response:
        return json.loads(response.read().decode("utf-8"))


def fetch_manifest(config=None):
    config = config or update_config()
    url = str(config.get("manifest_url") or "").strip()
    if not url:
        raise UpdateError("No update manifest URL configured")
    if not url.lower().startswith("https://"):
        raise UpdateError("Manifest URL must use https://")
    manifest = _fetch_json(url, config["timeout_seconds"])
    if not isinstance(manifest, dict):
        raise UpdateError("Manifest is not a JSON object")
    for field in ("version", "url", "sha256"):
        if not manifest.get(field):
            raise UpdateError(f"Manifest missing {field}")
    if not str(manifest["url"]).lower().startswith("https://"):
        raise UpdateError("Update ZIP URL must use https://")
    return manifest


def _download(url, dest, timeout, progress=None):
    req = urllib.request.Request(url, headers={"User-Agent": "its-a-plane-updater/1.0"})
    with urllib.request.urlopen(req, timeout=timeout) as response:
        total = int(response.headers.get("Content-Length") or 0)
        done = 0
        with open(dest, "wb") as f:
            while True:
                chunk = response.read(1024 * 64)
                if not chunk:
                    break
                f.write(chunk)
                done += len(chunk)
                if progress and total:
                    progress("Downloading update", 10 + int(45 * done / total))


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


def _safe_extract(zip_path, dest):
    with zipfile.ZipFile(zip_path) as zf:
        for member in zf.infolist():
            name = member.filename.replace("\\", "/")
            if name.startswith("/") or ".." in name.split("/"):
                raise UpdateError(f"Unsafe ZIP path: {member.filename}")
        zf.extractall(dest)


def _payload_root(extract_dir):
    entries = [entry for entry in os.listdir(extract_dir) if not entry.startswith("__MACOSX")]
    if len(entries) == 1:
        candidate = os.path.join(extract_dir, entries[0])
        if os.path.isdir(candidate) and os.path.exists(os.path.join(candidate, "its-a-plane.py")):
            return candidate
    return extract_dir


def _backup_current(version):
    os.makedirs(BACKUP_DIR, exist_ok=True)
    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    backup_path = os.path.join(BACKUP_DIR, f"{version}-{stamp}.zip")
    with zipfile.ZipFile(backup_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        for root, dirs, files in os.walk(BASE_DIR):
            rel_root = os.path.relpath(root, BASE_DIR)
            if rel_root == ".":
                rel_root = ""
            dirs[:] = [
                d for d in dirs
                if d not in PRESERVE_NAMES and not d.startswith(".git") and d != "__pycache__"
            ]
            for filename in files:
                rel = os.path.join(rel_root, filename) if rel_root else filename
                if rel.split(os.sep)[0] in PRESERVE_NAMES:
                    continue
                zf.write(os.path.join(root, filename), rel)
    return backup_path


def _copy_payload(payload_root):
    for name in os.listdir(payload_root):
        if name in PRESERVE_NAMES:
            continue
        src = os.path.join(payload_root, name)
        dst = os.path.join(BASE_DIR, name)
        if os.path.isdir(src):
            if os.path.exists(dst):
                shutil.rmtree(dst)
            shutil.copytree(src, dst)
        else:
            shutil.copy2(src, dst)


def _prune_backups(keep):
    if not os.path.isdir(BACKUP_DIR):
        return
    backups = sorted(
        (os.path.join(BACKUP_DIR, name) for name in os.listdir(BACKUP_DIR) if name.endswith(".zip")),
        key=os.path.getmtime,
        reverse=True,
    )
    for path in backups[keep:]:
        try:
            os.remove(path)
        except OSError:
            pass


def check_for_update():
    config = update_config()
    if not config.get("enabled"):
        state = _save_state(last_status="updates disabled", last_checked=_now_iso())
        return {"ok": True, "enabled": False, "state": state}
    manifest = fetch_manifest(config)
    installed = _state().get("installed_version") or _read_installed_version()
    available = is_newer(manifest.get("version"), installed)
    state = _save_state(
        last_status="update available" if available else "already current",
        last_checked=_now_iso(),
        latest_version=manifest.get("version"),
        latest_manifest=manifest,
        last_error=None,
    )
    return {
        "ok": True,
        "enabled": True,
        "update_available": available,
        "installed_version": installed,
        "latest_version": manifest.get("version"),
        "manifest": manifest,
        "state": state,
    }


def install_update(manifest=None, progress=None):
    with _operation_lock():
        config = update_config()
        if not config.get("enabled"):
            raise UpdateError("Updates are disabled")
        manifest = manifest or fetch_manifest(config)
        installed = _state().get("installed_version") or _read_installed_version()
        if not is_newer(manifest.get("version"), installed):
            return _save_state(last_status="already current", latest_version=manifest.get("version"))

        os.makedirs(DOWNLOAD_DIR, exist_ok=True)
        os.makedirs(BACKUP_DIR, exist_ok=True)
        zip_path = os.path.join(DOWNLOAD_DIR, f"plane-tracker-{manifest['version']}.zip")

        if progress:
            progress("Downloading update", 10)
        _download(manifest["url"], zip_path, config["timeout_seconds"], progress=progress)

        if progress:
            progress("Verifying update", 60)
        digest = _sha256(zip_path)
        if digest.lower() != str(manifest["sha256"]).lower():
            raise UpdateError("Downloaded ZIP checksum did not match manifest")

        with tempfile.TemporaryDirectory(prefix="plane-update-") as tmp:
            if progress:
                progress("Extracting update", 70)
            _safe_extract(zip_path, tmp)
            payload = _payload_root(tmp)
            if not os.path.exists(os.path.join(payload, "its-a-plane.py")):
                raise UpdateError("Update ZIP does not contain its-a-plane.py")

            if progress:
                progress("Backing up current code", 80)
            backup_path = _backup_current(installed)

            if progress:
                progress("Installing update", 90)
            _copy_payload(payload)

        _prune_backups(config["keep_backups"])
        state = _save_state(
            installed_version=manifest["version"],
            previous_version=installed,
            last_status="installed",
            last_install=_now_iso(),
            last_backup=backup_path,
            latest_version=manifest.get("version"),
            latest_manifest=manifest,
            needs_restart=True,
            pending_validation=True,
            validation_started=False,
            last_error=None,
        )
        if progress:
            progress("Update installed", 100)
        return state


def rollback(progress=None):
    with _operation_lock():
        state = _state()
        backup = state.get("last_backup")
        if not backup or not os.path.exists(backup):
            raise UpdateError("No rollback backup available")
        if progress:
            progress("Rolling back", 20)
        with tempfile.TemporaryDirectory(prefix="plane-rollback-") as tmp:
            _safe_extract(backup, tmp)
            if progress:
                progress("Restoring previous code", 70)
            _copy_payload(tmp)
        restored = state.get("previous_version") or "unknown"
        return _save_state(
            installed_version=restored,
            last_status="rolled back",
            last_rollback=_now_iso(),
            needs_restart=True,
            pending_validation=False,
            validation_started=False,
        )


def startup_update(progress=None):
    config = update_config()
    if not config.get("enabled") or not config.get("check_on_start", True):
        return {"ok": True, "installed": False, "message": "updates disabled"}
    try:
        if progress:
            progress("Checking for updates", 5)
        result = check_for_update()
        if not result.get("update_available"):
            return {"ok": True, "installed": False, "message": "already current"}
        if not config.get("auto_install", True):
            return {"ok": True, "installed": False, "message": "update available"}
        state = install_update(result["manifest"], progress=progress)
        return {"ok": True, "installed": True, "state": state}
    except Exception as exc:
        _save_state(last_status=f"update failed: {exc}", last_error=str(exc), last_checked=_now_iso())
        if progress:
            progress("Update failed, starting app", 100)
        return {"ok": False, "installed": False, "message": str(exc)}


def status():
    cfg = update_config()
    state = _state()
    return {
        "enabled": bool(cfg.get("enabled")),
        "check_on_start": bool(cfg.get("check_on_start", True)),
        "auto_install": bool(cfg.get("auto_install", True)),
        "manifest_url": cfg.get("manifest_url", ""),
        "installed_version": state.get("installed_version") or _read_installed_version(),
        "latest_version": state.get("latest_version"),
        "last_status": state.get("last_status"),
        "last_checked": state.get("last_checked"),
        "last_install": state.get("last_install"),
        "last_error": state.get("last_error"),
        "progress_percent": state.get("progress_percent"),
        "needs_restart": bool(state.get("needs_restart")),
        "rollback_available": bool(state.get("last_backup") and os.path.exists(state.get("last_backup"))),
    }


def mark_app_started():
    return _save_state(
        needs_restart=False,
        pending_validation=False,
        validation_started=False,
        last_status="running",
        last_app_start=_now_iso(),
        progress_percent=0,
        last_error=None,
    )


def record_progress(message, percent=None):
    updates = {"last_status": message}
    if percent is not None:
        updates["progress_percent"] = max(0, min(100, int(percent)))
    return _save_state(**updates)


def prepare_startup(progress=None):
    state = _state()
    if not state.get("pending_validation"):
        return {"rolled_back": False}
    if not state.get("validation_started"):
        _save_state(validation_started=True, last_status="validating installed update")
        return {"rolled_back": False, "validating": True}

    if progress:
        progress("Update failed, rolling back", 10)
    rollback(progress=progress)
    return {"rolled_back": True}
