"""One-time schema creation + data ingestion for Pipeline 3 (TigerGraph GraphRAG).

Run once before starting the server:
    python core/pipeline_3/setup.py

Requires TigerGraph community container running on localhost:14240:
    docker run -d --init -p 14240:14240 --name tigergraph tigergraph/community:4.2.2

Safe to re-run after interruption — resumes from last completed checkpoint.
"""
import json
import logging
import os
import requests
from pathlib import Path
from typing import Optional

import numpy as np
import pyTigerGraph as tg
from dotenv import load_dotenv
from sentence_transformers import SentenceTransformer

load_dotenv()

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

# ── Config ─────────────────────────────────────────────────────────────────────
TG_HOST      = os.environ.get("TG_HOST", "http://localhost")
TG_PORT      = int(os.environ.get("TG_PORT", "14240"))
TG_USER      = os.environ.get("TG_USER", "tigergraph")
TG_PASS      = os.environ.get("TG_PASS", "tigergraph")
TG_SECRET    = os.environ.get("TG_SECRET", "")
GRAPH        = "PaperGraph"
EMBED_MODEL  = "BAAI/bge-large-en-v1.5"
EMBED_DIM    = 1024
EMBED_BATCH  = 64
UPSERT_BATCH = 256

_PROJECT_ROOT  = Path(__file__).parents[2]
_DATA_DIR      = _PROJECT_ROOT / "data"
_PIPELINE_DIR  = Path(__file__).parent
_EMBED_CACHE   = _PIPELINE_DIR / "embeddings.npy"
_EMBED_PARTIAL = _PIPELINE_DIR / "embed_partials"
_CHECKPOINT    = _PIPELINE_DIR / "checkpoint.json"


# ── REST client (bypasses pyTigerGraph auth issues on TG Cloud) ────────────────
class _TGClient:
    def __init__(self, token: str):
        self._url  = f"{TG_HOST}:{TG_PORT}/restpp/graph/{GRAPH}"
        self._session = requests.Session()
        self._session.headers["Authorization"] = f"Bearer {token}"

    def upsertVertices(self, vtype: str, vertices: list) -> None:
        body = {"vertices": {vtype: {
            v_id: {k: {"value": v} for k, v in attrs.items()}
            for v_id, attrs in vertices
        }}}
        resp = self._session.post(self._url, json=body, timeout=120)
        resp.raise_for_status()

    def upsertEdges(self, src_type: str, edge_type: str, _tgt_type: str, edges: list) -> None:
        src_map: dict = {}
        for src_id, tgt_id, attrs in edges:
            src_map.setdefault(src_id, {}).setdefault(edge_type, {})[tgt_id] = {
                k: {"value": v} for k, v in attrs.items()
            }
        body = {"edges": {src_type: src_map}}
        resp = self._session.post(self._url, json=body, timeout=120)
        resp.raise_for_status()


def _get_conn():
    if TG_SECRET:
        tg_conn = tg.TigerGraphConnection(
            host=TG_HOST, graphname=GRAPH, gsqlSecret=TG_SECRET,
            restppPort=str(TG_PORT), gsPort=str(TG_PORT),
        )
        token = tg_conn.getToken(TG_SECRET)[0]
        return _TGClient(token)
    return tg.TigerGraphConnection(
        host=TG_HOST, graphname=GRAPH, username=TG_USER, password=TG_PASS,
        restppPort=str(TG_PORT), gsPort=str(TG_PORT),
    )


# ── Checkpoint helpers ─────────────────────────────────────────────────────────
def _load_checkpoint() -> dict:
    if _CHECKPOINT.exists():
        ckpt = json.loads(_CHECKPOINT.read_text())
        ckpt.setdefault("embed_batches", 0)
        return ckpt
    return {
        "embed_batches": 0,
        "papers":        0,
        "authors":       False,
        "topics":        False,
        "authored_by":   False,
        "has_topic":     False,
        "cites":         False,
    }


def _save_checkpoint(ckpt: dict) -> None:
    _CHECKPOINT.write_text(json.dumps(ckpt, indent=2))


# ── GSQL helper (raw REST — confirmed working with TG 4.2.2) ──────────────────
def _gsql(stmt: str) -> str:
    resp = requests.post(
        f"{TG_HOST}:{TG_PORT}/gsql/v1/statements",
        data=stmt.encode("utf-8"),
        headers={"Content-Type": "text/plain"},
        auth=(TG_USER, TG_PASS),
        timeout=120,
    )
    return resp.text


