#!/usr/bin/env python3
"""
task_worker.py
==============
PocketBase task worker — pulls ready tasks and dispatches to the
right sidecar. Two lanes:

  --lane deterministic
    - bizportal_search  → http://127.0.0.1:7001/search
    - bizportal_reserve → http://127.0.0.1:7001/reserve
    - pdf_generate      → shell out to pdf-creator
    - email_send        → SMTP via the user's email config
    (no LLM calls — runs fast, deterministic)

  --lane custom
    - custom            → write a JSON file to /tmp/agent_inbox/ for
                          the LLM agent to pick up. The agent is
                          expected to PATCH the task back via the
                          /api/tasks/[id]/complete endpoint.

Both lanes share the same atomic claim protocol. Only one worker
ever holds a given task.

Usage:
  python3 scripts/task_worker.py --lane deterministic
  python3 scripts/task_worker.py --lane custom
"""

import argparse
import json
import logging
import os
import smtplib
import subprocess
import sys
import time
import urllib.request
import urllib.error
from email.mime.text import MIMEText
from pathlib import Path
from typing import Any, Dict, Optional

# ── Config ──────────────────────────────────────────────────────

ROOT = Path(__file__).resolve().parent.parent
ENV_FILE = ROOT / ".env.local"

DEFAULT_PB_URL = "https://pb.agent.apphub.co.za"
DEFAULT_TASKS_API = "http://127.0.0.1:6001"  # the Next.js side
DEFAULT_BIZPORTAL_SERVICE = "http://127.0.0.1:7001"
DEFAULT_AGENT_INBOX = "/tmp/agent_inbox"
DEFAULT_PDF_CREATOR = "/Volumes/KIOXIA/.agents/skills/pdf-creator"
POLL_INTERVAL_S = 2.0
HTTP_TIMEOUT_S = 120

logging.basicConfig(
    level=os.getenv("LOG_LEVEL", "INFO"),
    format="%(asctime)s [%(levelname)s] %(name)s — %(message)s",
)
log = logging.getLogger(f"task-worker[{os.getpid()}]")


def load_env() -> Dict[str, str]:
    env = dict(os.environ)
    if ENV_FILE.exists():
        for line in ENV_FILE.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, v = line.split("=", 1)
            env[k.strip()] = v.strip().strip('"').strip("'")
    return env


# ── HTTP helpers ────────────────────────────────────────────────

def http_json(url: str, method: str = "GET", headers: Optional[Dict] = None, body: Any = None):
    data = None
    if body is not None:
        data = json.dumps(body).encode("utf-8")
    h = dict(headers or {})
    if data is not None and "Content-Type" not in h:
        h["Content-Type"] = "application/json"
    req = urllib.request.Request(url, data=data, headers=h, method=method)
    try:
        with urllib.request.urlopen(req, timeout=HTTP_TIMEOUT_S) as resp:
            payload = resp.read().decode("utf-8")
            return resp.status, json.loads(payload) if payload else None
    except urllib.error.HTTPError as e:
        body = e.read().decode("utf-8", errors="ignore")
        try:
            return e.code, json.loads(body)
        except Exception:
            return e.code, {"raw": body[:500]}
    except Exception as e:
        return 0, {"error": str(e)}


# ── Task lifecycle ─────────────────────────────────────────────

def list_ready(env: Dict[str, str], lane: str, limit: int = 1) -> list:
    """Fetch ready tasks. Filter by what the lane can handle."""
    pb_url = env.get("PB_URL", DEFAULT_PB_URL)
    # We go through the Next.js API so the auth flow is consistent.
    api = env.get("TASKS_API", DEFAULT_TASKS_API)
    qs = []
    qs.append("status=ready")
    if lane == "deterministic":
        # Deterministic lane takes everything except custom
        pass
    elif lane == "custom":
        # LLM lane takes only custom
        qs.append("task_type=custom")
    qs.append(f"limit={limit}")
    url = f"{api}/api/tasks?{'&'.join(qs)}"
    code, data = http_json(url, headers={"Authorization": env.get("PB_TOKEN", "")})
    if code != 200:
        log.warning(f"list ready failed: {code} {data}")
        return []
    if isinstance(data, dict) and "items" in data:
        return data["items"]
    return []


def claim_task(env: Dict[str, str], task_id: str) -> Optional[Dict]:
    api = env.get("TASKS_API", DEFAULT_TASKS_API)
    url = f"{api}/api/tasks/{task_id}/claim"
    code, data = http_json(
        url,
        method="POST",
        headers={"Authorization": env.get("PB_TOKEN", "")},
        body={"claimed_by": f"task_worker[{os.getpid()}]:{lane_label(env)}"},
    )
    if code == 200 and isinstance(data, dict) and data.get("ok"):
        return data.get("task")
    return None


def complete_task(env: Dict[str, str], task_id: str, result: Any, log_lines: list, duration_ms: int) -> bool:
    api = env.get("TASKS_API", DEFAULT_TASKS_API)
    url = f"{api}/api/tasks/{task_id}/complete"
    code, data = http_json(
        url,
        method="POST",
        headers={"Authorization": env.get("PB_TOKEN", "")},
        body={"result": result, "log": log_lines, "duration_ms": duration_ms},
    )
    return code == 200 and isinstance(data, dict) and data.get("ok")


