#!/usr/bin/env python3
"""
Stage 3 Worker — Graphiti Ingest with Bulk-vs-Single-Episode Routing
                  + Encoder Instructions (v1.0)

Polls stage_3_queue, routes each row to one of two ingest pathways based on
state-type classification produced by Stage 2:

  - BULK pathway (existing): supersedes_prior_state=false OR confidence=low
    OR routing fields missing. Fast, no temporal invalidation.

  - SINGLE-EPISODE pathway (new): supersedes_prior_state=true AND
    confidence in {medium, high}. Per-chunk POST to /episodes with shared
    saga tag, full edge invalidation, per-chunk timeout/retry independence.

Both pathways pass EXTRACTION_INSTRUCTIONS_V1 to the sidecar via
custom_extraction_instructions, which graphiti-core inserts into entity
and edge extraction prompts (NOT dedup prompts — that's intentional under
the encoder-stays-naive commitment).

Architectural posture: the encoder is content-naïve. It does not draw on
prior knowledge of the user, the substrate, or the cycle's accumulated
work. Schema and personality live in the cycle's consolidated substrate,
where the dream phase shapes them. The encoder produces source-grounded
ground truth for the cycle to work from. See EXTRACTION_INSTRUCTIONS_V1
below for the extraction guidance text.

Routing rationale: the single-episode pathway is the correct API per
graphiti-core's docs for content that supersedes prior facts (it does
edge invalidation that bulk skips). It costs more per chunk because of
the resolve_edge LLM call; the routing rule keeps that cost bounded to
content that actually needs it.

Chunking rationale (preserved from prior versions): 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.

Per-chunk heartbeat: single-episode pathway updates stage_3_queue.started_at
after each successful chunk POST so a long-running document doesn't cross
the 10-minute stale threshold mid-process and get re-dequeued by another
worker (or the same worker on next loop iteration). started_at thus means
"last activity timestamp" rather than "began at" — semantics that match
the dequeue query's intent (catch dead workers, not slow ones).

Saga-size limit (MAX_CHUNKS_PER_SAGA): 2026-05-01 incident showed bulk
sagas of 17 and 19 chunks deadlock the sidecar's Python-side coordination.
Documents producing more than MAX_CHUNKS_PER_SAGA chunks on the bulk
pathway are split into multiple bulk commits, each tagged with the same
saga value so Graphiti still links them. The single-episode pathway
doesn't need this split since each chunk is its own POST.

Wedge detection: 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.4"

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

# Routing rule: single-episode pathway requires both signals positive.
# Anything else (false, NULL, low confidence) routes to bulk — the
# safer-cheaper default. Mistral parse drift can't accidentally trigger
# the expensive pathway.
HIGH_TRUST_CONFIDENCE = ("medium", "high")

# Encoder extraction guidance v1.0 — see module docstring for posture rationale.
# Passed to graphiti-core via custom_extraction_instructions on both ingest
# pathways. Inserted into entity-extraction and edge-extraction prompts only;
# does NOT enter dedup prompts. Encoder-stays-naïve commitment is structural,
# not versioned: this text gets refined over time but the encoder does not
# acquire substrate context as the cycle matures.
EXTRACTION_INSTRUCTIONS_V1 = """\
EXTRACTION GUIDANCE — BirdAI cascade

The encoder's job is faithful capture from this chunk's text. It does
not draw on prior knowledge of the user, the substrate, or the cycle's
accumulated work. Schema, personality, and inferred context live in
the cycle's consolidated substrate, where the dream phase shapes them
through prediction-error replay and speculation. The encoder stays
content-naïve so the cycle has source-grounded ground truth to work
from.

The orientation produced by an upstream pass describes content shape,
not content interpretation. Use it as forward-facing guidance for what
to attend to in this document. Do not let it bound or limit what you
extract.

PREDICATE NAMING

Produce semantic predicates that describe the actual relationship the
text states. Use verbs or verb phrases — "wrote", "advised", "founded",
"works at", "led to", "contradicts", "is autobiographical to" — not
generic placeholders. Reserve generic forms (for example, "relates to"
or "mentions") for cases where the text genuinely does not specify a
more particular relationship. The verb is the load-bearing part of
the fact; preserving it is what makes the relationship queryable later.

EXTRACTION POSTURE

Extract from this chunk's text as if each entity is encountered fresh.
Do not try to reconcile entities you find here with entities that
might already exist elsewhere in the graph. Redundant entity instances
are acceptable. Cross-document entity resolution is downstream cycle
work, not extraction work.