def _graph_ready() -> bool:
    try:
        resp = requests.get(
            f"{TG_HOST}:{TG_PORT}/restpp/graph/{GRAPH}/vertices/Paper",
            params={"limit": 1},
            auth=(TG_USER, TG_PASS),
            timeout=10,
        )
        return resp.status_code == 200
    except Exception:
        return False


# ── Schema ─────────────────────────────────────────────────────────────────────
_SCHEMA = f"""
CREATE GRAPH {GRAPH}()
USE GRAPH {GRAPH}

CREATE SCHEMA_CHANGE JOB init_schema FOR GRAPH {GRAPH} {{
  ADD VERTEX Paper(
    PRIMARY_ID paper_id STRING,
    title STRING DEFAULT "",
    abstract STRING DEFAULT "",
    year INT DEFAULT 0,
    cited_by_count INT DEFAULT 0,
    doi STRING DEFAULT "",
    pdf_url STRING DEFAULT "",
    venue STRING DEFAULT "",
    corresponding_author STRING DEFAULT "",
    authors STRING DEFAULT "",
    funders STRING DEFAULT ""
  ) WITH primary_id_as_attribute="true";

  ADD VERTEX Author(
    PRIMARY_ID author_id STRING,
    display_name STRING DEFAULT ""
  ) WITH primary_id_as_attribute="true";

  ADD VERTEX Topic(
    PRIMARY_ID topic_id STRING,
    display_name STRING DEFAULT "",
    subfield STRING DEFAULT "",
    field STRING DEFAULT "",
    domain STRING DEFAULT ""
  ) WITH primary_id_as_attribute="true";

  ADD DIRECTED EDGE AUTHORED_BY(FROM Paper, TO Author);
  ADD DIRECTED EDGE HAS_TOPIC(FROM Paper, TO Topic, score FLOAT DEFAULT 0.0);
  ADD DIRECTED EDGE CITES(FROM Paper, TO Paper);
}}
RUN SCHEMA_CHANGE JOB init_schema
DROP JOB init_schema

CREATE SCHEMA_CHANGE JOB add_embedding FOR GRAPH {GRAPH} {{
  ALTER VERTEX Paper ADD VECTOR ATTRIBUTE embedding(dimension={EMBED_DIM});
}}
RUN SCHEMA_CHANGE JOB add_embedding
DROP JOB add_embedding
"""


def create_schema() -> None:
    if _graph_ready():
        logger.info("Graph already exists — skipping schema creation.")
        return
    logger.info("Creating graph schema…")
    out = _gsql(_SCHEMA)
    logger.info(f"Schema output:\n{out[:800]}")
    if "error" in out.lower() or "fail" in out.lower():
        raise RuntimeError(f"Schema creation failed:\n{out}")


# ── Embedding with per-batch checkpointing ─────────────────────────────────────
def _embed_papers(papers: list[dict], ckpt: dict) -> np.ndarray:
    if _EMBED_CACHE.exists():
        logger.info(f"Loading full embedding cache from {_EMBED_CACHE}…")
        embs = np.load(_EMBED_CACHE)
        logger.info(f"Loaded embeddings shape: {embs.shape}")
        return embs

    _EMBED_PARTIAL.mkdir(exist_ok=True)
    texts     = [f"Title: {p['title']}\n\nAbstract: {p['abstract']}" for p in papers]
    n_batches = (len(texts) + EMBED_BATCH - 1) // EMBED_BATCH
    completed = ckpt.get("embed_batches", 0)

    logger.info(f"Loading embedding model {EMBED_MODEL}…")
    model = SentenceTransformer(EMBED_MODEL)
    logger.info(f"Embedding: {completed}/{n_batches} batches already done, resuming…")

    for i in range(completed, n_batches):
        start      = i * EMBED_BATCH
        batch_embs = model.encode(
            texts[start : start + EMBED_BATCH],
            normalize_embeddings=True,
        )
        np.save(_EMBED_PARTIAL / f"batch_{i:04d}.npy", batch_embs)
        ckpt["embed_batches"] = i + 1
        _save_checkpoint(ckpt)
        if (i + 1) % 10 == 0 or i + 1 == n_batches:
            logger.info(f"  Embedded {i + 1}/{n_batches} batches ✓")

    # Merge partials into single cache file
    all_embs   = [np.load(_EMBED_PARTIAL / f"batch_{i:04d}.npy") for i in range(n_batches)]
    embeddings = np.vstack(all_embs)
    np.save(_EMBED_CACHE, embeddings)
    logger.info(f"Full embeddings saved → {_EMBED_CACHE}")

    for f in sorted(_EMBED_PARTIAL.glob("*.npy")):
        f.unlink()
    _EMBED_PARTIAL.rmdir()
    return embeddings


