#!/usr/bin/env python3
"""
Stage 2 Worker — Taxonomy-Free Mistral Orientation + State-Type Classification

Polls stage_2_queue, runs Mistral pass that produces:
  (a) orientation context (active frames, frame relationships, extraction focus)
  (b) state-type classification for Stage 3 routing (current/reference/historical,
      supersedes_prior_state boolean, confidence, rationale)

Enqueues Stage 3 with both concerns as explicit columns.

Routing:
- char_length < 2000 → skip Stage 3, mark complete (sparse content, cascade no benefit)
- char_length >= 2000 → enqueue Stage 3 with orientation + routing metadata

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

import os, json, time, subprocess, logging, 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 [stage2] %(levelname)s %(message)s",
    handlers=[
        logging.FileHandler("/var/log/aaronai/stage2.log", mode="a"),
    ]
)
log = logging.getLogger("stage2")

PG_DSN = os.getenv("PG_DSN")
OLLAMA_URL = "http://localhost:11434"
HEARTBEAT_FILE = Path("/var/log/aaronai/stage2-heartbeat")
CHAR_LENGTH_THRESHOLD = 2000
REQUEST_TIMEOUT = 300
RETRY_ATTEMPTS = 2
POLL_INTERVAL = 5
WORKER_VERSION = "2.2"

# Valid values for state-type fields. Mistral output validated against these;
# anything outside falls through to safe-cheap defaults (bulk routing).
VALID_STATE_TYPES = ("current", "reference", "historical")
VALID_CONFIDENCE  = ("low", "medium", "high")

# Safe-cheap defaults applied when Mistral output is missing or malformed.
# All route to bulk pathway (no temporal invalidation cost) per Phase A
# routing rule: route to single-episode only on supersedes_prior_state=true
# AND confidence in {medium, high}.
DEFAULT_STATE_TYPE       = "reference"
DEFAULT_CONFIDENCE       = "low"
DEFAULT_SUPERSEDES       = False
DEFAULT_RATIONALE        = "mistral output missing or malformed; default applied"

TAXFREE_PROMPT = (
    "You are a metadata extraction system. Given a document, produce a JSON object "
    "describing two distinct concerns about the document. Output JSON only — no prose, "
    "no explanation, no markdown.\n\n"

    "CONCERN 1 — ORIENTATION CONTEXT (for downstream knowledge-graph extraction):\n"
    "Describe the content shape. Do not summarize content. Do not extract entities. "
    "Do not assign a single category label.\n"
    "  - active_frames: which domains or frames are active in this content (there may "
    "be several simultaneously)\n"
    "  - frame_relationships: how those frames relate in this specific document, "
    "one sentence\n"
    "  - extraction_orientation: what kind of relational content a knowledge-graph "
    "extractor should look for, one sentence\n"
    "  - one_sentence_summary: a single-sentence content summary\n\n"

    "CONCERN 2 — STATE-TYPE CLASSIFICATION (for ingest routing):\n"
    "Classify the document's relationship to time and prior facts. This is independent "
    "of orientation: a document can be in a 'reference frame' (orientation) while "
    "describing 'current state' (state-type), or vice versa. Judge the document's "
    "ROLE, not its topic.\n"
    "  - state_type: one of\n"
    "      'current'    — describes the author's present state, recent decisions, or "
    "ongoing situations as of the document's date\n"
    "      'reference'  — timeless or slow-changing material: external books, "
    "documentation, technical reference, conceptual writing\n"
    "      'historical' — describes past events, prior states, or archived material "
    "the author is recording but not living in\n"
    "  - state_type_confidence: 'low' | 'medium' | 'high' — how confident you are in "
    "the classification. Use 'low' when genuinely uncertain.\n"
    "  - supersedes_prior_state: true if this document describes facts that should "
    "REPLACE previously-recorded facts about the same subjects (e.g. a journal entry "
    "saying 'I no longer work at X', a status update, a corrected belief). false "
    "otherwise. Default to false when uncertain.\n"
    "  - state_type_rationale: one sentence explaining the classification\n\n"

    "Output schema (flat, all eight fields at the top level):\n"
    '{"active_frames": ["<frame 1>", "<frame 2>"], '
    '"frame_relationships": "<one sentence>", '
    '"extraction_orientation": "<one sentence>", '
    '"one_sentence_summary": "<one sentence>", '
    '"state_type": "current|reference|historical", '
    '"state_type_confidence": "low|medium|high", '
    '"supersedes_prior_state": true|false, '
    '"state_type_rationale": "<one sentence>"}\n\n'

    "Document:\n"
)


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():
    log.warning("Mistral wedge detected — restarting Ollama")
    result = subprocess.run(["/usr/bin/sudo", "/bin/systemctl", "restart", "ollama"], capture_output=True, text=True)
    if result.returncode != 0:
        log.error(f"Ollama restart failed (rc={result.returncode}): stdout={result.stdout!r} stderr={result.stderr!r}")
        return False
    time.sleep(30)
    for _ in range(3):
        try:
            r = requests.get(f"{OLLAMA_URL}/api/tags", timeout=10)
            if r.status_code == 200:
                log.info("Ollama recovered")
                return True
        except Exception:
            time.sleep(5)
    log.error("Ollama recovery failed")
    return False


def run_mistral(doc_text):
    payload = {
        "model": "mistral:latest",
        "prompt": TAXFREE_PROMPT + doc_text[:12000],
        "stream": False,
        "format": "json",
    }
    resp = requests.post(f"{OLLAMA_URL}/api/generate", json=payload, timeout=REQUEST_TIMEOUT)
    resp.raise_for_status()
    raw = resp.json().get("response", "{}")
    try:
        return json.loads(raw)
    except Exception:
        return {"error": "parse_failed", "raw": raw[:200]}


def normalize_state_fields(meta):
    """Validate and normalize the four state-type fields from Mistral output.
    Anything missing or malformed falls through to safe-cheap defaults that
    route to bulk pathway (no temporal invalidation work)."""

    raw_state_type = meta.get("state_type")
    if isinstance(raw_state_type, str) and raw_state_type.lower() in VALID_STATE_TYPES:
        state_type = raw_state_type.lower()
    else:
        state_type = DEFAULT_STATE_TYPE

    raw_conf = meta.get("state_type_confidence")
    if isinstance(raw_conf, str) and raw_conf.lower() in VALID_CONFIDENCE:
        confidence = raw_conf.lower()
    else:
        confidence = DEFAULT_CONFIDENCE

    raw_supersedes = meta.get("supersedes_prior_state")
    if isinstance(raw_supersedes, bool):
        supersedes = raw_supersedes
    else:
        supersedes = DEFAULT_SUPERSEDES

    raw_rationale = meta.get("state_type_rationale")
    if isinstance(raw_rationale, str) and raw_rationale.strip():
        rationale = raw_rationale.strip()[:1000]
    else:
        rationale = DEFAULT_RATIONALE

    return state_type, confidence, supersedes, rationale


def build_orientation(meta):
    frames = ", ".join(meta.get("active_frames", []))
    rel = meta.get("frame_relationships", "")
    orient = meta.get("extraction_orientation", "")
    summary = meta.get("one_sentence_summary", "")
    return f"Active frames: {frames}. Frame relationships: {rel} Extraction focus: {orient} Summary: {summary}"


def enqueue_stage3(pg, source, full_text, orientation, metadata,
                   state_type, state_type_confidence, supersedes_prior_state,
                   state_type_rationale):
    """Write Stage 3 queue row with orientation + explicit routing columns.

    Routing columns (state_type, state_type_confidence, supersedes_prior_state,
    state_type_rationale) are first-class queue properties for Phase A.
    Stage 3 reads them on every dequeue to choose bulk vs single-episode pathway.
    The full Mistral metadata blob is also retained in stage2_metadata JSON for
    debugging and future cycle work."""
    cur = pg.cursor()
    cur.execute("""
        INSERT INTO stage_3_queue (
            source, full_text, orientation, stage2_metadata,
            state_type, state_type_confidence, supersedes_prior_state,
            state_type_rationale
        )
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
        ON CONFLICT (source) DO UPDATE SET
            full_text = EXCLUDED.full_text,
            orientation = EXCLUDED.orientation,
            stage2_metadata = EXCLUDED.stage2_metadata,
            state_type = EXCLUDED.state_type,
            state_type_confidence = EXCLUDED.state_type_confidence,
            supersedes_prior_state = EXCLUDED.supersedes_prior_state,
            state_type_rationale = EXCLUDED.state_type_rationale,
            enqueued_at = NOW(),
            completed_at = NULL,
            failed_at = NULL,
            attempts = 0
    """, (source, full_text, orientation, json.dumps(metadata),
          state_type, state_type_confidence, supersedes_prior_state,
          state_type_rationale))
    pg.commit()


def process_one(row):
    row_id, source, full_text, char_length = row
    log.info(f"Processing: {source} ({char_length} chars)")

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

    # Routing: sparse content skips Stage 3
    if char_length < CHAR_LENGTH_THRESHOLD:
        log.info(f"  Skipping Stage 3 (char_length={char_length} < {CHAR_LENGTH_THRESHOLD})")
        cur.execute("UPDATE stage_2_queue SET completed_at = NOW() WHERE id = %s", (row_id,))
        pg.commit()
        pg.close()
        return True

    # Run Mistral
    log.info(f"  Running Mistral taxonomy-free + state-type pass...")
    try:
        meta = run_mistral(full_text)
    except requests.exceptions.Timeout:
        log.warning(f"  Mistral timeout on {source}")
        cur.execute(
            "UPDATE stage_2_queue SET failed_at = NOW(), failure_reason = %s WHERE id = %s",
            (f"mistral_timeout_after_{REQUEST_TIMEOUT}s", row_id)
        )
        pg.commit()
        pg.close()
        return False
    except Exception as e:
        log.error(f"  Mistral error on {source}: {e}")
        cur.execute("UPDATE stage_2_queue SET failed_at = NOW(), failure_reason = %s WHERE id = %s",
                    (str(e)[:500], row_id))
        pg.commit()
        pg.close()
        return False

    if meta.get("error") == "parse_failed":
        log.warning(f"  Mistral parse failure on {source}: {meta.get('raw', '')[:100]}")
        cur.execute(
            "UPDATE stage_2_queue SET failed_at = NOW(), failure_reason = %s WHERE id = %s",
            ("mistral_parse_failure", row_id)
        )
        pg.commit()
        pg.close()
        return False

    frames = meta.get("active_frames", [])
    log.info(f"  Frames: {frames}")

    # Normalize state-type fields with safe-cheap defaults on malformed output.
    # Note: Mistral may return valid orientation but malformed state-type;
    # we accept the orientation and default the routing rather than fail
    # the whole row, since defaults route to bulk (cheap, safe).
    state_type, confidence, supersedes, rationale = normalize_state_fields(meta)
    log.info(
        f"  State-type: {state_type} (conf={confidence}, "
        f"supersedes={supersedes})"
    )

    orientation = build_orientation(meta)
    meta["_model"] = "mistral:latest"
    meta["_worker_version"] = WORKER_VERSION
    meta["_generated_at"] = datetime.now().isoformat()
    meta["char_length"] = char_length

    # Enqueue Stage 3 with explicit routing columns
    enqueue_stage3(pg, source, full_text, orientation, meta,
                   state_type, confidence, supersedes, rationale)
    cur.execute("UPDATE stage_2_queue SET completed_at = NOW() WHERE id = %s", (row_id,))
    pg.commit()
    pg.close()
    log.info(f"  Enqueued Stage 3: {source}")
    return True


def run():
    log.info(f"Stage 2 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, char_length
                FROM stage_2_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 Mistral 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(1)

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


if __name__ == "__main__":
    run()