When the same entity appears multiple times within this chunk with
slightly different spellings — a common artifact of voice transcription —
prefer the more frequent or more canonical-looking form. Do not invent
canonical forms; choose among the variants the text actually contains.

EXTRACT FROM THE SOURCE

Extract relationships the text states or strongly implies through
direct linguistic markers ("X led to Y", "X works for Y", "X met Y at
Z"). Do not extend extraction to relationships the text neither states
nor directly implies. Inferred relationships are produced by the
cycle's dream phase as speculative edges with explicit low-confidence
tagging, where they can be evaluated and either ratified or pruned by
subsequent cycle work. Encoding-time inference, mixed in with source-
grounded extraction, would lose the speculation/source distinction the
cycle's consolidation work relies on.

DO NOT PRE-EMPT CYCLE WORK

Do not omit relationships because they seem redundant with prior
extractions or with the existing graph. Cross-document entity
resolution and edge consolidation are downstream cycle operations;
redundant extraction at this stage is intentional. Extracting the
same fact from multiple sources gives the cycle's consolidation work
the recurrence signal it relies on.

EXTRACTION DEPTH

Use the orientation's frame_relationships and extraction_orientation
fields to inform what to attend to. If the orientation describes
cross-domain relational content, look for relationships that bridge
those domains explicitly, with named predicates for the bridging.
If the orientation describes single-domain technical content, look
for the structural relationships internal to that domain.

Extract every entity and every relationship the text states. Do not
summarize, do not filter, do not omit content because it seems
incidental. The orientation tells you what to look for; the source
text tells you what is there.
"""


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 heartbeat_row(row_id):
    """Refresh stage_3_queue.started_at to NOW() so a long-running single-episode
    ingest doesn't cross the 10-minute stale threshold mid-process. Called
    after each successful chunk POST. Best-effort: failures are logged but
    don't fail the chunk — the worst case is a stale-threshold re-dequeue,
    which graphiti's dedup will handle as a no-op."""
    try:
        pg = get_pg()
        cur = pg.cursor()
        cur.execute("UPDATE stage_3_queue SET started_at = NOW() WHERE id = %s", (row_id,))
        pg.commit()
        pg.close()
    except Exception as e:
        log.warning(f"  Heartbeat update failed (continuing): {e}")


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 post_episode(payload, episode_label=""):
    """Single POST to /episodes (singular) with consistent error handling.
    Used by the single-episode pathway, one call per chunk."""
    resp = requests.post(
        f"{GRAPHITI_URL}/episodes",
        json=payload,
        timeout=INGEST_TIMEOUT
    )
    if not resp.ok:
        prefix = f"{episode_label} " if episode_label else ""
        raise RuntimeError(f"{prefix}Sidecar {resp.status_code}: {resp.text[:500]}")
    return resp.json()


def ingest_bulk(source, full_text, orientation):
    """
    Bulk-pathway ingest: documents that don't supersede prior state.
    Skips edge invalidation. Cheap. Three sub-paths by document size:

    - Short documents (<SINGLE_EPISODE_THRESHOLD): single episode, no saga
      [note: 'single episode' here means one bulk call with one item, NOT
       the single-episode-pathway; naming overlap is unfortunate but local]
    - 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

    All three sub-paths pass EXTRACTION_INSTRUCTIONS_V1 to the sidecar.
    """
    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"  [bulk] Single episode ({char_length} chars)")
        return post_bulk({
            "episodes": episodes,
            "group_id": "aaron",
            "custom_extraction_instructions": EXTRACTION_INSTRUCTIONS_V1,
        })

    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"  [bulk] Chunked into {total_chunks} episodes ({char_length} chars)")
        return post_bulk({
            "episodes": episodes,
            "group_id": "aaron",
            "saga": source,
            "custom_extraction_instructions": EXTRACTION_INSTRUCTIONS_V1,
        })

    # 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"  [bulk] 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,
                "custom_extraction_instructions": EXTRACTION_INSTRUCTIONS_V1,
            },
            batch_label=batch_label,
        )
        log.info(f"    {batch_label} committed")
    return last_result