# ── Data parsing ───────────────────────────────────────────────────────────────
def _reconstruct_abstract(inverted: dict) -> str:
    if not inverted:
        return ""
    max_idx = max((max(v) for v in inverted.values() if v), default=0)
    words   = [""] * (max_idx + 1)
    for word, positions in inverted.items():
        for i in positions:
            words[i] = word
    return " ".join(words).strip()


def _parse_paper(path: Path) -> Optional[dict]:
    try:
        d        = json.loads(path.read_text())
        paper_id = path.stem

        abstract = _reconstruct_abstract(d.get("abstract_inverted_index") or {})

        loc    = d.get("primary_location") or {}
        source = loc.get("source") or {}
        venue  = source.get("display_name", "") or ""

        authorships  = d.get("authorships") or []
        author_names = [
            a["author"]["display_name"]
            for a in authorships[:5]
            if (a.get("author") or {}).get("display_name")
        ]
        corresponding_author = next(
            (
                a["author"].get("display_name", "")
                for a in authorships
                if a.get("is_corresponding") and a.get("author")
            ),
            author_names[0] if author_names else "",
        )

        funder_names = list({
            a["funder_display_name"]
            for a in (d.get("awards") or [])
            if a.get("funder_display_name")
        })

        best_oa = d.get("best_oa_location") or {}

        topics = [
            {
                "topic_id":     t["id"].split("/")[-1],
                "display_name": t.get("display_name", ""),
                "subfield":     (t.get("subfield") or {}).get("display_name", ""),
                "field":        (t.get("field")    or {}).get("display_name", ""),
                "domain":       (t.get("domain")   or {}).get("display_name", ""),
                "score":        float(t.get("score", 0.0)),
            }
            for t in (d.get("topics") or [])[:3]
            if t.get("id")
        ]

        author_vertices = [
            {
                "author_id":    a["author"]["id"].split("/")[-1],
                "display_name": a["author"].get("display_name", ""),
            }
            for a in authorships
            if (a.get("author") or {}).get("id")
        ]

        referenced = [r.split("/")[-1] for r in (d.get("referenced_works") or [])]

        return {
            "paper_id":             paper_id,
            "title":                d.get("title", "") or "",
            "abstract":             abstract,
            "year":                 d.get("publication_year") or 0,
            "cited_by_count":       d.get("cited_by_count", 0) or 0,
            "doi":                  d.get("doi", "") or "",
            "pdf_url":              best_oa.get("pdf_url", "") or "",
            "venue":                venue,
            "corresponding_author": corresponding_author,
            "authors":              ", ".join(author_names),
            "funders":              ", ".join(funder_names),
            "topics":               topics,
            "author_vertices":      author_vertices,
            "referenced":           referenced,
        }
    except Exception as e:
        logger.warning(f"Skipping {path.name}: {e}")
        return None


