# hiveos/chunk_runtime.py
from hiveos.db import get_db
import copy
import json
from hiveos.runtime.states import ChunkStatus

class RuntimeChunk:
    def __init__(self, chunk_id, task_id, required_capability, payload=None, depends_on=None, ttl=1, timing="ticks", zone_id=None, conditional=None, session_id=None):
        self.chunk_id = chunk_id
        self.task_id = task_id
        self.required_capability = required_capability
        self._payload = _safe_clone(payload if payload is not None else {})
        self.depends_on = depends_on
        self.ttl = float(ttl)
        self.timing = timing
        self.zone_id = zone_id
        self._conditional = _safe_clone(conditional if conditional is not None else {})
        self.session_id = session_id

        self.status = ChunkStatus.PENDING
        self.assigned_agent = None
        self.retry_count = 0
        self.ticks_completed = 0
        self._immutability_locked = True

    @property
    def payload(self):
        # Return a defensive copy so agents cannot mutate runtime-owned chunk intent.
        return _safe_clone(self._payload)

    @payload.setter
    def payload(self, value):
        if getattr(self, "_immutability_locked", False):
            raise AttributeError("chunk payload is immutable after injection")
        self._payload = _safe_clone(value if value is not None else {})

    @property
    def conditional(self):
        return _safe_clone(self._conditional)

    @conditional.setter
    def conditional(self, value):
        if getattr(self, "_immutability_locked", False):
            raise AttributeError("chunk conditional is immutable after injection")
        self._conditional = _safe_clone(value if value is not None else {})

    def is_ready(self, chunk_objects):
        if not self.depends_on:
            return True
        dep = chunk_objects.get(self.depends_on)
        return dep and dep.status == ChunkStatus.COMPLETED

    def assign_to_agent(self, agent_id):
        if self.status == ChunkStatus.COMPLETED:
            return False
        self.assigned_agent = agent_id
        self.status = ChunkStatus.ASSIGNED
        return True

    def mark_complete(self):
        self.status = ChunkStatus.COMPLETED
        
    def mark_stale(self):
        self.status = ChunkStatus.STALE
        
    def to_dict(self):
        return {
            "chunk_id": self.chunk_id,
            "task_id": self.task_id,
            "required_capability": self.required_capability,
            "status": self.status,
            "assigned_agent": self.assigned_agent,
            "ttl": self.ttl,
            "timing": self.timing,
            "depends_on": self.depends_on,
            "retry_count": self.retry_count,
            "conditional": self.conditional,
            "payload": self.payload,
            "session_id": self.session_id,
        }


def hydrate_from_db_row(row, session_id=None):
    chunk_id, task_id, status, agent, cap, zone, retry, payload, depends_on, ttl, timing, conditional, session_id = row
    
    payload_raw = row[7] or "{}"
    conditional_raw = row[11]
    
    try:
        if isinstance(payload_raw, str):
            payload = json.loads(payload_raw)
        else:
            payload = payload_raw
    except Exception as e:
        print(f"[Warning] Failed to deserialize payload: {e}")
        payload = payload_raw  # Fall back to raw if JSON fails
    
    try:
        conditional = json.loads(conditional_raw) if isinstance(conditional_raw, str) else conditional_raw
    except Exception as e:
        print(f"[Warning] Failed to deserialize conditional dict: {e}")
        conditional = {}
        
    if not isinstance(conditional, dict):
        conditional = {}
        
        
    
    chunk = RuntimeChunk(
        chunk_id=chunk_id,
        task_id=task_id,
        required_capability=cap,
        payload=payload,
        depends_on=depends_on,
        ttl=ttl,
        timing=timing,
        zone_id=zone,
        conditional=conditional,
        session_id=session_id
    )
    chunk.status = status
    chunk.assigned_agent = agent
    chunk.retry_count = retry
    # register(chunk)
    return chunk

def flush_chunk_db_state(chunk_objects):
    try:
        with get_db() as conn:
            for chunk in list(chunk_objects.values()):
                conn.execute(
                    "UPDATE task_chunks SET status = ?, assigned_agent = ?, retry_count = ?, ttl = ?, timing = ? WHERE chunk_id = ?",
                    (
                        chunk.status,
                        chunk.assigned_agent,
                        chunk.retry_count,
                        chunk.ttl,
                        chunk.timing,
                        chunk.chunk_id
                    )
                )
    except Exception as e:
        print(f"[WARN] Failed to flush chunk DB state: {e}")


def _safe_clone(value):
    if value is None:
        return None
    return copy.deepcopy(value)


