"""
Work Value Engine — protocol-generated task lifecycle and value settlement.

Tasks are auto-generated by protocol actions:
  register → validation task
  dispute  → arbitration task
  buy      → verification task
  flag     → moderation task

Worker nodes (validators/arbitrators) execute tasks using their AI compute,
get evaluated on quality, and earn OAS proportional to task value.
"""

from __future__ import annotations

import json
import os
import sqlite3
import threading
import time
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional


# ── Task types & status ─────────────────────────────────────────────

class TaskType(str, Enum):
    VALIDATION = "validation"
    ARBITRATION = "arbitration"
    VERIFICATION = "verification"
    MODERATION = "moderation"


class TaskStatus(str, Enum):
    PENDING = "pending"
    ASSIGNED = "assigned"
    COMPLETED = "completed"
    EVALUATED = "evaluated"
    SETTLED = "settled"
    FAILED = "failed"


# ── Base prices per task type (OAS) ─────────────────────────────────

TASK_BASE_PRICES: Dict[str, float] = {
    TaskType.VALIDATION: 2.0,
    TaskType.ARBITRATION: 50.0,
    TaskType.VERIFICATION: 5.0,
    TaskType.MODERATION: 3.0,
}


# ── Data model ──────────────────────────────────────────────────────

@dataclass
class WorkTask:
    task_id: str
    task_type: str                  # TaskType value
    trigger_tx: str                 # tx/asset that triggered this task
    status: str = TaskStatus.PENDING
    assigned_to: Optional[str] = None
    created_at: int = 0
    assigned_at: int = 0
    completed_at: int = 0
    settled_at: int = 0
    # Value fields
    base_value: float = 0.0
    quality_score: float = 0.0     # 0.0–1.0
    complexity: float = 1.0
    reputation_bonus: float = 0.0
    urgency: float = 1.0
    final_value: float = 0.0
    # Result
    result_data: str = ""          # JSON string of AI output

    def to_dict(self) -> Dict[str, Any]:
        return {
            "task_id": self.task_id,
            "task_type": self.task_type,
            "trigger_tx": self.trigger_tx,
            "status": self.status,
            "assigned_to": self.assigned_to,
            "created_at": self.created_at,
            "assigned_at": self.assigned_at,
            "completed_at": self.completed_at,
            "settled_at": self.settled_at,
            "base_value": self.base_value,
            "quality_score": self.quality_score,
            "complexity": self.complexity,
            "reputation_bonus": self.reputation_bonus,
            "urgency": self.urgency,
            "final_value": self.final_value,
            "result_data": self.result_data,
        }


# ── Value calculation ───────────────────────────────────────────────

def calculate_task_value(
    base_value: float,
    quality: float = 1.0,
    complexity: float = 1.0,
    reputation_bonus: float = 0.0,
    urgency: float = 1.0,
) -> float:
    """Calculate final OAS value for a completed task.

    Formula: base × quality × complexity × (1 + rep_bonus) × urgency
    """
    quality = max(0.0, min(1.0, quality))
    complexity = max(0.5, min(3.0, complexity))
    reputation_bonus = max(0.0, min(0.5, reputation_bonus))
    urgency = max(0.5, min(2.0, urgency))
    return round(base_value * quality * complexity * (1.0 + reputation_bonus) * urgency, 6)


# ── Engine ──────────────────────────────────────────────────────────

