import threading
import time
from typing import Optional
from app.config import settings

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):
        # ---- Memory pressure limits from config ----
        self.STREAM_HEAD_CHARS  = settings.STREAM_HEAD_CHARS
        self.STREAM_TAIL_CHARS  = settings.STREAM_TAIL_CHARS
        self.STREAM_MAX_CHARS   = self.STREAM_HEAD_CHARS + self.STREAM_TAIL_CHARS

        self.THOUGHT_HEAD_COUNT = settings.THOUGHT_HEAD_COUNT
        self.THOUGHT_TAIL_COUNT = settings.THOUGHT_TAIL_COUNT
        self.THOUGHT_MAX_COUNT  = self.THOUGHT_HEAD_COUNT + self.THOUGHT_TAIL_COUNT

        # Sharded locking for high concurrency
        self.NUM_SHARDS = 16
        self.shards = [{"tasks": {}, "lock": threading.Lock()} for _ in range(self.NUM_SHARDS)]

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

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

    def _get_shard(self, task_id: str):
        """Returns the shard (dict and lock) for a given task_id."""
        return self.shards[hash(task_id) % self.NUM_SHARDS]

    def _trim_stream(self, buf: str, chunk: str) -> str:
        """Append chunk to buf, enforcing the head+tail memory limit."""
        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."""
        thoughts = list(thoughts)
        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()
        shard = self._get_shard(task_id)
        with shard["lock"]:
            shard["tasks"][task_id] = {
                "event":         event,
                "prompt_event":  prompt_event,
                "result":        None,
                "node_id":       node_id,
                "stream_buffer": "",
                "thought_history": [],
                "created_at":    time.time(),
                "completed_at":  None,
            }
        return event

    def get_task_event(self, task_id: str) -> Optional[threading.Event]:
        """Returns the completion event for an existing task, or None if not found."""
        shard = self._get_shard(task_id)
        with shard["lock"]:
            task = shard["tasks"].get(task_id)
            return task["event"] if task else None

    def add_thought(self, task_id: str, thought: str) -> bool:
        """Adds an AI reasoning entry to the task's history (head+tail bounded)."""
        shard = self._get_shard(task_id)
        with shard["lock"]:
            if task_id not in shard["tasks"]:
                return False
            entry = {"time": time.time(), "thought": thought}
            shard["tasks"][task_id]["thought_history"] = self._trim_thoughts(
                shard["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)."""
        shard = self._get_shard(task_id)
        with shard["lock"]:
            if task_id not in shard["tasks"]:
                return False
            shard["tasks"][task_id]["stream_buffer"] = self._trim_stream(
                shard["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."""
        shard = self._get_shard(task_id)
        with shard["lock"]:
            if task_id in shard["tasks"]:
                shard["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."""
        shard = self._get_shard(task_id)
        with shard["lock"]:
            if task_id not in shard["tasks"]:
                return False
            if isinstance(result, dict) and result.get("stdout"):
                result["stdout"] = self._trim_final_stdout(result["stdout"])
            shard["tasks"][task_id]["result"]       = result
            shard["tasks"][task_id]["completed_at"] = time.time()
            shard["tasks"][task_id]["event"].set()
            shard["tasks"][task_id]["prompt_event"].set()
            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."""
        total_failed = 0
        for shard in self.shards:
            with shard["lock"]:
                to_fail = [
                    tid for tid, t in shard["tasks"].items()
                    if t.get("node_id") == node_id and t.get("result") is None
                ]
                for tid in to_fail:
                    shard["tasks"][tid]["result"]       = {"error": error_msg, "status": "ERROR"}
                    shard["tasks"][tid]["completed_at"] = time.time()
                    shard["tasks"][tid]["event"].set()
                    shard["tasks"][tid]["prompt_event"].set()
                total_failed += len(to_fail)
        return total_failed

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

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

            if res is None:
                return {
                    "stdout":         data["stream_buffer"],
                    "stderr":         "",
                    "status":         "RUNNING",
                    "partial":        True,
                    "thought_history": history,
                }

            if isinstance(res, dict):
                res = dict(res)
                res["thought_history"] = history
            return res

    def pop(self, task_id: str):
        """Removes the task's state from the journal."""
        shard = self._get_shard(task_id)
        with shard["lock"]:
            return shard["tasks"].pop(task_id, None)

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