#!/usr/bin/env python3
"""
Stage 3 Worker — Graphiti Ingest with Taxonomy-Free Orientation
Polls stage_3_queue, chunks documents, ingests as episodic saga to Graphiti.

Chunking rationale: Large documents sent as single episodes cause FalkorDB
write lock contention during entity deduplication. Chunking at ~500 words
(matching Stage 1) produces smaller deduplication passes that don't block.
Each document's chunks are linked via Graphiti's saga mechanism, preserving
document structure in the graph.

Saga-size limit (MAX_CHUNKS_PER_SAGA): 2026-05-01 incident showed sagas of
17 and 19 chunks deadlock the sidecar's Python-side coordination. Documents
producing more than MAX_CHUNKS_PER_SAGA chunks are split into multiple bulk
commits, each tagged with the same saga value so Graphiti still links them.

Wedge detection: 2026-05-01 incident also surfaced the asymmetry with Stage 2 —
Stage 3 had no recovery path when the sidecar deadlocked. Now mirrors Stage 2's
consecutive_failures pattern with sidecar restart on threshold.

Runs as systemd service: aaronai-stage3.service
"""

import os, json, time, logging, subprocess, requests
from pathlib import Path
from datetime import datetime
from dotenv import load_dotenv
import psycopg2

load_dotenv(Path.home() / "aaronai" / ".env", override=True)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [stage3] %(levelname)s %(message)s",
    handlers=[
        logging.FileHandler("/var/log/aaronai/stage3.log", mode="a"),
    ]
)
log = logging.getLogger("stage3")

PG_DSN = os.getenv("PG_DSN")
GRAPHITI_URL = "http://localhost:8001"
HEARTBEAT_FILE = Path("/var/log/aaronai/stage3-heartbeat")
RETRY_ATTEMPTS = 2
POLL_INTERVAL = 5
INGEST_TIMEOUT = 600
WORKER_VERSION = "2.2"

# Match Stage 1 chunking parameters
CHUNK_SIZE_WORDS = 500
CHUNK_OVERLAP_WORDS = 50
# Documents under this threshold ingested as single episode (no chunking overhead)
SINGLE_EPISODE_THRESHOLD = 1500
# Sagas larger than this many chunks split into multiple commits
# (2026-05-01 incident: 17 and 19 chunk sagas deadlocked sidecar)
MAX_CHUNKS_PER_SAGA = 10


def get_pg():
    return psycopg2.connect(PG_DSN)


def write_heartbeat():
    try:
        HEARTBEAT_FILE.parent.mkdir(parents=True, exist_ok=True)
        HEARTBEAT_FILE.write_text(datetime.now().isoformat())
    except Exception:
        pass


def recover_wedge():
    """Restart Graphiti sidecar when consecutive failures suggest deadlock.
    Mirrors Stage 2's recover_wedge() for ollama. Requires passwordless sudo
    for `systemctl restart aaronai-graphiti.service` for the worker's user."""
    log.warning("Graphiti wedge detected — restarting sidecar")
    result = subprocess.run(
        ["/usr/bin/sudo", "/bin/systemctl", "restart", "aaronai-graphiti.service"],
        capture_output=True, text=True
    )
    if result.returncode != 0:
        log.error(f"Sidecar restart failed (rc={result.returncode}): stdout={result.stdout!r} stderr={result.stderr!r}")
        return False
    # Sidecar needs longer than ollama for model loading (sentence-transformers
    # + BGE reranker + Graphiti library init)
    time.sleep(45)
    for _ in range(3):
        try:
            r = requests.get(f"{GRAPHITI_URL}/health", timeout=10)
            if r.status_code == 200:
                log.info("Graphiti recovered")
                return True
        except Exception:
            time.sleep(10)
    log.error("Graphiti recovery failed")
    return False


def chunk_text(text, chunk_size=CHUNK_SIZE_WORDS, overlap=CHUNK_OVERLAP_WORDS):
    """Split text into word-based chunks matching Stage 1 chunking."""
    words = text.split()
    chunks = []
    start = 0
    while start < len(words):
        end = start + chunk_size
        chunk = " ".join(words[start:end])
        if chunk.strip():
            chunks.append(chunk)
        start += chunk_size - overlap
    return chunks


def post_bulk(payload, batch_label=""):
    """Single POST to /episodes/bulk with consistent error handling."""
    resp = requests.post(
        f"{GRAPHITI_URL}/episodes/bulk",
        json=payload,
        timeout=INGEST_TIMEOUT
    )
    if not resp.ok:
        prefix = f"{batch_label} " if batch_label else ""
        raise RuntimeError(f"{prefix}Sidecar {resp.status_code}: {resp.text[:500]}")
    return resp.json()