# ── Ingestion ──────────────────────────────────────────────────────────────────
def ingest(data_dir: Path = _DATA_DIR) -> None:
    ckpt = _load_checkpoint()
    logger.info(f"Resuming from checkpoint: {ckpt}")

    files      = sorted(data_dir.glob("*.json"))
    corpus_ids = {f.stem for f in files}
    logger.info(f"Found {len(files)} paper files.")

    papers = [p for f in files if (p := _parse_paper(f))]
    logger.info(f"Parsed {len(papers)} papers.")

    embeddings    = _embed_papers(papers, ckpt)
    conn          = _get_conn()
    total_batches = (len(papers) + UPSERT_BATCH - 1) // UPSERT_BATCH

    # Paper vertices ───────────────────────────────────────────────────────────
    completed = ckpt["papers"]
    if completed < total_batches:
        logger.info(f"Upserting Paper vertices — resuming from batch {completed}/{total_batches}…")
        for batch_idx in range(completed, total_batches):
            start = batch_idx * UPSERT_BATCH
            chunk = papers[start : start + UPSERT_BATCH]
            embs  = embeddings[start : start + UPSERT_BATCH]
            conn.upsertVertices("Paper", [
                (p["paper_id"], {
                    "title":                p["title"],
                    "abstract":             p["abstract"],
                    "year":                 p["year"],
                    "cited_by_count":       p["cited_by_count"],
                    "doi":                  p["doi"],
                    "pdf_url":              p["pdf_url"],
                    "venue":                p["venue"],
                    "corresponding_author": p["corresponding_author"],
                    "authors":              p["authors"],
                    "funders":              p["funders"],
                    "embedding":            emb.tolist(),
                })
                for p, emb in zip(chunk, embs)
            ])
            ckpt["papers"] = batch_idx + 1
            _save_checkpoint(ckpt)
            logger.info(f"  Papers batch {batch_idx + 1}/{total_batches} ✓")
    else:
        logger.info("Paper vertices already complete — skipping.")

    # Collect Author + Topic data ──────────────────────────────────────────────
    all_authors: dict[str, dict] = {}
    all_topics:  dict[str, dict] = {}
    for p in papers:
        for av in p["author_vertices"]:
            all_authors.setdefault(av["author_id"], av)
        for t in p["topics"]:
            all_topics.setdefault(t["topic_id"], t)

    # Author vertices ──────────────────────────────────────────────────────────
    if not ckpt["authors"]:
        author_list = list(all_authors.values())
        logger.info(f"Upserting {len(author_list)} Author vertices…")
        for start in range(0, len(author_list), UPSERT_BATCH):
            chunk = author_list[start : start + UPSERT_BATCH]
            conn.upsertVertices("Author", [
                (a["author_id"], {"display_name": a["display_name"]})
                for a in chunk
            ])
        ckpt["authors"] = True
        _save_checkpoint(ckpt)
        logger.info("Author vertices ✓")
    else:
        logger.info("Author vertices already complete — skipping.")

    # Topic vertices ───────────────────────────────────────────────────────────
    if not ckpt["topics"]:
        topic_list = list(all_topics.values())
        logger.info(f"Upserting {len(topic_list)} Topic vertices…")
        for start in range(0, len(topic_list), UPSERT_BATCH):
            chunk = topic_list[start : start + UPSERT_BATCH]
            conn.upsertVertices("Topic", [
                (t["topic_id"], {
                    "display_name": t["display_name"],
                    "subfield":     t["subfield"],
                    "field":        t["field"],
                    "domain":       t["domain"],
                })
                for t in chunk
            ])
        ckpt["topics"] = True
        _save_checkpoint(ckpt)
        logger.info("Topic vertices ✓")
    else:
        logger.info("Topic vertices already complete — skipping.")

    # AUTHORED_BY edges ────────────────────────────────────────────────────────
    if not ckpt["authored_by"]:
        authored_buf = [
            (p["paper_id"], av["author_id"], {})
            for p in papers
            for av in p["author_vertices"]
            if av["author_id"]
        ]
        logger.info(f"Upserting {len(authored_buf)} AUTHORED_BY edges…")
        for start in range(0, len(authored_buf), UPSERT_BATCH):
            conn.upsertEdges("Paper", "AUTHORED_BY", "Author",
                             authored_buf[start : start + UPSERT_BATCH])
        ckpt["authored_by"] = True
        _save_checkpoint(ckpt)
        logger.info("AUTHORED_BY edges ✓")
    else:
        logger.info("AUTHORED_BY edges already complete — skipping.")

    # HAS_TOPIC edges ──────────────────────────────────────────────────────────
    if not ckpt["has_topic"]:
        topic_buf = [
            (p["paper_id"], t["topic_id"], {"score": t["score"]})
            for p in papers
            for t in p["topics"]
            if t["topic_id"]
        ]
        logger.info(f"Upserting {len(topic_buf)} HAS_TOPIC edges…")
        for start in range(0, len(topic_buf), UPSERT_BATCH):
            conn.upsertEdges("Paper", "HAS_TOPIC", "Topic",
                             topic_buf[start : start + UPSERT_BATCH])
        ckpt["has_topic"] = True
        _save_checkpoint(ckpt)
        logger.info("HAS_TOPIC edges ✓")
    else:
        logger.info("HAS_TOPIC edges already complete — skipping.")

    # CITES edges (in-corpus only) ─────────────────────────────────────────────
    if not ckpt["cites"]:
        cites_buf = [
            (p["paper_id"], ref_id, {})
            for p in papers
            for ref_id in p["referenced"]
            if ref_id in corpus_ids
        ]
        logger.info(f"Upserting {len(cites_buf)} CITES edges (in-corpus)…")
        for start in range(0, len(cites_buf), UPSERT_BATCH):
            conn.upsertEdges("Paper", "CITES", "Paper",
                             cites_buf[start : start + UPSERT_BATCH])
        ckpt["cites"] = True
        _save_checkpoint(ckpt)
        logger.info("CITES edges ✓")
    else:
        logger.info("CITES edges already complete — skipping.")

    logger.info("Ingestion complete.")


if __name__ == "__main__":
    create_schema()
    ingest()