"""
common/task.py

Task atom contract between n8n (TaskScheduler) and pool workers.

Redis layout:
    task:{uuid}          → JSON  (input task atom, written by n8n before calling /execute)
    task:{uuid}:result   → JSON  (output written by worker after Execute() completes)

Task atom fields (task:{uuid}):
    doc_id        str   UUID hex identifying the document across all pipeline stages
    stage         str   input stage name (e.g. "raw", "tokenized")
    storage_key   str   Garage S3 key for the input payload
    context_uuid  str   profile UUID or zero UUID for global
    metadata      str   JSON-encoded metadata dict {title, author, year, ...}

Result atom fields (task:{uuid}:result):
    out_stage     str   output stage name (e.g. "tokenized", "hidden")
    out_key       str   Garage S3 key for the output payload
    status        str   "ok" | "error"
    error         str | null

TTL: 86400 seconds (24h) for both task and result atoms. Matches
     pipeline-dispatcher's resume window (DISPATCHER_TASK_TTL=86400) so
     a dispatcher crash + restart can short-circuit a stage whose worker
     already produced a result. Task atoms can be deleted by the
     dispatcher after observing the result.

Pool IDs:
    1  tokenizer
    2  forward_pass
    3  chunker
    4  extractor
    5  tagger
    6  storage
    7  compressor
"""
from __future__ import annotations

import json
import os
from typing import Any

import redis as _redis_mod

REDIS_URL    = os.getenv("REDIS_URL", "redis://10.121.1.60:6379/0")
RESULT_TTL   = int(os.getenv("TASK_RESULT_TTL", "86400"))


class CheckpointDrainRequested(Exception):
    """Raised inside a worker's compute loop (e.g. windowed_forward_pass) when a
    checkpoint-drain was requested (upgrade protocol, operator 2026-06-19): the work
    has been durably checkpointed up to `.window`, so the loop stops HERE and the
    caller re-queues the task (pinned to its own endpoint) for the relaunched worker
    to resume. `.window` = the next window to compute on resume."""
    def __init__(self, window: int):
        super().__init__(f"checkpoint-drain at window {window}")
        self.window = window

# Pool ID registry — the PIPELINE-STAGE pools only. This is NOT the complete list: the qwen worker's
# pools live in qwen_service (POOL_FORWARD_PASS=2 mirrored here; POOL_QWEN_EMBEDDER/POOL_EMBEDDING=8 are
# ONLY there, NOT here). The AUTHORITATIVE cross-file map is wiki memory/reference_pool_id_registry —
# grep BOTH common/ and qwen_service/ before assigning a new id (g#194: OCR took 8, colliding w/ embed).
POOL_TOKENIZER    = 1
POOL_FORWARD_PASS = 2
POOL_CHUNKER      = 3
POOL_EXTRACTOR    = 4
POOL_TAGGER       = 5
POOL_STORAGE      = 6
POOL_COMPRESSOR   = 7
POOL_OCR          = 9   # scanned PDF/image -> text (pre-tokenize for scanned corpora; #185). WAS 8, but
                        # that COLLIDES with qwen_service POOL_QWEN_EMBEDDER=8 (embedder, :8801) — OCR
                        # tasks landed on qwen embed workers and errored (g#194). Moved to 9 (next free).
POOL_DISSECT      = 10  # OCR workflow stage 1 (fan-out): container atom (PDF) -> one image atom per page.
                        # rm#34 / g#195; workflows/ocr.json; wiki memory/reference_pool_id_registry. Next free: 12.
POOL_ASSEMBLE     = 11  # OCR workflow stage 3 (fan-in barrier): stitch per-page texts -> transcript. rm#34/g#195.


def get_redis() -> _redis_mod.Redis:
    return _redis_mod.from_url(
        REDIS_URL,
        decode_responses=True,
        socket_connect_timeout=5,
        socket_timeout=60,
        retry_on_timeout=True,
    )


def read_task(r: _redis_mod.Redis, task_uuid: str) -> dict | None:
    """
    Fetch task atom from Redis.
    Returns parsed dict or None if key missing.
    """
    raw = r.get(f"task:{task_uuid}")
    if not raw:
        return None
    return json.loads(raw)


def parse_task(task: dict) -> tuple[str, str, str, str, dict]:
    """
    Returns (doc_id, stage, storage_key, context_uuid, metadata).
    metadata is already a dict (parsed from JSON string).
    """
    return (
        task["doc_id"],
        task["stage"],
        task["storage_key"],
        task["context_uuid"],
        json.loads(task.get("metadata", "{}")),
    )


def write_result(r: _redis_mod.Redis, task_uuid: str,
                 out_stage: str, out_key: str,
                 status: str = "ok", error: str | None = None,
                 token_count: int | None = None) -> None:
    """Write result atom to Redis with TTL. token_count (when the stage knows it,
    e.g. tokenize) rides the done signal so the advancer can size-route the next
    stage without a Garage read."""
    r.set(
        f"task:{task_uuid}:result",
        json.dumps({
            "out_stage":   out_stage,
            "out_key":     out_key,
            "status":      status,
            "error":       error,
            "token_count": token_count,
        }),
        ex=RESULT_TTL,
    )


def write_task(r: _redis_mod.Redis, task_uuid: str,
               doc_id: str, stage: str, storage_key: str,
               context_uuid: str, metadata: dict,
               ttl: int = 86400) -> None:
    """
    Write a task atom to Redis.
    Called by pipeline-dispatcher via its own poolservice.write_task_atom
    (which passes DISPATCHER_TASK_TTL explicitly — this default is the
    fallback for any other caller).
    """
    r.set(
        f"task:{task_uuid}",
        json.dumps({
            "doc_id":       doc_id,
            "stage":        stage,
            "storage_key":  storage_key,
            "context_uuid": context_uuid,
            "metadata":     json.dumps(metadata),
        }),
        ex=ttl,
    )


def delete_task(r: _redis_mod.Redis, task_uuid: str) -> None:
    """Clean up task + result atoms after n8n has read the result."""
    r.delete(f"task:{task_uuid}", f"task:{task_uuid}:result")
