"""
BizNode Distributed AI Execution Engine (DAEE)
================================================
PostgreSQL-backed distributed task engine for all bot + AI workloads.

Architecture:
    Telegram/API → Task Ingestion → PostgreSQL Queue → Dispatcher → Workers
    Workers: rag, email, payment, automation, blockchain, video

Task flow:
    1. Create task (API or bot command)
    2. Task queued in PostgreSQL with priority
    3. Worker fetches with SELECT ... FOR UPDATE SKIP LOCKED
    4. Worker executes and updates status
    5. Result returned to caller
"""

import json
import logging
import os
import sys
import threading
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

_PROJECT_ROOT = str(Path(__file__).resolve().parent.parent)
if _PROJECT_ROOT not in sys.path:
    sys.path.insert(0, _PROJECT_ROOT)

log = logging.getLogger(__name__)

# Task type → handler mapping
_TASK_HANDLERS: Dict[str, Callable] = {}
_worker_threads: Dict[str, threading.Thread] = {}
_stop_event = threading.Event()


# Priority levels by tier
TIER_PRIORITY = {
    "admin": 100,
    "onezbnode": 90,
    "biznode": 80,
    "bizmax": 70,
    "bizplus": 60,
    "pro": 50,
    "standard": 30,
    "basic": 20,
    "trial": 10,
}

# Daily task limits by tier
TIER_TASK_LIMITS = {
    "admin": 999999,
    "onezbnode": 500,
    "biznode": 500,
    "bizmax": 200,
    "bizplus": 100,
    "pro": 100,
    "standard": 50,
    "basic": 20,
    "trial": 5,
}


def register_handler(task_type: str, handler: Callable):
    """Register a handler function for a task type."""
    _TASK_HANDLERS[task_type] = handler
    log.info("Task handler registered: %s", task_type)


def get_priority(agent_id: str = "", tier: str = "basic") -> int:
    """Get task priority based on tier."""
    return TIER_PRIORITY.get(tier, 10)


# ------------------------------------------------------------------
# Task creation
# ------------------------------------------------------------------

def submit_task(
    task_type: str,
    payload: Dict = None,
    agent_id: str = "",
    user_id: str = "",
    priority: int = None,
    tier: str = "basic",
    scheduled_at: str = "",
    max_retries: int = 3,
) -> Dict[str, Any]:
    """Submit a new task to the queue."""
    from memory.database import create_task, log_task_event

    if task_type not in _TASK_HANDLERS and task_type not in (
        "rag", "email", "payment", "automation", "blockchain", "video",
        "kb_search", "tool_call", "bot_response", "report", "notification",
    ):
        return {"success": False, "error": f"Unknown task type: {task_type}"}

    task_id = f"task-{uuid.uuid4().hex[:12]}"
    if priority is None:
        priority = get_priority(agent_id, tier)

    payload_json = json.dumps(payload or {})

    try:
        create_task(
            task_id=task_id,
            task_type=task_type,
            payload=payload_json,
            priority=priority,
            agent_id=agent_id,
            user_id=user_id,
            max_retries=max_retries,
            scheduled_at=scheduled_at,
        )
        log_task_event(task_id, f"Task created: {task_type}", "queued")
        return {"success": True, "task_id": task_id, "status": "queued", "priority": priority}
    except Exception as exc:
        log.error("Failed to create task: %s", exc)
        return {"success": False, "error": str(exc)}


def get_task_status(task_id: str) -> Dict[str, Any]:
    """Get task status and result."""
    from memory.database import get_task, get_task_logs

    task = get_task(task_id)
    if not task:
        return {"success": False, "error": "Task not found"}

    logs = get_task_logs(task_id)
    result = task.get("result")
    if result:
        try:
            result = json.loads(result)
        except (json.JSONDecodeError, TypeError):
            pass

    return {
        "success": True,
        "task_id": task_id,
        "task_type": task.get("task_type"),
        "status": task.get("status"),
        "priority": task.get("priority"),
        "result": result,
        "error": task.get("error_message"),
        "retries": task.get("retries"),
        "created_at": task.get("created_at"),
        "started_at": task.get("started_at"),
        "completed_at": task.get("completed_at"),
        "logs": [{"message": l.get("log_message"), "status": l.get("status"), "at": l.get("created_at")} for l in logs],
    }


def get_queue_stats() -> Dict[str, Any]:
    """Get overall queue statistics."""
    from memory.database import get_task_stats, get_active_workers

    stats = get_task_stats()
    workers = get_active_workers()
    return {
        "success": True,
        "queue": stats,
        "workers": [
            {
                "node_id": w["node_id"],
                "status": w["status"],
                "capabilities": w.get("capabilities"),
                "tasks_completed": w.get("tasks_completed", 0),
                "last_heartbeat": w.get("last_heartbeat"),
            }
            for w in workers
        ],
    }


