"""
BM25 pickle checkpoint: corruption detection, backups, atomic save, background rebuild.
"""
from __future__ import annotations

import glob
import logging
import os
import pickle
import shutil
import threading
import time
from contextlib import contextmanager
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple

from app.core.config import settings

_logger = logging.getLogger("askme.bm25.checkpoint")

# Intentional clear_all: do not restore from .bak on next load
_skip_restore = threading.local()

_REBUILD_LOCK = threading.Lock()
_REBUILD_STATE = {
    "status": "idle",
    "recovered_from": "none",
    "chunks_processed": 0,
    "chunks_total": 0,
    "last_error": None,
    "started_at": None,
    "finished_at": None,
    "pkl_path": None,
}


def set_skip_restore(value: bool) -> None:
    _skip_restore.flag = bool(value)


def should_skip_restore() -> bool:
    return bool(getattr(_skip_restore, "flag", False))


def chroma_path_from_pkl(pkl_path: str) -> str:
    if pkl_path.endswith("_bm25.pkl"):
        return pkl_path[: -len("_bm25.pkl")]
    return os.path.splitext(pkl_path)[0]


def pkl_path_for_chroma(chroma_path: str) -> str:
    return f"{chroma_path}_bm25.pkl"


def checkpoint_enabled() -> bool:
    return bool(getattr(settings, "BM25_CHECKPOINT_ENABLED", True))


def validate_pkl_data(data: Any) -> Tuple[bool, str]:
    if not isinstance(data, dict):
        return False, f"expected dict, got {type(data).__name__}"
    chunk_ids = list(data.get("chunk_ids", []))
    documents = list(data.get("documents", []))
    metadatas = list(data.get("metadatas", []))
    tokenized = list(data.get("tokenized_corpus", []))
    n = len(chunk_ids)
    if not (len(documents) == n and len(metadatas) == n and len(tokenized) == n):
        return (
            False,
            f"length mismatch: chunk_ids={n}, documents={len(documents)}, "
            f"metadatas={len(metadatas)}, tokenized_corpus={len(tokenized)}",
        )
    return True, ""


def is_pkl_corrupted(path: str) -> Tuple[bool, str]:
    """Return (is_corrupt, reason). Missing file is not corrupt."""
    if not path:
        return False, ""
    if not os.path.exists(path):
        return False, ""
    try:
        if os.path.getsize(path) == 0:
            return True, "zero-byte file"
    except OSError as exc:
        return True, str(exc)

    try:
        with open(path, "rb") as f:
            data = pickle.load(f)
    except EOFError:
        return True, "truncated pickle (EOFError)"
    except pickle.UnpicklingError as exc:
        return True, f"unpickling error: {exc}"
    except Exception as exc:
        return True, f"load error: {exc}"

    ok, msg = validate_pkl_data(data)
    if not ok:
        return True, msg
    return False, ""


def load_pkl_dict(path: str) -> Dict[str, Any]:
    with open(path, "rb") as f:
        data = pickle.load(f)
    ok, msg = validate_pkl_data(data)
    if not ok:
        raise ValueError(msg)
    return data


def _backup_glob(path: str) -> str:
    return f"{path}.bak.*"


def list_backup_paths(path: str) -> List[str]:
    pattern = _backup_glob(path)
    files = glob.glob(pattern)
    files.sort(key=lambda p: os.path.getmtime(p) if os.path.exists(p) else 0, reverse=True)
    return files


def backup_pkl(path: str) -> Optional[str]:
    if not path or not os.path.exists(path):
        return None
    corrupt, _ = is_pkl_corrupted(path)
    if corrupt:
        return None
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    dest = f"{path}.bak.{ts}"
    shutil.copy2(path, dest)
    keep = max(1, int(getattr(settings, "BM25_BACKUP_KEEP", 3) or 3))
    backups = list_backup_paths(path)
    for old in backups[keep:]:
        try:
            os.remove(old)
        except OSError as exc:
            _logger.warning("Could not prune backup %s: %s", old, exc)
    return dest


def quarantine_pkl(path: str) -> Optional[str]:
    if not path or not os.path.exists(path):
        return None
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    dest = f"{path}.bad.{ts}"
    try:
        os.replace(path, dest)
        _logger.warning("Quarantined BM25 pickle: %s -> %s", path, dest)
        return dest
    except OSError as exc:
        _logger.warning("Could not quarantine %s: %s", path, exc)
        try:
            os.remove(path)
        except OSError:
            pass
        return None


def prune_all_backups(path: str) -> None:
    for p in list_backup_paths(path):
        try:
            os.remove(p)
        except OSError:
            pass


@contextmanager
def pkl_lock(path: str, timeout: float = 30.0):
    lock_path = f"{path}.lock"
    deadline = time.time() + timeout
    acquired = False
    while time.time() < deadline:
        try:
            fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
            try:
                os.write(fd, str(os.getpid()).encode())
            finally:
                os.close(fd)
            acquired = True
            break
        except FileExistsError:
            time.sleep(0.05)
        except OSError:
            time.sleep(0.05)
    if not acquired:
        raise TimeoutError(f"Could not acquire BM25 lock within {timeout}s: {lock_path}")
    try:
        yield
    finally:
        try:
            os.remove(lock_path)
        except OSError:
            pass


def atomic_save_pkl(path: str, data: Dict[str, Any]) -> None:
    directory = os.path.dirname(path) or "."
    os.makedirs(directory, exist_ok=True)
    tmp = os.path.join(directory, f".{os.path.basename(path)}.tmp.{os.getpid()}")
    try:
        with open(tmp, "wb") as f:
            pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
        os.replace(tmp, path)
    finally:
        if os.path.exists(tmp):
            try:
                os.remove(tmp)
            except OSError:
                pass


