"""
step5_embed_wardiary.py
───────────────────────
Reads step3/step4 JSON output files from a TestRun (or any outputs folder) and
embeds them into the pgvector document_chunks table.

Sources ingested:
  *_narratives.json     → chunk_type=narrative   (one chunk per diary entry)
  *_appx_records.json   → chunk_type=appx_record (one chunk per document page)

Positions JSON (*_positions.json, *_appx_positions.json) are NOT embedded here —
they belong in a PostGIS geometry table, not a vector search table.

Usage:
  python scripts/step5_embed_wardiary.py --input_dir TestRun/outputs --nationality canadian
  python scripts/step5_embed_wardiary.py --input_dir TestRun/outputs --nationality canadian --dry_run
"""

import os
import sys
import json
import argparse
from pathlib import Path
from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()

# ── Embedding config (same as GoogleDocumentOCR/embed_and_store.py) ──────────
EMBEDDING_MODEL = "BAAI/bge-base-en-v1.5"
EMBEDDING_DIMS  = 768
BATCH_SIZE      = 100

client = OpenAI(
    api_key=os.getenv("DEEPINFRA_API_KEY"),
    base_url="https://api.deepinfra.com/v1/openai",
)

# ── DB (reuse GoogleDocumentOCR/db.py) ───────────────────────────────────────
sys.path.insert(0, str(Path(__file__).parent.parent / "GoogleDocumentOCR"))
from db import init_db, get_conn
from psycopg2.extras import execute_values


# ─────────────────────────────────────────────────────────────────────────────
# Embedding
# ─────────────────────────────────────────────────────────────────────────────

def embed_texts(texts: list[str]) -> list[list[float]]:
    all_embeddings = []
    for i in range(0, len(texts), BATCH_SIZE):
        batch = texts[i:i + BATCH_SIZE]
        response = client.embeddings.create(input=batch, model=EMBEDDING_MODEL)
        all_embeddings.extend([item.embedding for item in response.data])
    return all_embeddings


# ─────────────────────────────────────────────────────────────────────────────
# DB insert (extended schema with chunk_type, date_entry, unit)
# ─────────────────────────────────────────────────────────────────────────────

def init_extended_schema():
    """Add chunk_type, date_entry, unit columns if they don't exist yet."""
    conn = get_conn()
    cur  = conn.cursor()
    for col, typedef in [
        ("chunk_type",  "TEXT DEFAULT 'unknown'"),
        ("date_entry",  "TEXT"),
        ("unit",        "TEXT"),
    ]:
        cur.execute(f"ALTER TABLE document_chunks ADD COLUMN IF NOT EXISTS {col} {typedef};")
    conn.commit()
    cur.close()
    conn.close()


def upsert_rows(rows: list[dict]):
    conn = get_conn()
    cur  = conn.cursor()
    execute_values(
        cur,
        """
        INSERT INTO document_chunks
            (filename, page_num, chunk_index, text, embedding,
             nationality, corpus, chunk_type, date_entry, unit)
        VALUES %s
        """,
        [
            (
                r["filename"], r["page_num"], r["chunk_index"], r["text"], r["embedding"],
                r["nationality"], r["corpus"], r["chunk_type"],
                r.get("date_entry"), r.get("unit"),
            )
            for r in rows
        ],
        template="(%s, %s, %s, %s, %s::vector, %s, %s, %s, %s, %s)"
    )
    conn.commit()
    cur.close()
    conn.close()


# ─────────────────────────────────────────────────────────────────────────────
# Text builders — turn each JSON record into a single embeddable string
# ─────────────────────────────────────────────────────────────────────────────

def narrative_to_text(entry: dict) -> str:
    parts = []
    if entry.get("date"):
        parts.append(f"Date: {entry['date']}")
    if entry.get("place"):
        parts.append(f"Location: {entry['place']}")
    if entry.get("weather"):
        parts.append(f"Weather: {entry['weather']}")
    if entry.get("summary"):
        parts.append(entry["summary"])
    return "\n".join(parts).strip()


def appx_record_to_text(record: dict) -> str:
    parts = []
    if record.get("page_type"):
        parts.append(f"Document type: {record['page_type']}")
    if record.get("document_id"):
        parts.append(f"Document ID: {record['document_id']}")
    if record.get("date"):
        parts.append(f"Date: {record['date']}")
    if record.get("time"):
        parts.append(f"Time: {record['time']}")
    if record.get("from_unit"):
        parts.append(f"From: {record['from_unit']}")
    if record.get("to_units"):
        parts.append(f"To: {', '.join(record['to_units'])}")
    if record.get("subject"):
        parts.append(f"Subject: {record['subject']}")
    if record.get("summary"):
        parts.append(record["summary"])
    if record.get("units_mentioned"):
        parts.append(f"Units mentioned: {', '.join(record['units_mentioned'])}")
    return "\n".join(parts).strip()


# ─────────────────────────────────────────────────────────────────────────────
# Ingest functions
# ─────────────────────────────────────────────────────────────────────────────