# ------------------------------------------------------------------
# Worker engine
# ------------------------------------------------------------------

class TaskWorker:
    """Base worker that pulls and executes tasks from PostgreSQL."""

    def __init__(self, node_id: str, task_types: List[str] = None, poll_interval: int = 2):
        self.node_id = node_id
        self.task_types = task_types
        self.poll_interval = poll_interval

    def start(self):
        """Start the worker in a background thread."""
        from memory.database import register_worker
        caps = json.dumps(self.task_types or ["all"])
        register_worker(self.node_id, caps)
        log.info("Worker %s started (types: %s)", self.node_id, self.task_types or "all")

        thread = threading.Thread(target=self._run_loop, daemon=True, name=f"worker-{self.node_id}")
        thread.start()
        _worker_threads[self.node_id] = thread

    def _run_loop(self):
        """Main worker loop — fetch and execute tasks."""
        from memory.database import (
            fetch_next_task, complete_task, fail_task,
            log_task_event, worker_heartbeat,
        )

        while not _stop_event.is_set():
            try:
                worker_heartbeat(self.node_id, status="polling")
                task = fetch_next_task(self.node_id, self.task_types)

                if not task:
                    time.sleep(self.poll_interval)
                    continue

                task_id = task["id"]
                task_type = task["task_type"]
                payload = task.get("payload", "{}")

                try:
                    payload_dict = json.loads(payload) if isinstance(payload, str) else payload
                except (json.JSONDecodeError, TypeError):
                    payload_dict = {}

                worker_heartbeat(self.node_id, status="executing", current_task=task_id)
                log_task_event(task_id, f"Worker {self.node_id} started execution", "running", self.node_id)

                # Find and execute handler
                handler = _TASK_HANDLERS.get(task_type)
                if handler:
                    try:
                        result = handler(payload_dict)
                        result_json = json.dumps(result) if isinstance(result, (dict, list)) else str(result)
                        complete_task(task_id, result_json)
                        log_task_event(task_id, "Task completed successfully", "completed", self.node_id)
                        log.info("Task %s completed by %s", task_id, self.node_id)
                    except Exception as exc:
                        fail_task(task_id, str(exc))
                        log_task_event(task_id, f"Task failed: {exc}", "failed", self.node_id)
                        log.error("Task %s failed: %s", task_id, exc)
                else:
                    # Try built-in handlers
                    result = self._builtin_handler(task_type, payload_dict)
                    if result is not None:
                        complete_task(task_id, json.dumps(result))
                        log_task_event(task_id, "Completed via built-in handler", "completed", self.node_id)
                    else:
                        fail_task(task_id, f"No handler for task type: {task_type}")
                        log_task_event(task_id, f"No handler found for: {task_type}", "failed", self.node_id)

                worker_heartbeat(self.node_id, status="idle", current_task="")

            except Exception as exc:
                log.error("Worker %s error: %s", self.node_id, exc)
                time.sleep(self.poll_interval)

    def _builtin_handler(self, task_type: str, payload: Dict) -> Optional[Dict]:
        """Built-in handlers for common task types."""
        if task_type == "email":
            return self._handle_email(payload)
        elif task_type == "bot_response":
            return self._handle_bot_response(payload)
        elif task_type == "kb_search":
            return self._handle_kb_search(payload)
        elif task_type == "notification":
            return self._handle_notification(payload)
        elif task_type == "report":
            return self._handle_report(payload)
        elif task_type == "tool_call":
            return self._handle_tool_call(payload)
        return None

    def _handle_tool_call(self, payload: Dict) -> Dict:
        """Execute a published tool and notify the user via Telegram."""
        tool_id = payload.get("tool_id", "")
        caller_agent_id = payload.get("caller_agent_id", "")
        chat_id = payload.get("chat_id", "")
        user_name = payload.get("user_name", "User")

        # Execute tool
        try:
            from services.tool_publishing_service import ToolPublishingService
            svc = ToolPublishingService()
            result = svc.call_tool(tool_id, caller_agent_id, payload.get("params", {}))
        except Exception as exc:
            result = {"success": False, "error": str(exc)}

        # Send result back to user via Telegram if chat_id provided
        if chat_id:
            try:
                import requests as _req
                from memory.database import get_hosted_agent
                agent = get_hosted_agent(caller_agent_id)
                bot_token = agent.get("bot_token", "") if agent else ""
                if bot_token:
                    if result.get("success"):
                        text = f"Action completed!\n\n{json.dumps(result.get('result', ''), indent=2)[:500]}"
                    else:
                        text = f"Action could not be completed: {result.get('error', 'Unknown error')}"
                    _req.post(
                        f"https://api.telegram.org/bot{bot_token}/sendMessage",
                        json={"chat_id": chat_id, "text": text},
                        timeout=10,
                    )
            except Exception as exc:
                log.debug("Tool result notification failed: %s", exc)

        return result

    def _handle_email(self, payload: Dict) -> Dict:
        """Send email task."""
        from services.email_service import EmailService
        es = EmailService()
        result = es.send_html_email(
            payload.get("to", ""),
            payload.get("subject", "BizNode Notification"),
            payload.get("html", payload.get("body", "")),
        )
        return result

    def _handle_bot_response(self, payload: Dict) -> Dict:
        """Generate bot response via LLM."""
        import requests
        try:
            resp = requests.post(
                "http://localhost:11434/api/chat",
                json={
                    "model": payload.get("model", "qwen3.5:4b"),
                    "messages": payload.get("messages", []),
                    "stream": False,
                },
                timeout=60,
            )
            if resp.ok:
                return {"response": resp.json().get("message", {}).get("content", "")}
        except Exception as exc:
            return {"error": str(exc)}
        return {"error": "LLM unavailable"}

    def _handle_kb_search(self, payload: Dict) -> Dict:
        """Search knowledge base."""
        import requests
        query = payload.get("query", "")
        agent_id = payload.get("agent_id", "")
        try:
            # Get embedding
            emb_resp = requests.post(
                "http://localhost:11434/api/embeddings",
                json={"model": "nomic-embed-text", "prompt": query},
                timeout=15,
            )
            if not emb_resp.ok:
                return {"results": [], "error": "Embedding failed"}

            embedding = emb_resp.json().get("embedding", [])
            search_body = {"vector": embedding, "limit": 5, "with_payload": True}
            if agent_id:
                search_body["filter"] = {"must": [{"key": "agent_id", "match": {"value": agent_id}}]}

            resp = requests.post(
                "http://localhost:6333/collections/biznode_support/points/search",
                json=search_body, timeout=10,
            )
            if resp.ok:
                results = [p.get("payload", {}).get("text", "") for p in resp.json().get("result", [])]
                return {"results": results}
        except Exception as exc:
            return {"results": [], "error": str(exc)}
        return {"results": []}

    def _handle_notification(self, payload: Dict) -> Dict:
        """Send Telegram notification."""
        import requests
        bot_token = payload.get("bot_token", os.getenv("TELEGRAM_BOT_TOKEN", ""))
        chat_id = payload.get("chat_id", "")
        text = payload.get("text", "")
        if not bot_token or not chat_id:
            return {"error": "Missing bot_token or chat_id"}
        try:
            resp = requests.post(
                f"https://api.telegram.org/bot{bot_token}/sendMessage",
                json={"chat_id": chat_id, "text": text},
                timeout=10,
            )
            return {"sent": resp.ok}
        except Exception as exc:
            return {"error": str(exc)}

    def _handle_report(self, payload: Dict) -> Dict:
        """Generate a report task."""
        report_type = payload.get("type", "daily")
        agent_id = payload.get("agent_id", "")
        if report_type == "daily":
            try:
                from services.hosted_bot_runner import HostedBotRunner
                runner = HostedBotRunner()
                return runner.generate_daily_reports()
            except Exception as exc:
                return {"error": str(exc)}
        return {"error": f"Unknown report type: {report_type}"}


# ------------------------------------------------------------------
# Engine management
# ------------------------------------------------------------------

def start_engine(num_workers: int = 2, task_types: List[str] = None):
    """Start the task engine with multiple workers."""
    _stop_event.clear()

    # Register built-in handlers
    _register_builtin_handlers()

    for i in range(num_workers):
        node_id = f"worker-{i+1}"
        worker = TaskWorker(node_id, task_types)
        worker.start()

    log.info("Task engine started with %d workers", num_workers)


def stop_engine():
    """Stop all workers."""
    _stop_event.set()
    for node_id, thread in _worker_threads.items():
        thread.join(timeout=5)
    _worker_threads.clear()
    log.info("Task engine stopped")


def _register_builtin_handlers():
    """Register built-in task handlers from the tool registry."""
    try:
        from automation.registry import get_registry
        registry = get_registry()
        for tool_name in registry.list_tools():
            tool = registry.get(tool_name)
            def make_handler(t):
                return lambda payload: t.run(**payload)
            register_handler(f"tool:{tool_name}", make_handler(tool))
    except Exception:
        pass
