"""
Doit Agent - Task Engine
Async queue-based task execution with persistence, retry, timeouts, and resource guards.
"""
from __future__ import annotations

import asyncio
import logging
import time
import uuid
from datetime import datetime, timezone
from typing import Any, Callable, Optional

import psutil

from core.config import (
    MAX_WORKERS, TASK_TIMEOUT_DEFAULT, TASK_MAX_RETRIES, TASK_RETRY_DELAY,
    MAX_CPU_PERCENT, MAX_MEMORY_MB
)
from persistence.database import get_db
from tools.registry import get_registry

logger = logging.getLogger("doit.tasks")


class ResourceGuard:
    """Monitor and enforce system resource limits."""

    def __init__(self):
        self.max_cpu = MAX_CPU_PERCENT
        self.max_memory_mb = MAX_MEMORY_MB

    async def check(self) -> tuple[bool, str]:
        """Check if system resources are within limits."""
        cpu = psutil.cpu_percent(interval=0.5)
        mem = psutil.virtual_memory()
        mem_used_mb = mem.used / 1e6

        if cpu > self.max_cpu:
            return False, f"CPU usage too high: {cpu:.1f}% (limit: {self.max_cpu}%)"
        if mem_used_mb > self.max_memory_mb * 10:  # System-wide memory guard (soft limit)
            return False, f"System memory critical: {mem_used_mb:.0f}MB used"
        return True, "ok"

    async def wait_for_resources(self, timeout: float = 30) -> bool:
        """Wait until resources are available. Returns False if timeout exceeded."""
        start = time.time()
        while time.time() - start < timeout:
            ok, _ = await self.check()
            if ok:
                return True
            await asyncio.sleep(2)
        return False


class Task:
    def __init__(
        self,
        task_type: str,
        payload: dict,
        task_id: str = None,
        max_retries: int = TASK_MAX_RETRIES,
        timeout: int = TASK_TIMEOUT_DEFAULT,
    ):
        self.id = task_id or str(uuid.uuid4())
        self.type = task_type
        self.payload = payload
        self.status = "pending"
        self.retry_count = 0
        self.max_retries = max_retries
        self.timeout = timeout
        self.result: Any = None
        self.error: str = None
        self.created_at = datetime.now(timezone.utc).isoformat()
        self.started_at: str = None
        self.completed_at: str = None

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "type": self.type,
            "payload": self.payload,
            "status": self.status,
            "retry_count": self.retry_count,
            "max_retries": self.max_retries,
            "timeout": self.timeout,
            "result": self.result,
            "error": self.error,
            "created_at": self.created_at,
            "started_at": self.started_at,
            "completed_at": self.completed_at,
        }