class WorkValueEngine:
    """Manages the full task lifecycle: create → assign → complete → evaluate → settle."""

    def __init__(self, db_path: Optional[str] = None):
        if db_path is None:
            db_path = os.path.join(os.path.expanduser("~"), ".oasyce", "work.db")
        self.db_path = db_path
        os.makedirs(os.path.dirname(db_path), exist_ok=True)
        self._lock = threading.Lock()
        self._conn = sqlite3.connect(db_path, check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._create_tables()

    def _create_tables(self) -> None:
        with self._lock, self._conn:
            self._conn.execute("""
                CREATE TABLE IF NOT EXISTS work_tasks (
                    task_id       TEXT PRIMARY KEY,
                    task_type     TEXT NOT NULL,
                    trigger_tx    TEXT NOT NULL,
                    status        TEXT NOT NULL DEFAULT 'pending',
                    assigned_to   TEXT,
                    created_at    INTEGER NOT NULL,
                    assigned_at   INTEGER DEFAULT 0,
                    completed_at  INTEGER DEFAULT 0,
                    settled_at    INTEGER DEFAULT 0,
                    base_value    REAL NOT NULL,
                    quality_score REAL DEFAULT 0.0,
                    complexity    REAL DEFAULT 1.0,
                    reputation_bonus REAL DEFAULT 0.0,
                    urgency       REAL DEFAULT 1.0,
                    final_value   REAL DEFAULT 0.0,
                    result_data   TEXT DEFAULT ''
                )
            """)
            self._conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_work_status ON work_tasks(status)
            """)
            self._conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_work_assigned ON work_tasks(assigned_to)
            """)

    # ── Row ↔ WorkTask ──────────────────────────────────────────────

    @staticmethod
    def _row_to_task(row: sqlite3.Row) -> WorkTask:
        return WorkTask(
            task_id=row["task_id"],
            task_type=row["task_type"],
            trigger_tx=row["trigger_tx"],
            status=row["status"],
            assigned_to=row["assigned_to"],
            created_at=row["created_at"],
            assigned_at=row["assigned_at"],
            completed_at=row["completed_at"],
            settled_at=row["settled_at"],
            base_value=row["base_value"],
            quality_score=row["quality_score"],
            complexity=row["complexity"],
            reputation_bonus=row["reputation_bonus"],
            urgency=row["urgency"],
            final_value=row["final_value"],
            result_data=row["result_data"] or "",
        )

    # ── Create ──────────────────────────────────────────────────────

    def create_task(
        self,
        task_type: str,
        trigger_tx: str,
        complexity: float = 1.0,
        urgency: float = 1.0,
    ) -> WorkTask:
        """Create a new protocol-generated task."""
        task_id = f"wt_{uuid.uuid4().hex[:12]}"
        base_value = TASK_BASE_PRICES.get(task_type, 2.0)
        now = int(time.time())

        task = WorkTask(
            task_id=task_id,
            task_type=task_type,
            trigger_tx=trigger_tx,
            status=TaskStatus.PENDING,
            created_at=now,
            base_value=base_value,
            complexity=complexity,
            urgency=urgency,
        )

        with self._lock, self._conn:
            self._conn.execute(
                """INSERT INTO work_tasks
                   (task_id, task_type, trigger_tx, status, created_at,
                    base_value, complexity, urgency)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
                (task_id, task_type, trigger_tx, TaskStatus.PENDING,
                 now, base_value, complexity, urgency),
            )
        return task

    # ── Assign ──────────────────────────────────────────────────────

    def assign_task(self, task_id: str, worker_id: str) -> Optional[WorkTask]:
        """Assign a pending task to a worker node."""
        now = int(time.time())
        with self._lock, self._conn:
            row = self._conn.execute(
                "SELECT * FROM work_tasks WHERE task_id = ?", (task_id,)
            ).fetchone()
            if row is None or row["status"] != TaskStatus.PENDING:
                return None
            self._conn.execute(
                """UPDATE work_tasks
                   SET status = ?, assigned_to = ?, assigned_at = ?
                   WHERE task_id = ?""",
                (TaskStatus.ASSIGNED, worker_id, now, task_id),
            )
        return self.get_task(task_id)

    # ── Complete ────────────────────────────────────────────────────

    def complete_task(self, task_id: str, result_data: str = "") -> Optional[WorkTask]:
        """Mark a task as completed with AI result data."""
        now = int(time.time())
        with self._lock, self._conn:
            row = self._conn.execute(
                "SELECT * FROM work_tasks WHERE task_id = ?", (task_id,)
            ).fetchone()
            if row is None or row["status"] != TaskStatus.ASSIGNED:
                return None
            self._conn.execute(
                """UPDATE work_tasks
                   SET status = ?, completed_at = ?, result_data = ?
                   WHERE task_id = ?""",
                (TaskStatus.COMPLETED, now, result_data, task_id),
            )
        return self.get_task(task_id)

    # ── Evaluate ────────────────────────────────────────────────────

    def evaluate_task(
        self,
        task_id: str,
        quality_score: float,
        reputation_bonus: float = 0.0,
    ) -> Optional[WorkTask]:
        """Evaluate a completed task and calculate final value."""
        with self._lock, self._conn:
            row = self._conn.execute(
                "SELECT * FROM work_tasks WHERE task_id = ?", (task_id,)
            ).fetchone()
            if row is None or row["status"] != TaskStatus.COMPLETED:
                return None

            final = calculate_task_value(
                base_value=row["base_value"],
                quality=quality_score,
                complexity=row["complexity"],
                reputation_bonus=reputation_bonus,
                urgency=row["urgency"],
            )
            self._conn.execute(
                """UPDATE work_tasks
                   SET status = ?, quality_score = ?, reputation_bonus = ?,
                       final_value = ?
                   WHERE task_id = ?""",
                (TaskStatus.EVALUATED, quality_score, reputation_bonus,
                 final, task_id),
            )
        return self.get_task(task_id)

    # ── Settle ──────────────────────────────────────────────────────

    def settle_task(self, task_id: str) -> Optional[WorkTask]:
        """Mark an evaluated task as settled (OAS paid out)."""
        now = int(time.time())
        with self._lock, self._conn:
            row = self._conn.execute(
                "SELECT * FROM work_tasks WHERE task_id = ?", (task_id,)
            ).fetchone()
            if row is None or row["status"] != TaskStatus.EVALUATED:
                return None
            self._conn.execute(
                """UPDATE work_tasks
                   SET status = ?, settled_at = ?
                   WHERE task_id = ?""",
                (TaskStatus.SETTLED, now, task_id),
            )
        return self.get_task(task_id)

    # ── Fail ────────────────────────────────────────────────────────

    def fail_task(self, task_id: str) -> Optional[WorkTask]:
        """Mark a task as failed (worker didn't deliver or quality too low)."""
        with self._lock, self._conn:
            self._conn.execute(
                "UPDATE work_tasks SET status = ? WHERE task_id = ?",
                (TaskStatus.FAILED, task_id),
            )
        return self.get_task(task_id)

    # ── Query ───────────────────────────────────────────────────────

    def get_task(self, task_id: str) -> Optional[WorkTask]:
        with self._lock:
            row = self._conn.execute(
                "SELECT * FROM work_tasks WHERE task_id = ?", (task_id,)
            ).fetchone()
        return self._row_to_task(row) if row else None

    def list_tasks(
        self,
        status: Optional[str] = None,
        worker_id: Optional[str] = None,
        task_type: Optional[str] = None,
        limit: int = 50,
    ) -> List[WorkTask]:
        clauses: List[str] = []
        params: List[Any] = []
        if status:
            clauses.append("status = ?")
            params.append(status)
        if worker_id:
            clauses.append("assigned_to = ?")
            params.append(worker_id)
        if task_type:
            clauses.append("task_type = ?")
            params.append(task_type)
        where = " WHERE " + " AND ".join(clauses) if clauses else ""
        params.append(limit)
        with self._lock:
            rows = self._conn.execute(
                f"SELECT * FROM work_tasks{where} ORDER BY created_at DESC LIMIT ?",
                params,
            ).fetchall()
        return [self._row_to_task(r) for r in rows]

    def pending_tasks(self, task_type: Optional[str] = None, limit: int = 20) -> List[WorkTask]:
        """Get pending tasks available for assignment."""
        return self.list_tasks(status=TaskStatus.PENDING, task_type=task_type, limit=limit)

    def worker_stats(self, worker_id: str) -> Dict[str, Any]:
        """Get stats for a worker node."""
        with self._lock:
            total = self._conn.execute(
                "SELECT COUNT(*) AS c FROM work_tasks WHERE assigned_to = ?",
                (worker_id,),
            ).fetchone()["c"]
            settled = self._conn.execute(
                "SELECT COUNT(*) AS c FROM work_tasks WHERE assigned_to = ? AND status = ?",
                (worker_id, TaskStatus.SETTLED),
            ).fetchone()["c"]
            earned_row = self._conn.execute(
                "SELECT COALESCE(SUM(final_value), 0) AS s FROM work_tasks "
                "WHERE assigned_to = ? AND status = ?",
                (worker_id, TaskStatus.SETTLED),
            ).fetchone()
            earned = earned_row["s"]
            avg_quality_row = self._conn.execute(
                "SELECT AVG(quality_score) AS q FROM work_tasks "
                "WHERE assigned_to = ? AND status IN (?, ?)",
                (worker_id, TaskStatus.EVALUATED, TaskStatus.SETTLED),
            ).fetchone()
            avg_quality = avg_quality_row["q"] or 0.0
            failed = self._conn.execute(
                "SELECT COUNT(*) AS c FROM work_tasks WHERE assigned_to = ? AND status = ?",
                (worker_id, TaskStatus.FAILED),
            ).fetchone()["c"]
        return {
            "worker_id": worker_id,
            "total_tasks": total,
            "settled": settled,
            "failed": failed,
            "total_earned": round(earned, 6),
            "avg_quality": round(avg_quality, 4),
        }

    def global_stats(self) -> Dict[str, Any]:
        """Get global work system stats."""
        with self._lock:
            total = self._conn.execute(
                "SELECT COUNT(*) AS c FROM work_tasks"
            ).fetchone()["c"]
            by_status = {}
            for row in self._conn.execute(
                "SELECT status, COUNT(*) AS c FROM work_tasks GROUP BY status"
            ).fetchall():
                by_status[row["status"]] = row["c"]
            total_value = self._conn.execute(
                "SELECT COALESCE(SUM(final_value), 0) AS s FROM work_tasks WHERE status = ?",
                (TaskStatus.SETTLED,),
            ).fetchone()["s"]
        return {
            "total_tasks": total,
            "by_status": by_status,
            "total_value_settled": round(total_value, 6),
        }

    def close(self) -> None:
        self._conn.close()