def ingest_to_graphiti(source, full_text, orientation):
    """
    Ingest document to Graphiti as chunked episodes linked by saga.

    Three paths:
    - Short documents (<SINGLE_EPISODE_THRESHOLD): single episode, no saga
    - Medium documents (chunks <= MAX_CHUNKS_PER_SAGA): one bulk commit, saga-linked
    - Large documents (chunks > MAX_CHUNKS_PER_SAGA): split into batches of
      MAX_CHUNKS_PER_SAGA, each its own bulk commit, all sharing the same saga tag
      so Graphiti links them as one document unit
    """
    char_length = len(full_text)

    if char_length < SINGLE_EPISODE_THRESHOLD:
        episodes = [{
            "name": source,
            "content": full_text,
            "source_description": orientation,
            "timestamp": datetime.now().isoformat(),
        }]
        log.info(f"  Single episode ({char_length} chars)")
        return post_bulk({"episodes": episodes, "group_id": "aaron"})

    chunks = chunk_text(full_text)
    total_chunks = len(chunks)

    if total_chunks <= MAX_CHUNKS_PER_SAGA:
        episodes = [
            {
                "name": f"{source} [{i+1}/{total_chunks}]",
                "content": chunk,
                "source_description": orientation,
                "timestamp": datetime.now().isoformat(),
            }
            for i, chunk in enumerate(chunks)
        ]
        log.info(f"  Chunked into {total_chunks} episodes ({char_length} chars)")
        return post_bulk(
            {"episodes": episodes, "group_id": "aaron", "saga": source}
        )

    # Large document: split into batches sharing the same saga tag
    batch_count = (total_chunks + MAX_CHUNKS_PER_SAGA - 1) // MAX_CHUNKS_PER_SAGA
    log.info(
        f"  Chunked into {total_chunks} episodes ({char_length} chars); "
        f"splitting into {batch_count} batches of up to {MAX_CHUNKS_PER_SAGA}"
    )
    last_result = None
    for batch_idx in range(batch_count):
        start = batch_idx * MAX_CHUNKS_PER_SAGA
        end = min(start + MAX_CHUNKS_PER_SAGA, total_chunks)
        batch_chunks = chunks[start:end]
        episodes = [
            {
                "name": f"{source} [{start + i + 1}/{total_chunks}]",
                "content": chunk,
                "source_description": orientation,
                "timestamp": datetime.now().isoformat(),
            }
            for i, chunk in enumerate(batch_chunks)
        ]
        batch_label = f"batch {batch_idx + 1}/{batch_count} (chunks {start + 1}-{end})"
        log.info(f"    {batch_label} starting")
        last_result = post_bulk(
            {"episodes": episodes, "group_id": "aaron", "saga": source},
            batch_label=batch_label,
        )
        log.info(f"    {batch_label} committed")
    return last_result


def process_one(row):
    row_id, source, full_text, orientation = row
    log.info(f"Ingesting to Graphiti: {source}")

    pg = get_pg()
    cur = pg.cursor()
    cur.execute(
        "UPDATE stage_3_queue SET started_at = NOW(), attempts = attempts + 1 WHERE id = %s",
        (row_id,)
    )
    pg.commit()

    try:
        result = ingest_to_graphiti(source, full_text, orientation)
        cur.execute("UPDATE stage_3_queue SET completed_at = NOW() WHERE id = %s", (row_id,))
        pg.commit()
        pg.close()
        log.info(f"  Committed to Graphiti: {source}")
        return True
    except Exception as e:
        log.error(f"  Graphiti ingest failed for {source}: {e}")
        cur.execute("""
            UPDATE stage_3_queue
            SET failed_at = NOW(), failure_reason = %s
            WHERE id = %s
        """, (str(e)[:500], row_id))
        pg.commit()
        pg.close()
        return False


def run():
    log.info(f"Stage 3 worker starting (v{WORKER_VERSION})")
    consecutive_failures = 0

    while True:
        write_heartbeat()

        try:
            pg = get_pg()
            cur = pg.cursor()
            cur.execute("""
                SELECT id, source, full_text, orientation
                FROM stage_3_queue
                WHERE completed_at IS NULL
                  AND failed_at IS NULL
                  AND (started_at IS NULL OR started_at < NOW() - INTERVAL '10 minutes')
                  AND attempts < %s
                ORDER BY enqueued_at ASC
                LIMIT 1
            """, (RETRY_ATTEMPTS + 1,))
            row = cur.fetchone()
            pg.close()

            if not row:
                consecutive_failures = 0
                time.sleep(POLL_INTERVAL)
                continue

            success = process_one(row)

            if not success:
                consecutive_failures += 1
                if consecutive_failures >= 2:
                    log.warning("Multiple consecutive failures — checking for Graphiti wedge")
                    recovered = recover_wedge()
                    if not recovered:
                        log.error("Wedge recovery failed — continuing anyway")
                    consecutive_failures = 0
                time.sleep(10)
            else:
                consecutive_failures = 0
                time.sleep(2)

        except Exception as e:
            log.error(f"Worker loop error: {e}")
            time.sleep(10)


if __name__ == "__main__":
    run()