def ingest_narratives(json_path: Path, nationality: str, unit: str, dry_run: bool):
    data = json.loads(json_path.read_text(encoding="utf-8"))
    print(f"  {json_path.name}: {len(data)} narrative entries")

    chunks = []
    for idx, entry in enumerate(data):
        text = narrative_to_text(entry)
        if not text:
            continue
        chunks.append({
            "filename":    json_path.name,
            "page_num":    entry.get("source_page"),
            "chunk_index": idx,
            "text":        text,
            "nationality": nationality,
            "corpus":      "war_diary_narrative",
            "chunk_type":  "narrative",
            "date_entry":  entry.get("date"),
            "unit":        unit,
        })

    if not chunks:
        print("    No embeddable entries.")
        return

    if dry_run:
        print(f"    DRY RUN — would embed+store {len(chunks)} chunks")
        print(f"    Sample: {chunks[0]['text'][:120]}...")
        return

    print(f"    Embedding {len(chunks)} chunks...")
    embeddings = embed_texts([c["text"] for c in chunks])
    for c, emb in zip(chunks, embeddings):
        c["embedding"] = emb

    upsert_rows(chunks)
    print(f"    ✓ Stored {len(chunks)} narrative chunks")


def ingest_appx_records(json_path: Path, nationality: str, unit: str, dry_run: bool):
    data = json.loads(json_path.read_text(encoding="utf-8"))
    print(f"  {json_path.name}: {len(data)} appendix records")

    chunks = []
    for idx, record in enumerate(data):
        text = appx_record_to_text(record)
        if not text:
            continue
        chunks.append({
            "filename":    json_path.name,
            "page_num":    record.get("source_page"),
            "chunk_index": idx,
            "text":        text,
            "nationality": nationality,
            "corpus":      "war_diary_appendix",
            "chunk_type":  "appx_record",
            "date_entry":  record.get("date"),
            "unit":        unit,
        })

    if not chunks:
        print("    No embeddable records.")
        return

    if dry_run:
        print(f"    DRY RUN — would embed+store {len(chunks)} chunks")
        print(f"    Sample: {chunks[0]['text'][:120]}...")
        return

    print(f"    Embedding {len(chunks)} chunks...")
    embeddings = embed_texts([c["text"] for c in chunks])
    for c, emb in zip(chunks, embeddings):
        c["embedding"] = emb

    upsert_rows(chunks)
    print(f"    ✓ Stored {len(chunks)} appendix record chunks")


# ─────────────────────────────────────────────────────────────────────────────
# Unit name extraction from filename
# ─────────────────────────────────────────────────────────────────────────────

UNIT_PATTERNS = {
    "calgary-highlanders":  "Calgary Highlanders",
    "calgary_highlanders":  "Calgary Highlanders",
    "blackwatch":           "Black Watch (RHC)",
    "black-watch":          "Black Watch (RHC)",
    "5cib":                 "5 Canadian Infantry Brigade",
}

def extract_unit(filename: str) -> str:
    lower = filename.lower()
    for key, name in UNIT_PATTERNS.items():
        if key in lower:
            return name
    return "unknown"


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Embed war diary JSON outputs into pgvector")
    parser.add_argument("--input_dir",   required=True,  help="Folder containing step3/step4 JSON files")
    parser.add_argument("--nationality", default="canadian", help="Nationality tag (default: canadian)")
    parser.add_argument("--dry_run",     action="store_true", help="Print what would be done without writing to DB")
    args = parser.parse_args()

    input_dir = Path(args.input_dir)
    if not input_dir.exists():
        print(f"ERROR: input_dir not found: {input_dir}")
        sys.exit(1)

    if not args.dry_run:
        init_db()
        init_extended_schema()

    # Find all JSON files that are narratives or appx_records (skip positions)
    narrative_files    = sorted(input_dir.glob("*_narratives.json"))
    appx_record_files  = sorted(input_dir.glob("*_appx_records.json"))

    skipped = [f.name for f in input_dir.glob("*.json")
               if f not in narrative_files + appx_record_files]

    print(f"\n{'DRY RUN — ' if args.dry_run else ''}War Diary Embedding Ingest")
    print(f"Input dir : {input_dir}")
    print(f"Nationality: {args.nationality}")
    print(f"Narratives : {len(narrative_files)} file(s)")
    print(f"Appx records: {len(appx_record_files)} file(s)")
    if skipped:
        print(f"Skipping (positions/other): {', '.join(skipped)}\n")

    for f in narrative_files:
        unit = extract_unit(f.name)
        print(f"\n[narrative] {f.name} -> unit: {unit}")
        ingest_narratives(f, args.nationality, unit, args.dry_run)

    for f in appx_record_files:
        unit = extract_unit(f.name)
        print(f"\n[appx_record] {f.name} -> unit: {unit}")
        ingest_appx_records(f, args.nationality, unit, args.dry_run)

    print("\nDone.")


if __name__ == "__main__":
    main()



