import threading
import time


class TaskJournal:
    """
    State machine for tracking tasks through their asynchronous lifecycle.

    Memory Pressure Strategy — Head + Tail buffers:
    - For stream_buffer (stdout): keep first HEAD bytes (command echo + initial
      context, which is the most important for understanding what was requested)
      and last TAIL bytes (most recent output). The middle is replaced with a
      single '[... N bytes omitted ...]' marker.
    - For thought_history: same idea — keep first HEAD_COUNT entries (initial
      reasoning context) and last TAIL_COUNT entries (most recent AI decisions).
    """

    def __init__(self):
        self.lock = threading.Lock()
        self.tasks = {}  # task_id -> { event, result, node_id, buffers... }

        # ---- Memory pressure limits ----
        # stream_buffer hard limit = HEAD + TAIL = 40KB
        self.STREAM_HEAD_CHARS  = 10_000   # ~10KB — preserve initial command/echo
        self.STREAM_TAIL_CHARS  = 30_000   # ~30KB — preserve most recent output
        self.STREAM_MAX_CHARS   = self.STREAM_HEAD_CHARS + self.STREAM_TAIL_CHARS

        # thought_history hard limit = HEAD + TAIL = 20 entries
        self.THOUGHT_HEAD_COUNT = 5        # first 5 thoughts (task inception context)
        self.THOUGHT_TAIL_COUNT = 15       # last 15 thoughts (most recent AI reasoning)
        self.THOUGHT_MAX_COUNT  = self.THOUGHT_HEAD_COUNT + self.THOUGHT_TAIL_COUNT

        threading.Thread(
            target=self._cleanup_loop, daemon=True, name="JournalCleanup"
        ).start()

    # ------------------------------------------------------------------
    #  Internal helpers
    # ------------------------------------------------------------------

    def _trim_stream(self, buf: str, chunk: str) -> str:
        """
        Append chunk to buf, enforcing the head+tail memory limit.

        When the total length exceeds STREAM_MAX_CHARS, the middle section is
        replaced with one human-readable '[... N bytes omitted ...]' marker.
        The head (STREAM_HEAD_CHARS) and tail (STREAM_TAIL_CHARS) are always kept.
        """
        combined = buf + chunk
        if len(combined) <= self.STREAM_MAX_CHARS:
            return combined

        head    = combined[:self.STREAM_HEAD_CHARS]
        tail    = combined[-self.STREAM_TAIL_CHARS:]
        omitted = len(combined) - self.STREAM_HEAD_CHARS - self.STREAM_TAIL_CHARS
        marker  = f"\n\n[... {omitted:,} bytes omitted — memory pressure limit reached ...]\n\n"
        return head + marker + tail

    def _trim_thoughts(self, thoughts: list, new_entry: dict) -> list:
        """
        Append new_entry to thoughts, enforcing the head+tail thought limit.

        When total entries exceed THOUGHT_MAX_COUNT, the middle entries are
        collapsed into a single sentinel entry.
        """
        thoughts = list(thoughts)   # shallow copy to avoid mutating in-place
        thoughts.append(new_entry)
        if len(thoughts) <= self.THOUGHT_MAX_COUNT:
            return thoughts

        head    = thoughts[:self.THOUGHT_HEAD_COUNT]
        tail    = thoughts[-self.THOUGHT_TAIL_COUNT:]
        omitted = len(thoughts) - self.THOUGHT_HEAD_COUNT - self.THOUGHT_TAIL_COUNT
        sentinel = {
            "time":    time.time(),
            "thought": f"[... {omitted} intermediate thoughts omitted — memory limit reached ...]"
        }
        return head + [sentinel] + tail

    def _trim_final_stdout(self, stdout: str) -> str:
        """Apply head+tail trim to a final stdout blob received from the node."""
        if len(stdout) <= self.STREAM_MAX_CHARS:
            return stdout
        head    = stdout[:self.STREAM_HEAD_CHARS]
        tail    = stdout[-self.STREAM_TAIL_CHARS:]
        omitted = len(stdout) - self.STREAM_HEAD_CHARS - self.STREAM_TAIL_CHARS
        marker  = f"\n\n[... {omitted:,} bytes omitted — output too large ...]\n\n"
        return head + marker + tail

    # ------------------------------------------------------------------
    #  Lifecycle
    # ------------------------------------------------------------------

    def _cleanup_loop(self):
        while True:
            time.sleep(300)     # every 5 minutes
            try:
                self.cleanup()
            except Exception:
                pass

    def register(self, task_id: str, node_id: str = None) -> threading.Event:
        """Initializes state for a new task and returns its completion event."""
        event        = threading.Event()
        prompt_event = threading.Event()
        with self.lock:
            self.tasks[task_id] = {
                "event":         event,
                "prompt_event":  prompt_event,
                "result":        None,
                "node_id":       node_id,
                "stream_buffer": "",          # head+tail bounded raw stdout
                "thought_history": [],        # head+tail bounded AI reasoning log
                "created_at":    time.time(),
                "completed_at":  None,
            }
        return event

    def add_thought(self, task_id: str, thought: str) -> bool:
        """Adds an AI reasoning entry to the task's history (head+tail bounded)."""
        with self.lock:
            if task_id not in self.tasks:
                return False
            entry = {"time": time.time(), "thought": thought}
            self.tasks[task_id]["thought_history"] = self._trim_thoughts(
                self.tasks[task_id]["thought_history"], entry
            )
            return True

    def append_stream(self, task_id: str, chunk: str) -> bool:
        """Appends a real-time output chunk to the task's stream buffer (head+tail bounded)."""
        with self.lock:
            if task_id not in self.tasks:
                return False
            self.tasks[task_id]["stream_buffer"] = self._trim_stream(
                self.tasks[task_id]["stream_buffer"], chunk
            )
            return True

    def signal_prompt(self, task_id: str) -> bool:
        """Signals that an agent node detected an interactive prompt."""
        with self.lock:
            if task_id in self.tasks:
                self.tasks[task_id]["prompt_event"].set()
                return True
        return False

    def fulfill(self, task_id: str, result: dict) -> bool:
        """
        Stores the final result from a node and wakes all waiting threads.
        Also trims the final stdout blob if it exceeds the memory limit.
        """
        with self.lock:
            if task_id not in self.tasks:
                return False
            # Trim oversized final stdout before storing
            if isinstance(result, dict) and result.get("stdout"):
                result["stdout"] = self._trim_final_stdout(result["stdout"])
            self.tasks[task_id]["result"]       = result
            self.tasks[task_id]["completed_at"] = time.time()
            self.tasks[task_id]["event"].set()
            self.tasks[task_id]["prompt_event"].set()   # wake sub-agent edge sleeps
            return True

    def fail_node_tasks(self, node_id: str, error_msg: str = "Node disconnected") -> int:
        """Fulfills all pending tasks for a disconnected node with an error."""
        with self.lock:
            to_fail = [
                tid for tid, t in self.tasks.items()
                if t.get("node_id") == node_id and t.get("result") is None
            ]
            for tid in to_fail:
                self.tasks[tid]["result"]       = {"error": error_msg, "status": "ERROR"}
                self.tasks[tid]["completed_at"] = time.time()
                self.tasks[tid]["event"].set()
                self.tasks[tid]["prompt_event"].set()
        return len(to_fail)

    def get_result(self, task_id: str):
        """Returns the enriched result for a task, or a partial snapshot if still running."""
        with self.lock:
            data = self.tasks.get(task_id)
            if data is None:
                return None

            res     = data["result"]
            history = data.get("thought_history", [])

            if res is None:
                # Task still running — return partial stream buffer
                return {
                    "stdout":         data["stream_buffer"],
                    "stderr":         "",
                    "status":         "RUNNING",
                    "partial":        True,
                    "thought_history": history,
                }

            # Enrich final result with reasoning history
            if isinstance(res, dict):
                res = dict(res)             # don't mutate stored result
                res["thought_history"] = history
            return res

    def pop(self, task_id: str):
        """Removes the task's state from the journal (call after result is consumed)."""
        with self.lock:
            return self.tasks.pop(task_id, None)

    def cleanup(self, max_age_s: int = 900):
        """Purges stale tasks to prevent slow memory accumulation."""
        now = time.time()
        with self.lock:
            to_remove = [
                tid for tid, t in self.tasks.items()
                if (t["completed_at"] and (now - t["completed_at"]) > 120)  # finished: keep 2m
                or (now - t["created_at"]) > 900                            # pending:  keep 15m
            ]
            for tid in to_remove:
                del self.tasks[tid]
