#!/usr/bin/env python3
"""
Stage 2 Worker — Taxonomy-Free Mistral Orientation
Polls stage_2_queue, runs Mistral taxonomy-free pass, enqueues Stage 3.
Runs as systemd service: aaronai-stage2.service

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

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.1"

TAXFREE_PROMPT = (
    "You are a metadata extraction system. Given a document, describe its content "
    "shape for use as orientation context in a knowledge graph extraction pass.\n\n"
    "Do not summarize content. Do not extract entities. Do not assign a single category label.\n\n"
    "Instead, describe:\n"
    "- What domains or frames are active in this content (there may be several simultaneously)\n"
    "- How those frames relate to each other in this specific document\n"
    "- What kind of relational content a knowledge graph extractor should look for\n\n"
    "Output JSON only. No prose, no explanation, no markdown.\n\n"
    "Schema:\n"
    '{"active_frames": ["<frame 1>", "<frame 2>"], '
    '"frame_relationships": "<one sentence>", '
    '"extraction_orientation": "<one sentence>", '
    '"one_sentence_summary": "<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 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):
    cur = pg.cursor()
    cur.execute("""
        INSERT INTO stage_3_queue (source, full_text, orientation, stage2_metadata)
        VALUES (%s, %s, %s, %s)
        ON CONFLICT (source) DO UPDATE SET
            full_text = EXCLUDED.full_text,
            orientation = EXCLUDED.orientation,
            stage2_metadata = EXCLUDED.stage2_metadata,
            enqueued_at = NOW(),
            completed_at = NULL,
            failed_at = NULL,
            attempts = 0
    """, (source, full_text, orientation, json.dumps(metadata)))
    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 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}")

    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
    enqueue_stage3(pg, source, full_text, orientation, meta)
    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()
