# backend/tasks/executor.py
import asyncio
from dataclasses import dataclass, field
from typing import Any
from uuid import uuid4

@dataclass
class Task:
    id: str = field(default_factory=lambda: str(uuid4()))
    name: str = ""
    steps: list[dict] = field(default_factory=list)
    current_step: int = 0
    status: str = "pending"  # pending | running | paused | completed | failed | cancelled
    result: Any = None
    error: str = ""
    permission_tier: str = "approval_required"  # pre_approved | approval_required | explicit_each_time | biometric

class TaskExecutor:
    def __init__(self, app_handle, permission_gateway):
        self.tasks: dict[str, Task] = {}
        self.app_handle = app_handle
        self.permission_gateway = permission_gateway
        
        # Load db path
        import sys, os
        if getattr(sys, 'frozen', False):
            self.db_path = os.path.join(os.environ.get('APPDATA', os.path.expanduser('~')), 'JARVIS_OS', 'memory.db')
        else:
            self.db_path = os.path.join(os.getcwd(), 'memory.db')
            
        self._recover_tasks()

    def _recover_tasks(self):
        import sqlite3
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.execute('PRAGMA journal_mode=WAL')
                cursor = conn.cursor()
                cursor.execute("SELECT id, name, status, steps, current_step, result, error, permission_tier FROM tasks")
                for row in cursor.fetchall():
                    status = row[2]
                    # Safety: Fail any tasks that were running/paused during a hard crash
                    if status in ["running", "paused"]:
                        status = "failed"
                        cursor.execute("UPDATE tasks SET status=?, error=? WHERE id=?", 
                                       ("failed", "System rebooted mid-execution", row[0]))
                        conn.commit()
                    
                    t = Task(
                        id=row[0], name=row[1], status=status, 
                        steps=[], # Callables cannot be easily serialized to JSON, so we leave steps empty on recovery
                        current_step=row[4], result=row[5], error=row[6], permission_tier=row[7]
                    )
                    self.tasks[t.id] = t
        except Exception as e:
            import logging; logging.getLogger(__name__).error(f"Swallowed exception: {e}") # Tables might not exist yet

    def _save_task(self, task: Task):
        import sqlite3, json
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.execute('PRAGMA journal_mode=WAL')
                cursor = conn.cursor()
                # We do not serialize the actual 'handler' since it's a function reference. 
                # We just serialize a dummy representation of steps for UI if needed.
                steps_json = json.dumps([{"name": getattr(s.get('handler'), '__name__', 'step')} for s in task.steps])
                cursor.execute('''
                    INSERT INTO tasks (id, name, status, steps, current_step, result, error, permission_tier)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                    ON CONFLICT(id) DO UPDATE SET
                        status=excluded.status,
                        current_step=excluded.current_step,
                        result=excluded.result,
                        error=excluded.error
                ''', (task.id, task.name, task.status, steps_json, task.current_step, str(task.result) if task.result else "", task.error, task.permission_tier))
                conn.commit()
        except Exception as e:
            import logging; logging.getLogger(__name__).error(f"Swallowed exception: {e}")

    async def submit(self, task: Task) -> str:
        self.tasks[task.id] = task
        self._save_task(task)
        asyncio.create_task(self._run(task))
        return task.id

    async def _run(self, task: Task):
        task.status = "running"
        self._save_task(task)
        self._emit("task:updated", task)
        for i, step in enumerate(task.steps):
            task.current_step = i
            self._save_task(task)
            self._emit("task:step_started", {"task_id": task.id, "step": step})
            allowed = await self.permission_gateway.check(task.permission_tier, step)
            if not allowed:
                task.status = "paused"
                self._save_task(task)
                self._emit("task:awaiting_permission", {"task_id": task.id, "step": step})
                allowed = await self.permission_gateway.wait_for_user_response(task.id)
                if not allowed:
                    task.status = "cancelled"
                    self._save_task(task)
                    self._emit("task:updated", task)
                    return
            try:
                result = await step["handler"](**step["params"])
                task.result = result
                self._save_task(task)
                self._emit("task:step_completed", {"task_id": task.id, "step": step, "result": result})
            except Exception as e:
                task.status = "failed"
                task.error = str(e)
                self._save_task(task)
                self._emit("task:updated", task)
                return
        task.status = "completed"
        self._save_task(task)
        self._emit("task:updated", task)

    def cancel(self, task_id: str):
        task = self.tasks.get(task_id)
        if task: 
            task.status = "cancelled"
            self._save_task(task)
        self._emit("task:updated", task)

    def _emit(self, event: str, payload):
        try:
            from backend.ws.agent_ws import ws_manager
            import json
            from dataclasses import asdict
            
            # Serialize payload to dict safely
            if hasattr(payload, '__dataclass_fields__'):
                safe_payload = asdict(payload)
            else:
                # Fallback via json roundtrip to handle non-serializable objects cleanly
                safe_payload = json.loads(json.dumps(payload, default=str))
                
            asyncio.create_task(ws_manager.broadcast({"event": event, "payload": safe_payload}))
        except Exception as e:
            import logging; logging.getLogger(__name__).error(f"Swallowed exception: {e}")