#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""Main entry point for the PULSE Worker daemon."""
from __future__ import annotations

import argparse
import logging
import os
import signal
import sys
import time
from pathlib import Path

if os.name == "nt" and __name__ == "__main__":
    try:
        sys.stdout.reconfigure(encoding="utf-8")
        sys.stderr.reconfigure(encoding="utf-8")
    except (AttributeError, OSError):
        pass

ROOT_DIR = Path(__file__).resolve().parent.parent.parent.parent
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))
STATE_DIR = ROOT_DIR / "state"

from pulse.tasks.task_ops import update_registry
from pulse.workers.worker_llm import WorkerLLM
from pulse.workers.runner.worker_state import register_worker, deregister_worker, send_heartbeat
from pulse.workers.runner.worker_router import claim_task_atomic
from pulse.workers.runner.worker_execute import execute_task
from pulse.workers.runner.worker_memory import is_memory_task, execute_memory_task

POLL_INTERVAL = 5
MAX_TOOL_ITERATIONS = 25
HEARTBEAT_INTERVAL = 30
LOG_DIR = STATE_DIR / "worker_logs"

TOOL_NAMES = [
    "read_file", "write_file", "run_command", "brain_query",
    "list_files", "search_code", "edit_file", "submit_task", "complete_task",
]

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(name)s] %(levelname)s %(message)s", datefmt="%H:%M:%S")


class PulseWorker:
    """Main loop orchestrator for the worker daemon."""

    def __init__(self, provider: str, model: str, tier: str, worker_id: str, skills_used: bool = False):
        self.provider = provider
        self.model = model
        self.tier = tier
        self.worker_id = worker_id
        self.skills_used = skills_used
        self.running = True
        self.log = logging.getLogger(worker_id)
        self._last_heartbeat = 0.0
        self._llm = WorkerLLM(provider, model)
        LOG_DIR.mkdir(parents=True, exist_ok=True)

    def _heartbeat_cb(self, task_id=None, iteration=None, status="idle"):
        self._last_heartbeat = send_heartbeat(
            self.worker_id, self.model, self.tier, self._last_heartbeat,
            HEARTBEAT_INTERVAL, STATE_DIR, task_id, iteration, status
        )

    def run(self):
        register_worker(self.worker_id, self.model, self.tier, self.log, TOOL_NAMES)
        self.log.info("Worker started. Polling every %ds...", POLL_INTERVAL)
        while self.running:
            try:
                self._heartbeat_cb()
                task = claim_task_atomic(self.worker_id, self.tier)
                if task:
                    if is_memory_task(task):
                        execute_memory_task(task, self.worker_id, self.log)
                    else:
                        execute_task(
                            task, self.worker_id, self.model, self.tier, self.provider,
                            self.skills_used, self._llm, self.log, TOOL_NAMES,
                            MAX_TOOL_ITERATIONS, self._heartbeat_cb, LOG_DIR
                        )
                else:
                    time.sleep(POLL_INTERVAL)
            except KeyboardInterrupt:
                self.running = False
            except Exception as e:
                self.log.error("Loop error: %s", e, exc_info=True)
                time.sleep(POLL_INTERVAL)
        deregister_worker(self.worker_id)
        self.log.info("Worker deregistered and stopped.")

    def stop(self):
        self.running = False


def main():
    parser = argparse.ArgumentParser(description="PULSE Background Worker")
    parser.add_argument("--provider", required=True, choices=["anthropic", "google", "openai", "ollama"])
    parser.add_argument("--model", required=True)
    parser.add_argument("--tier", default="fast", choices=["fast", "standard", "premium"])
    parser.add_argument("--id", default=None)
    parser.add_argument("--skills-used", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()
    worker_id = args.id or f"worker_{args.provider}_{os.getpid()}"
    worker = PulseWorker(provider=args.provider, model=args.model, tier=args.tier,
                         worker_id=worker_id, skills_used=args.skills_used)
    if args.dry_run:
        print(f"PULSE Worker (dry-run) ID: {worker.worker_id}")
        return
    signal.signal(signal.SIGINT, lambda s, f: worker.stop())
    signal.signal(signal.SIGTERM, lambda s, f: worker.stop())
    worker.run()


if __name__ == "__main__":
    main()
