# hiveos/task_runtime.py
from hiveos.db import get_db
#from hiveos.runtime.chunk_runtime import chunk_objects
import json
from hiveos.runtime.states import ChunkStatus, TaskStatus

#task_objects = {}

class RuntimeTask:
    def __init__(self, task_id, name, execution_mode="sequential", has_conditionals=False, zone_id=None, capabilities=None, session_id=None):
        self.task_id = task_id
        self.name = name
        self.zone_id = zone_id
        self.execution_mode = execution_mode
        self.status = TaskStatus.PENDING
        self.has_conditionals = has_conditionals
        self._capabilities = tuple(capabilities or ())
        self.session_id = session_id
        self._capabilities_locked = True

    @property
    def capabilities(self):
        return self._capabilities

    @capabilities.setter
    def capabilities(self, value):
        if getattr(self, "_capabilities_locked", False):
            raise AttributeError("task capabilities are immutable after injection")
        self._capabilities = tuple(value or ())

    def is_complete(self, chunk_objects):
        return all(
            chunk.status in (ChunkStatus.COMPLETED, ChunkStatus.STALE)
            for chunk in chunk_objects.values()
            if chunk.task_id == self.task_id
        )

    def mark_complete(self):
        self.status = TaskStatus.COMPLETED
        
    def to_dict(self):
        return {
            "task_id": self.task_id,
            "name": self.name,
            "status": self.status,
            "execution_mode": self.execution_mode,
            "has_conditionals": self.has_conditionals,
            "zone_id": self.zone_id,
            "session_id": self.session_id,
        }
        
    def to_yaml_dict(self, chunk_objects):
        yaml_task = {
            "task_id": self.task_id,
            "task_name": self.name,
            "execution_mode": self.execution_mode,
            "capability_required": list(self.capabilities),
            "chunks": []
        }

        if self.has_conditionals:
            yaml_task["has_conditionals"] = True

        # Include associated chunks
        for chunk in chunk_objects.values():
            if chunk.task_id != self.task_id:
                continue

            chunk_yaml = {
                "chunk_id": chunk.chunk_id.rsplit("_", 1)[0],  # remove suffix
                "required_capability": chunk.required_capability,
                "payload": json.loads(chunk.payload) if isinstance(chunk.payload, str) else (chunk.payload or {}),
                "ttl": chunk.ttl,
                "timing": chunk.timing,
            }

            if chunk.depends_on:
                base_depends = chunk.depends_on.rsplit("_", 1)[0]
                chunk_yaml["depends_on"] = base_depends

            if chunk.conditional:
                if isinstance(chunk.conditional, str):
                    cond = json.loads(chunk.conditional)
                else:
                    cond = chunk.conditional
                if cond:
                    chunk_yaml["conditional"] = cond

            yaml_task["chunks"].append(chunk_yaml)

        return yaml_task


# def mark_task_complete(task_id):
#     task = task_objects.get(task_id)
#     if task:
#         task.mark_complete()
        
"""
def all_chunks_completed(task_id):
    return all(
        chunk.status == "completed"
        for chunk in chunk_objects.values()
        if chunk.task_id == task_id
    )
"""

# def register(task_obj):
#     task_objects[task_obj.task_id] = task_obj

# def get_by_id(task_id):
#     return task_objects.get(task_id)

def hydrate_from_db_row(task_obj, session_id, capabilities=None):
    task = RuntimeTask(
        task_id=task_obj.task_id,
        name=task_obj.name,
        execution_mode=task_obj.execution_mode,
        has_conditionals=getattr(task_obj, "has_conditionals", False),
        capabilities=capabilities or [],
        session_id=session_id
    )
    # register(task)
    return task
    
def has_conditionals(self):
    return self.has_conditionals


def flush_task_db_state(task_objects: dict):
    try:
        with get_db() as conn:
            for task in task_objects.values():
                conn.execute(
                    "UPDATE tasks SET status = ? WHERE task_id = ?",
                    (task.status, task.task_id)
                )
    except Exception as e:
        print(f"[WARN] Failed to flush task DB state: {e}")