def fail_task(env: Dict[str, str], task_id: str, error_msg: str, log_lines: list) -> bool:
    api = env.get("TASKS_API", DEFAULT_TASKS_API)
    url = f"{api}/api/tasks/{task_id}/fail"
    code, data = http_json(
        url,
        method="POST",
        headers={"Authorization": env.get("PB_TOKEN", "")},
        body={"error": error_msg, "log": log_lines, "retry": True},
    )
    return code == 200 and isinstance(data, dict) and data.get("ok")


def lane_label(env: Dict[str, str]) -> str:
    return env.get("WORKER_LANE", "deterministic")


# ── Handlers ────────────────────────────────────────────────────

def handle_bizportal_search(env: Dict[str, str], payload: Dict) -> tuple[Any, list]:
    log = [f"bizportal_search payload={payload}"]
    service = env.get("BIZPORTAL_SERVICE_URL", DEFAULT_BIZPORTAL_SERVICE)
    url = f"{service}/search"
    code, data = http_json(url, method="POST", body=payload)
    log.append(f"service responded {code}")
    if code != 200 or not isinstance(data, dict):
        raise RuntimeError(f"bizportal service error: {code} {data}")
    if not data.get("ok"):
        raise RuntimeError(data.get("error") or "search failed")
    log.append(f"parsed {len(data.get('results') or [])} companies")
    return data, log


def handle_bizportal_reserve(env: Dict[str, str], payload: Dict) -> tuple[Any, list]:
    log = [f"bizportal_reserve payload={payload}"]
    service = env.get("BIZPORTAL_SERVICE_URL", DEFAULT_BIZPORTAL_SERVICE)
    url = f"{service}/reserve"
    code, data = http_json(url, method="POST", body=payload)
    log.append(f"service responded {code}")
    if code != 200 or not isinstance(data, dict):
        raise RuntimeError(f"bizportal service error: {code} {data}")
    if not data.get("ok"):
        raise RuntimeError(data.get("error") or "reserve failed")
    log.append(f"reservation status={data.get('status')}")
    return data, log


def handle_pdf_generate(env: Dict[str, str], payload: Dict) -> tuple[Any, list]:
    """Shell out to the pdf-creator skill. The skill is a directory
    that contains a runnable script — we invoke it with a JSON
    payload via stdin and read the result file path from stdout."""
    log = [f"pdf_generate payload={payload}"]
    skill = env.get("PDF_CREATOR_PATH", DEFAULT_PDF_CREATOR)
    if not Path(skill).exists():
        raise RuntimeError(f"pdf-creator not found at {skill}")
    # Look for an entry-point script
    entry = None
    for candidate in ("run.sh", "main.py", "pdf_creator.py", "index.js"):
        p = Path(skill) / candidate
        if p.exists():
            entry = p
            break
    if entry is None:
        raise RuntimeError(f"no entry-point script found in {skill}")
    log.append(f"invoking {entry}")
    proc = subprocess.run(
        [str(entry)],
        input=json.dumps(payload),
        capture_output=True,
        text=True,
        timeout=300,
    )
    log.append(f"rc={proc.returncode}")
    if proc.returncode != 0:
        log.append(f"stderr={proc.stderr[:500]}")
        raise RuntimeError(f"pdf-creator failed (rc={proc.returncode}): {proc.stderr[:200]}")
    result = {"stdout": proc.stdout.strip(), "stderr": proc.stderr.strip()}
    return result, log


def handle_email_send(env: Dict[str, str], payload: Dict) -> tuple[Any, list]:
    log = [f"email_send payload keys={list(payload.keys())}"]
    to = payload.get("to")
    subject = payload.get("subject", "")
    body = payload.get("body", "")
    if not to or not isinstance(to, str):
        raise RuntimeError("email_send: 'to' is required")
    smtp_host = env.get("SMTP_HOST")
    smtp_user = env.get("SMTP_USER")
    smtp_pass = env.get("SMTP_PASSWORD") or env.get("SMTP_PASS")
    smtp_from = env.get("SMTP_FROM", smtp_user or "noreply@example.com")
    smtp_port = int(env.get("SMTP_PORT", "587"))
    if not smtp_host:
        raise RuntimeError("SMTP_HOST not configured in .env.local")

    msg = MIMEText(body)
    msg["Subject"] = subject
    msg["From"] = smtp_from
    msg["To"] = to

    log.append(f"connecting to {smtp_host}:{smtp_port} as {smtp_user}")
    with smtplib.SMTP(smtp_host, smtp_port, timeout=30) as smtp:
        smtp.starttls()
        if smtp_user and smtp_pass:
            smtp.login(smtp_user, smtp_pass)
        smtp.send_message(msg)
    log.append("sent OK")
    return {"to": to, "subject": subject, "sent": True}, log


