#!/usr/bin/env python3

import json
import os
import signal
import sys
from pathlib import Path
from typing import Any, Dict, List, Tuple

# Keep HuggingFace/SentenceTransformer model loading deterministic in CLI jobs.
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "10")
os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "30")

MODEL_LOAD_TIMEOUT_SECONDS = int(os.environ.get("RETRIEX_EMBEDDING_MODEL_LOAD_TIMEOUT_SECONDS", "60"))


def log_event(event: str, **payload: Any) -> None:
    print(json.dumps({"event": event, **payload}, ensure_ascii=False), file=sys.stderr, flush=True)


def fail(message: str, code: int) -> None:
    print(f"ERROR: {message}", file=sys.stderr, flush=True)
    sys.exit(code)


# ---------------------------------------------------------
# Positional args
# 1 tags.ndjson
# 2 out_index_path (can be .tmp)
# ---------------------------------------------------------
if len(sys.argv) < 3:
    fail("usage: vector_ingest_tags.py <tags.ndjson> <out.index>", 2)

tags_path = Path(sys.argv[1]).resolve()
out_path = Path(sys.argv[2]).resolve()
meta_path = Path(str(out_path) + ".meta.json")


# ---------------------------------------------------------
# Dependency checks
# ---------------------------------------------------------
try:
    import faiss
except Exception:
    fail("Python module 'faiss' not found.", 10)

try:
    from sentence_transformers import SentenceTransformer
except Exception:
    fail("Python module 'sentence-transformers' not found.", 11)

import numpy as np


# ---------------------------------------------------------
# File checks
# ---------------------------------------------------------
if not tags_path.is_file():
    fail(f"tags.ndjson not found at {tags_path}", 20)

out_path.parent.mkdir(parents=True, exist_ok=True)


# ---------------------------------------------------------
# Helpers
# ---------------------------------------------------------
def cleanup_outputs() -> None:
    if out_path.exists():
        out_path.unlink()
    if meta_path.exists():
        meta_path.unlink()


def normalize_text(value: Any) -> str:
    text = str(value).strip()
    text = " ".join(text.split())

    if len(text) > 4000:
        text = text[:4000].rstrip()

    return text


def resolve_embedding_model_from_meta() -> str:
    # Local model path wins. This avoids implicit network/cache lookup in production.
    override = os.environ.get("RETRIEX_EMBEDDING_MODEL_PATH", "").strip()
    if override:
        return override

    base_path = Path(__file__).resolve().parents[2]
    index_meta_path = base_path / "var" / "knowledge" / "index_meta.json"

    if not index_meta_path.exists():
        fail("index_meta.json not found", 30)

    try:
        meta = json.loads(index_meta_path.read_text(encoding="utf-8"))
    except Exception:
        fail("index_meta.json is invalid", 30)

    embedding_model = meta.get("embedding_model")
    if not isinstance(embedding_model, str) or embedding_model.strip() == "":
        fail("embedding_model missing in index_meta.json", 31)

    return embedding_model.strip()


def load_sentence_transformer(model_name_or_path: str) -> SentenceTransformer:
    def timeout_handler(_signum: int, _frame: Any) -> None:
        raise TimeoutError(
            "Embedding model load timed out. "
            "Cache the model locally or set RETRIEX_EMBEDDING_MODEL_PATH."
        )

    log_event(
        "tag_embedding_model_load_start",
        model=model_name_or_path,
        timeout_seconds=MODEL_LOAD_TIMEOUT_SECONDS,
        hf_hub_disable_xet=os.environ.get("HF_HUB_DISABLE_XET"),
    )

    previous_handler = signal.signal(signal.SIGALRM, timeout_handler)
    signal.alarm(MODEL_LOAD_TIMEOUT_SECONDS)

    try:
        loaded_model = SentenceTransformer(model_name_or_path)
    except TimeoutError as exc:
        fail(str(exc), 32)
    except Exception as exc:
        fail(f"Unable to load embedding model '{model_name_or_path}': {exc}", 33)
    finally:
        signal.alarm(0)
        signal.signal(signal.SIGALRM, previous_handler)

    log_event("tag_embedding_model_load_done", model=model_name_or_path)
    return loaded_model


# ---------------------------------------------------------
# Streaming read NDJSON
# ---------------------------------------------------------
def load_rows(path: Path) -> Tuple[List[str], List[str], Dict[str, int]]:
    texts: List[str] = []
    ids: List[str] = []
    seen_ids = set()

    stats = {
        "lines_total": 0,
        "lines_empty": 0,
        "lines_invalid_json": 0,
        "rows_missing_fields": 0,
        "rows_duplicate_tag_id": 0,
        "rows_accepted": 0,
    }

    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            stats["lines_total"] += 1
            line = line.strip()

            if line == "":
                stats["lines_empty"] += 1
                continue

            try:
                entry = json.loads(line)
            except Exception:
                stats["lines_invalid_json"] += 1
                continue

            if not isinstance(entry, dict):
                stats["rows_missing_fields"] += 1
                continue

            tag_id = str(entry.get("tag_id", "")).strip()
            text = normalize_text(entry.get("text", ""))

            if tag_id == "" or text == "":
                stats["rows_missing_fields"] += 1
                continue

            if tag_id in seen_ids:
                stats["rows_duplicate_tag_id"] += 1
                continue

            seen_ids.add(tag_id)
            ids.append(tag_id)
            texts.append(f"passage: {text}")
            stats["rows_accepted"] += 1

    return texts, ids, stats


texts, ids, stats = load_rows(tags_path)
log_event("tag_rows_loaded", **stats)

if not texts:
    cleanup_outputs()
    sys.exit(0)


# ---------------------------------------------------------
# Load model only after we know that usable tags exist
# ---------------------------------------------------------
embedding_model = resolve_embedding_model_from_meta()
model = load_sentence_transformer(embedding_model)


# ---------------------------------------------------------
# Build embeddings
# ---------------------------------------------------------
log_event("tag_embedding_encode_start", rows=len(texts))

embeddings = model.encode(
    texts,
    normalize_embeddings=True,
    show_progress_bar=False,
    batch_size=128,
)

log_event("tag_embedding_encode_done", rows=len(texts))

embeddings = np.array(embeddings, dtype="float32")

if embeddings.ndim != 2 or embeddings.shape[0] != len(ids) or embeddings.shape[0] == 0:
    cleanup_outputs()
    fail("tag embeddings have invalid shape", 40)

if embeddings.shape[1] <= 0:
    cleanup_outputs()
    fail("tag embeddings have invalid dimension", 41)

dim = int(embeddings.shape[1])

index = faiss.IndexFlatIP(dim)
index.add(embeddings)

if int(index.ntotal) != len(ids):
    cleanup_outputs()
    fail("FAISS tag index count does not match meta ids", 42)

faiss.write_index(index, str(out_path))
meta_path.write_text(
    json.dumps(ids, ensure_ascii=False),
    encoding="utf-8",
)

log_event(
    "tag_vector_index_written",
    index=str(out_path),
    meta=str(meta_path),
    rows=len(ids),
    dimension=dim,
)

sys.exit(0)