def ingest_single_episode(row_id, source, full_text, orientation):
    """
    Single-episode pathway: documents that supersede prior state with
    medium-or-high confidence. Each chunk is its own POST to /episodes
    with shared saga tag. Each call independent: own timeout, own retry
    envelope, own failure semantics.

    Each chunk POST passes EXTRACTION_INSTRUCTIONS_V1 to the sidecar.

    Partial-success behavior: if chunk N of total fails, chunks 1..N-1
    stay committed (graphiti has already accepted them) and the function
    raises with detail about which chunk failed and how many succeeded.
    The caller marks the row failed_at with that detail; the operator
    decides whether to re-enqueue. Re-ingestion will re-POST chunks 1..N-1
    against the graph; graphiti's dedup will handle them as no-ops.

    Heartbeats stage_3_queue.started_at after each successful chunk so the
    row doesn't cross the 10-minute stale threshold while actively progressing.
    """
    char_length = len(full_text)

    # Short documents: one POST, no chunking, no saga
    if char_length < SINGLE_EPISODE_THRESHOLD:
        payload = {
            "name": source,
            "content": full_text,
            "source_description": orientation,
            "group_id": "aaron",
            "timestamp": datetime.now().isoformat(),
            "custom_extraction_instructions": EXTRACTION_INSTRUCTIONS_V1,
        }
        log.info(f"  [single-ep] Single episode, no chunking ({char_length} chars)")
        return post_episode(payload, episode_label="single-ep")

    chunks = chunk_text(full_text)
    total_chunks = len(chunks)
    log.info(
        f"  [single-ep] Chunked into {total_chunks} episodes ({char_length} chars); "
        f"per-chunk POSTs with shared saga"
    )

    succeeded = 0
    for i, chunk in enumerate(chunks):
        chunk_num = i + 1
        payload = {
            "name": f"{source} [{chunk_num}/{total_chunks}]",
            "content": chunk,
            "source_description": orientation,
            "group_id": "aaron",
            "saga": source,
            "timestamp": datetime.now().isoformat(),
            "custom_extraction_instructions": EXTRACTION_INSTRUCTIONS_V1,
        }
        try:
            post_episode(payload, episode_label=f"chunk {chunk_num}/{total_chunks}")
            succeeded += 1
            log.info(f"    chunk {chunk_num}/{total_chunks} committed")
            heartbeat_row(row_id)
        except Exception as e:
            # Annotate the exception with partial-success detail so the
            # caller can write a clean failure_reason. Re-raise to abort
            # the document; previously-committed chunks stay in the graph.
            raise RuntimeError(
                f"single_episode_partial: chunk {chunk_num}/{total_chunks} failed "
                f"(succeeded: {succeeded}); error: {str(e)[:300]}"
            ) from e

    log.info(f"  [single-ep] All {total_chunks} chunks committed")
    return {"ok": True, "chunks_committed": total_chunks}


def should_route_single_episode(supersedes_prior_state, state_type_confidence):
    """Routing decision for Phase A.

    Single-episode pathway requires BOTH:
      - supersedes_prior_state is true (Mistral judged it temporally superseding)
      - confidence is medium or high (Mistral was confident enough to trust)

    Anything else routes to bulk: false supersedes, NULL fields (legacy rows
    pre-dating Stage 2 v2.2), low confidence even on supersedes=true. This
    is the safer-cheaper default — bulk skips temporal invalidation, which
    is the right behavior when we're not confident the content needs it.
    """
    if not supersedes_prior_state:
        return False
    if state_type_confidence not in HIGH_TRUST_CONFIDENCE:
        return False
    return True


def process_one(row):
    (row_id, source, full_text, orientation,
     state_type, state_type_confidence, supersedes_prior_state,
     state_type_rationale) = row

    # Route decision
    use_single_episode = should_route_single_episode(
        supersedes_prior_state, state_type_confidence
    )
    pathway = "single-episode" if use_single_episode else "bulk"

    log.info(
        f"Ingesting to Graphiti: {source} "
        f"[pathway={pathway}, state_type={state_type}, "
        f"conf={state_type_confidence}, supersedes={supersedes_prior_state}]"
    )
    if state_type_rationale:
        log.info(f"  rationale: {state_type_rationale[:200]}")

    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()
    pg.close()

    try:
        if use_single_episode:
            result = ingest_single_episode(row_id, source, full_text, orientation)
        else:
            result = ingest_bulk(source, full_text, orientation)

        pg = get_pg()
        cur = pg.cursor()
        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}")
        pg = get_pg()
        cur = pg.cursor()
        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,
                       state_type, state_type_confidence, supersedes_prior_state,
                       state_type_rationale
                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()