def handle_custom(env: Dict[str, str], task: Dict) -> tuple[Any, list]:
    """LLM lane — write the task to the agent's inbox directory. The
    agent is expected to PATCH the task back via /complete or /fail.
    We return a placeholder so the worker can mark this batch as
    "dispatched" but the actual completion comes from the agent."""
    inbox = Path(env.get("AGENT_INBOX", DEFAULT_AGENT_INBOX))
    inbox.mkdir(parents=True, exist_ok=True)
    payload = {
        "task_id": task["id"],
        "title": task["title"],
        "description": task.get("description"),
        "payload": task.get("payload"),
        "result_endpoint": f"{env.get('TASKS_API', DEFAULT_TASKS_API)}/api/tasks/{task['id']}/complete",
        "fail_endpoint": f"{env.get('TASKS_API', DEFAULT_TASKS_API)}/api/tasks/{task['id']}/fail",
    }
    out = inbox / f"{task['id']}.json"
    out.write_text(json.dumps(payload, indent=2))
    log_lines = [f"wrote {out}"]
    # Do NOT mark complete — the LLM agent will do that. The task
    # stays in 'in_progress' until the agent calls back.
    raise _InboxDispatched(out, log_lines)


class _InboxDispatched(Exception):
    """Special exception signalling the LLM agent has been notified
    and the worker should NOT call /complete — the agent will."""

    def __init__(self, out_path: Path, log_lines: list):
        self.out_path = out_path
        self.log_lines = log_lines


HANDLERS = {
    "bizportal_search": handle_bizportal_search,
    "bizportal_reserve": handle_bizportal_reserve,
    "pdf_generate": handle_pdf_generate,
    "email_send": handle_email_send,
    "custom": handle_custom,
}

# The custom lane only takes `custom` task type; the deterministic
# lane takes everything except `custom`.
DETERMINISTIC_TYPES = {"bizportal_search", "bizportal_reserve", "pdf_generate", "email_send"}


# ── Main loop ──────────────────────────────────────────────────

def run_one(env: Dict[str, str], lane: str, task: Dict) -> None:
    """Claim and execute a single task. Returns nothing — all state
    transitions go through the API."""
    task_type = task["task_type"]
    title = task["title"]

    if lane == "deterministic" and task_type not in DETERMINISTIC_TYPES:
        log.debug(f"skip {task['id']} (type={task_type} not in deterministic lane)")
        return
    if lane == "custom" and task_type != "custom":
        log.debug(f"skip {task['id']} (type={task_type} not in custom lane)")
        return

    log.info(f"claim {task_type} {task['id'][:12]}… '{title}'")
    claimed = claim_task(env, task["id"])
    if not claimed:
        log.info(f"  → already claimed by someone else")
        return

    handler = HANDLERS.get(task_type)
    if handler is None:
        fail_task(env, task["id"], f"no handler for task_type={task_type}", [f"unknown type {task_type}"])
        return

    t0 = time.time()
    log_lines: list = []
    try:
        payload = claimed.get("payload") or {}
        result, handler_log = handler(env, payload)
        log_lines.extend(handler_log)
        duration_ms = int((time.time() - t0) * 1000)
        complete_task(env, task["id"], result, log_lines, duration_ms)
        log.info(f"  → done in {duration_ms}ms")
    except _InboxDispatched as d:
        # LLM agent has been notified — log and move on
        log_lines.extend(d.log_lines)
        log.info(f"  → dispatched to agent inbox, awaiting callback")
    except Exception as e:
        log_lines.append(f"ERROR: {e}")
        duration_ms = int((time.time() - t0) * 1000)
        log.error(f"  → failed after {duration_ms}ms: {e}")
        fail_task(env, task["id"], str(e)[:4900], log_lines)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--lane", choices=["deterministic", "custom", "all"], default="all",
                        help="Which lane to run. 'all' = both as one process.")
    parser.add_argument("--once", action="store_true", help="Run one pass and exit (for testing).")
    args = parser.parse_args()

    env = load_env()
    env["WORKER_LANE"] = args.lane
    lanes = ["deterministic", "custom"] if args.lane == "all" else [args.lane]
    log.info(f"worker pid={os.getpid()} lane={lanes} poll={POLL_INTERVAL_S}s")
    log.info(f"PB_URL={env.get('PB_URL', DEFAULT_PB_URL)}")
    log.info(f"TASKS_API={env.get('TASKS_API', DEFAULT_TASKS_API)}")
    log.info(f"BIZPORTAL_SERVICE_URL={env.get('BIZPORTAL_SERVICE_URL', DEFAULT_BIZPORTAL_SERVICE)}")

    while True:
        any_claimed = False
        for lane in lanes:
            try:
                ready = list_ready(env, lane, limit=1)
            except Exception as e:
                log.warning(f"list_ready error: {e}")
                continue
            for task in ready:
                any_claimed = True
                try:
                    run_one(env, lane, task)
                except Exception as e:
                    log.exception(f"run_one crashed: {e}")
        if args.once:
            break
        if not any_claimed:
            time.sleep(POLL_INTERVAL_S)


if __name__ == "__main__":
    main()