"""
common/pipeline_pool.py

Intermediate base class: PoolWorker + pipeline stage logic.

Bridges pool_worker.PoolWorker (HTTP dispatch, heartbeat, state machine)
with the existing pipeline stage pattern (Garage S3 payload, checkpoint,
read_stage / write_stage / write_result).

Subclasses declare:
    output_stage: str     — name of the stage this worker produces
    process(doc_id, context_uuid, metadata, payload) -> Any

And optionally override:
    on_cleanup()          — release task-local resources before FREE

Everything else — checkpoint detection, Garage read/write, result atom
writing to Redis, error handling — is handled here.
"""
from __future__ import annotations

import os
import sys
import syslog

sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))

from pool_worker import PoolWorker
from common.pipeline import (
    stage_exists, read_stage, write_stage, storage_key
)
from common.task import read_task, parse_task, write_result


class PipelinePoolWorker(PoolWorker):
    """
    Abstract base for pipeline stage pool workers.

    Subclasses MUST define:
        output_stage: str
        process(doc_id, context_uuid, metadata, payload) -> Any

    Subclasses MAY override:
        on_cleanup()       — release task-local resources (tensors, buffers)
                             model/DB connections must NOT be released here
    """

    output_stage: str  # must be set by subclass

    def Execute(self, task_uuid: str) -> bool:
        """
        Standard pipeline Execute contract:
          1. Fetch task atom from Redis
          2. Checkpoint: if output already exists, skip and write result
          3. Load input payload from Garage S3
          4. Call process() — pure computation, no side effects
          5. Write output to Garage S3
          6. Write result atom to Redis for n8n to read
        """
        task = read_task(self._redis, task_uuid)
        if not task:
            syslog.syslog(syslog.LOG_ERR,
                          f"[{self.__class__.__name__}] "
                          f"task not found: {task_uuid}")
            return False

        doc_id, stage, in_key, context_uuid, metadata = parse_task(task)
        cls_name = self.__class__.__name__

        # Checkpoint — if output already in Garage, skip computation
        if stage_exists(context_uuid, doc_id, self.output_stage):
            out_key = storage_key(context_uuid, doc_id, self.output_stage)
            syslog.syslog(syslog.LOG_INFO,
                          f"[{cls_name}] checkpoint hit "
                          f"doc={doc_id[:8]} out={self.output_stage}")
            write_result(self._redis, task_uuid, self.output_stage, out_key)
            return True

        # Load input payload from Garage
        payload = read_stage(in_key)

        # Run stage computation
        syslog.syslog(syslog.LOG_INFO,
                      f"[{cls_name}] processing doc={doc_id[:8]}")
        output = self.process(doc_id, context_uuid, metadata, payload)

        # Write output to Garage
        out_key = write_stage(context_uuid, doc_id, self.output_stage, output)

        # Write result atom for n8n
        write_result(self._redis, task_uuid, self.output_stage, out_key)

        syslog.syslog(syslog.LOG_INFO,
                      f"[{cls_name}] done doc={doc_id[:8]} → {out_key}")
        return True

    def process(self, doc_id: str, context_uuid: str,
                metadata: dict, payload) -> object:
        """Override in subclass. Returns output payload dict."""
        raise NotImplementedError

    def on_cleanup(self):
        """Override to release task-local resources. Default: no-op."""