class TaskEngine:
    """
    Persistent async task queue.
    - Recovers pending tasks from DB on startup
    - Worker pool with configurable concurrency
    - Per-task timeout enforcement
    - Retry with exponential backoff
    - Resource guard integration
    - Result callbacks (for Telegram notifications)
    """

    def __init__(self):
        self._queue: asyncio.Queue = asyncio.Queue()
        self._workers: list[asyncio.Task] = []
        self._running: dict[str, Task] = {}
        self._guard = ResourceGuard()
        self._result_callbacks: list[Callable] = []
        self._shutdown = False
        self._safe_mode = False

    def on_result(self, callback: Callable) -> None:
        """Register a callback for task completion (e.g., send Telegram message)."""
        self._result_callbacks.append(callback)

    async def start(self) -> None:
        """Start worker pool and recover pending tasks."""
        logger.info("Starting task engine with %d workers", MAX_WORKERS)
        # Recover pending tasks from DB
        db = await get_db()
        pending = await db.tasks_pending()
        for row in pending:
            task = Task(
                task_type=row["type"],
                payload=self._parse_payload(row.get("payload", "{}")),
                task_id=row["id"],
                max_retries=row["max_retries"],
                timeout=row["timeout"],
            )
            task.retry_count = row["retry_count"]
            await self._queue.put(task)
            logger.info("Recovered pending task: %s (%s)", task.id, task.type)

        # Start workers
        for i in range(MAX_WORKERS):
            worker = asyncio.create_task(self._worker(i))
            self._workers.append(worker)

    async def stop(self) -> None:
        """Gracefully stop: finish running tasks, persist state."""
        logger.info("Task engine stopping...")
        self._shutdown = True

        # Cancel workers
        for w in self._workers:
            w.cancel()
        await asyncio.gather(*self._workers, return_exceptions=True)
        self._workers.clear()
        logger.info("Task engine stopped")

    def enter_safe_mode(self):
        self._safe_mode = True
        logger.info("Task engine: SAFE MODE ON")

    def exit_safe_mode(self):
        self._safe_mode = False
        logger.info("Task engine: SAFE MODE OFF")

    async def submit(self, task_type: str, payload: dict, **kwargs) -> Task:
        """Submit a new task to the queue."""
        task = Task(task_type, payload, **kwargs)
        db = await get_db()
        await db.task_insert(task.to_dict())
        await db.audit("task_submit", target=task.id, details={"type": task_type})
        await self._queue.put(task)
        logger.info("Task submitted: %s (%s)", task.id, task_type)
        return task

    async def submit_action(self, action: dict) -> Optional[Task]:
        """Submit a task from an AI action dict."""
        tool = action.get("tool")
        if not tool:
            # Multi-step
            steps = action.get("steps", [])
            if steps:
                # Submit as a pipeline task
                return await self.submit("pipeline", {"steps": steps})
            return None

        # Special non-tool actions
        if tool in ("clarify", "error"):
            return None

        return await self.submit("tool_call", {
            "tool": tool,
            "args": action.get("args", {}),
            "description": action.get("description", ""),
        })

    async def _worker(self, worker_id: int) -> None:
        logger.debug("Worker %d started", worker_id)
        while not self._shutdown:
            try:
                task = await asyncio.wait_for(self._queue.get(), timeout=1.0)
            except asyncio.TimeoutError:
                continue
            except asyncio.CancelledError:
                break

            if self._safe_mode:
                # Re-queue and wait
                await asyncio.sleep(2)
                await self._queue.put(task)
                continue

            # Resource check
            ok, reason = await self._guard.check()
            if not ok:
                logger.warning("Resource limit: %s — delaying task %s", reason, task.id)
                await asyncio.sleep(5)
                await self._queue.put(task)
                continue

            self._running[task.id] = task
            await self._execute_task(task)
            self._running.pop(task.id, None)
            self._queue.task_done()

    async def _execute_task(self, task: Task) -> None:
        db = await get_db()
        task.status = "running"
        task.started_at = datetime.now(timezone.utc).isoformat()
        await db.task_update(task.id, status="running", started_at=task.started_at)
        logger.info("Executing task %s (%s)", task.id, task.type)

        try:
            result = await asyncio.wait_for(
                self._run_task(task),
                timeout=task.timeout
            )
            task.status = "completed"
            task.result = result
            task.completed_at = datetime.now(timezone.utc).isoformat()
            await db.task_update(
                task.id, status="completed",
                result=result, completed_at=task.completed_at
            )
            await db.audit("task_complete", target=task.id, details={"type": task.type})
            logger.info("Task %s completed", task.id)

        except asyncio.TimeoutError:
            task.status = "failed"
            task.error = f"Task timed out after {task.timeout}s"
            await self._handle_failure(task, db)

        except Exception as e:
            task.status = "failed"
            task.error = str(e)
            logger.exception("Task %s failed", task.id)
            await self._handle_failure(task, db)

        # Notify callbacks
        for cb in self._result_callbacks:
            try:
                await cb(task)
            except Exception as e:
                logger.warning("Result callback error: %s", e)

    async def _handle_failure(self, task: Task, db) -> None:
        task.retry_count += 1
        if task.retry_count <= task.max_retries:
            delay = TASK_RETRY_DELAY * (2 ** (task.retry_count - 1))
            logger.info("Retrying task %s in %ds (attempt %d/%d)",
                        task.id, delay, task.retry_count, task.max_retries)
            await db.task_update(
                task.id, status="pending",
                retry_count=task.retry_count, error=task.error
            )
            await asyncio.sleep(delay)
            task.status = "pending"
            await self._queue.put(task)
        else:
            task.status = "failed"
            task.completed_at = datetime.now(timezone.utc).isoformat()
            await db.task_update(
                task.id, status="failed",
                error=task.error,
                retry_count=task.retry_count,
                completed_at=task.completed_at,
            )
            await db.audit("task_fail", target=task.id, details={"error": task.error}, status="error")

    async def _run_task(self, task: Task) -> Any:
        registry = get_registry()

        if task.type == "tool_call":
            tool_name = task.payload.get("tool")
            args = task.payload.get("args", {})
            if not registry.exists(tool_name):
                raise ValueError(f"Unknown tool: {tool_name}")
            return await registry.execute(tool_name, args)

        elif task.type == "pipeline":
            results = []
            steps = task.payload.get("steps", [])
            stop_on_error = task.payload.get("stop_on_error", True)
            for step in steps:
                tool_name = step.get("tool")
                args = step.get("args", {})
                description = step.get("description", tool_name)
                if not registry.exists(tool_name):
                    results.append({"tool": tool_name, "error": f"Unknown tool: {tool_name}"})
                    if stop_on_error:
                        break
                    continue
                result = await registry.execute(tool_name, args)
                results.append({"tool": tool_name, "description": description, "result": result})
                # Stop pipeline on tool error unless configured otherwise
                if "error" in result and stop_on_error:
                    results.append({"stopped": True, "reason": f"Step '{tool_name}' failed"})
                    break
            return {"steps": results, "total": len(steps), "completed": len(results)}

        else:
            raise ValueError(f"Unknown task type: {task.type}")

    def _parse_payload(self, payload_str: str) -> dict:
        import json
        try:
            return json.loads(payload_str)
        except Exception:
            return {}

    def status(self) -> dict:
        return {
            "queue_size": self._queue.qsize(),
            "running": len(self._running),
            "running_tasks": [
                {"id": t.id, "type": t.type}
                for t in self._running.values()
            ],
            "workers": MAX_WORKERS,
            "safe_mode": self._safe_mode,
        }


# Singleton
_engine: Optional[TaskEngine] = None


def get_engine() -> TaskEngine:
    global _engine
    if _engine is None:
        _engine = TaskEngine()
    return _engine