def _update_rebuild_state(**kwargs: Any) -> None:
    with _REBUILD_LOCK:
        _REBUILD_STATE.update(kwargs)


def get_bm25_status(pkl_path: Optional[str] = None) -> Dict[str, Any]:
    path = pkl_path or _REBUILD_STATE.get("pkl_path")
    corrupt = False
    corrupt_reason = ""
    file_size = 0
    if path and os.path.exists(path):
        file_size = os.path.getsize(path)
        corrupt, corrupt_reason = is_pkl_corrupted(path)

    with _REBUILD_LOCK:
        st = dict(_REBUILD_STATE)

    total = int(st.get("chunks_total") or 0)
    done = int(st.get("chunks_processed") or 0)
    progress_pct = round(100.0 * done / total, 1) if total > 0 else 0.0

    return {
        "status": st.get("status", "idle"),
        "recovered_from": st.get("recovered_from", "none"),
        "chunks_processed": done,
        "chunks_total": total,
        "progress_pct": progress_pct,
        "last_error": st.get("last_error"),
        "started_at": st.get("started_at"),
        "finished_at": st.get("finished_at"),
        "pkl_path": path,
        "pkl_size_bytes": file_size,
        "pkl_corrupt": corrupt,
        "pkl_corrupt_reason": corrupt_reason if corrupt else None,
        "message": _status_message(st),
    }


def _status_message(st: Dict[str, Any]) -> str:
    status = st.get("status", "idle")
    if status == "rebuilding":
        return "Sparse index rebuild in progress"
    if status == "failed":
        return st.get("last_error") or "Sparse index rebuild failed"
    if status == "idle" and st.get("recovered_from") == "chroma":
        return "Sparse index ready (recovered from Chroma)"
    return "Sparse index ready"


def schedule_background_rebuild(
    chroma_path: str,
    openai_api_key: str,
    *,
    pkl_path: Optional[str] = None,
) -> bool:
    """Start daemon rebuild if not already running. Returns True if started."""
    if not checkpoint_enabled():
        return False
    if not openai_api_key or not str(openai_api_key).strip():
        return False

    path = pkl_path or pkl_path_for_chroma(chroma_path)
    with _REBUILD_LOCK:
        if _REBUILD_STATE.get("status") == "rebuilding":
            return False

    _update_rebuild_state(
        status="rebuilding",
        recovered_from="none",
        chunks_processed=0,
        chunks_total=0,
        last_error=None,
        started_at=datetime.now().isoformat(),
        finished_at=None,
        pkl_path=path,
    )

    thread = threading.Thread(
        target=_run_background_rebuild,
        args=(chroma_path, openai_api_key, path),
        name="bm25-rebuild",
        daemon=True,
    )
    thread.start()
    return True


def _run_background_rebuild(chroma_path: str, openai_api_key: str, pkl_path: str) -> None:
    try:
        from app.utils.chroma_vector_db import ChromaVectorDB
        from app.services.document_nodes import BM25Index

        vector_db = ChromaVectorDB(chroma_path, openai_api_key)
        total = vector_db.count() if vector_db else 0
        _update_rebuild_state(chunks_total=total)

        if total <= 0:
            _update_rebuild_state(
                status="failed",
                last_error="Chroma collection empty; re-ingest documents required",
                finished_at=datetime.now().isoformat(),
            )
            _logger.error("bm25_recovered_from=empty chroma_path=%s", chroma_path)
            return

        batch_size = max(50, int(getattr(settings, "BM25_REBUILD_BATCH_SIZE", 500) or 500))
        model_every = max(1, int(getattr(settings, "BM25_REBUILD_MODEL_EVERY", 10) or 10))

        index = BM25Index(
            pkl_path,
            chroma_path=chroma_path,
            openai_api_key=openai_api_key,
            skip_load=True,
        )
        processed = index.rebuild_from_chroma(
            vector_db,
            batch_size=batch_size,
            model_every=model_every,
            progress_callback=lambda n, t: _update_rebuild_state(
                chunks_processed=n, chunks_total=t
            ),
        )

        _update_rebuild_state(
            status="idle",
            recovered_from="chroma",
            chunks_processed=processed,
            chunks_total=total,
            finished_at=datetime.now().isoformat(),
            last_error=None,
        )
        _logger.info(
            "bm25_recovered_from=chroma chroma_path=%s chunks=%s",
            chroma_path,
            processed,
        )
    except Exception as exc:
        _update_rebuild_state(
            status="failed",
            last_error=str(exc),
            finished_at=datetime.now().isoformat(),
        )
        _logger.exception("bm25 rebuild failed: %s", exc)


def check_startup_and_maybe_schedule(
    chroma_path: str = "file_vectors_db",
    openai_api_key: Optional[str] = None,
) -> None:
    """Non-blocking: if main pkl corrupt, optionally schedule background rebuild."""
    if not checkpoint_enabled():
        return
    api_key = (openai_api_key or getattr(settings, "OPENAI_API_KEY", "") or "").strip()
    if not api_key:
        return
    path = pkl_path_for_chroma(chroma_path)
    if not os.path.exists(path):
        return
    corrupt, reason = is_pkl_corrupted(path)
    if not corrupt:
        return
    _logger.warning("BM25 pickle corrupt at startup: %s (%s)", path, reason)
    if getattr(settings, "BM25_AUTO_REBUILD", False):
        schedule_background_rebuild(chroma_path, api_key, pkl_path=path)
