#!/usr/bin/env python3

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


def fail(message: str, code: int) -> None:
    print(f"ERROR: {message}", file=sys.stderr)
    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


# ---------------------------------------------------------
# Load embedding model from index_meta.json (Single Source of Truth)
# ---------------------------------------------------------
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)

model = SentenceTransformer(embedding_model.strip())


# ---------------------------------------------------------
# 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


# ---------------------------------------------------------
# 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)

print(
    json.dumps(
        {
            "event": "tag_rows_loaded",
            **stats,
        },
        ensure_ascii=False,
    ),
    file=sys.stderr,
)

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


# ---------------------------------------------------------
# Build embeddings
# ---------------------------------------------------------
embeddings = model.encode(
    texts,
    normalize_embeddings=True,
    show_progress_bar=True,
    batch_size=128,
)

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",
)

sys.exit(0)