import os
import pty
import select
import threading
import time
import termios
import struct
import fcntl
import tempfile
from agent_node.skills.base import BaseSkill
from protos import agent_pb2

class ShellSkill(BaseSkill):
    """Admin Console Skill: Persistent stateful Bash via PTY."""
    def __init__(self, sync_mgr=None):
        self.sync_mgr = sync_mgr
        self.sessions = {} # session_id -> {fd, pid, thread, last_activity, ...}
        self.lock = threading.Lock()
        
        # Phase 3: Prompt Patterns for Edge Intelligence
        self.PROMPT_PATTERNS = [
            r"[\r\n].*[@\w\.\-]+:.*[#$]\s*$",  # bash/zsh: user@host:~$
            r">>>\s*$",                        # python
            r"\.\.\.\s*$",                      # python multi-line
            r">\s*$",                           # node/js
        ]
        
        # --- M7: Idle Session Reaper ---
        # Automatically kills dormant bash processes to free up system resources.
        self.reaper_thread = threading.Thread(target=self._session_reaper, daemon=True, name="ShellReaper")
        self.reaper_thread.start()

    def _session_reaper(self):
        """Background thread that cleans up unused PTY sessions."""
        while True:
            time.sleep(60)
            with self.lock:
                now = time.time()
                for sid, sess in list(self.sessions.items()):
                    # Avoid reaping currently active tasks
                    if sess.get("active_task"):
                        continue
                    
                    # 10 minute idle timeout
                    if now - sess.get("last_activity", 0) > 600:
                        print(f"    [🐚🧹] Reaping idle shell session: {sid}")
                        try:
                            os.close(sess["fd"])
                            os.kill(sess["pid"], 9)
                        except: pass
                        self.sessions.pop(sid, None)

    def _ensure_session(self, session_id, cwd, on_event):
        with self.lock:
            if session_id in self.sessions:
                self.sessions[session_id]["last_activity"] = time.time()
                return self.sessions[session_id]

            print(f"    [🐚] Initializing Persistent Shell Session: {session_id}")
            # Spawn bash in a pty
            pid, fd = pty.fork()
            if pid == 0: # Child
                # Environment prep
                os.environ["TERM"] = "xterm-256color"
                
                # Change to CWD
                if cwd and os.path.exists(cwd):
                    os.chdir(cwd)
                
                # Launch shell
                shell_path = "/bin/bash"
                if not os.path.exists(shell_path):
                    shell_path = "/bin/sh"
                os.execv(shell_path, [shell_path, "--login"])
            
            # Parent
            # Set non-blocking
            fl = fcntl.fcntl(fd, fcntl.F_GETFL)
            fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)

            sess = {
                "fd": fd, 
                "pid": pid, 
                "last_activity": time.time(),
                "buffer_file": None,
                "tail_buffer": "",
                "active_task": None
            }

            def reader():
                while True:
                    try:
                        r, _, _ = select.select([fd], [], [], 0.1)
                        if fd in r:
                            data = os.read(fd, 4096)
                            if not data: break
                            
                            decoded = data.decode("utf-8", errors="replace")
                            
                            # Streaming/Sync logic (Detect completion marker)
                            with self.lock:
                                active_tid = sess.get("active_task")
                                marker = sess.get("marker")
                                if active_tid and marker and sess.get("buffer_file"):
                                    # Phase 2: Persistence Offloading
                                    # Write directly to disk instead of heap memory
                                    sess["buffer_file"].write(decoded)
                                    sess["buffer_file"].flush()
                                    
                                    # Keep a tiny 4KB tail in RAM for marker detection and prompt scanning
                                    sess["tail_buffer"] = (sess.get("tail_buffer", "") + decoded)[-4096:]

                                    if marker in sess["tail_buffer"]:
                                        # Marker found! Extract exit code
                                        try:
                                            # The tail buffer has the marker
                                            after_marker = sess["tail_buffer"].split(marker)[1].strip().split()
                                            exit_code = int(after_marker[0]) if after_marker else 0

                                            # Formulate final stdout summary from the disk file
                                            bf = sess["buffer_file"]
                                            bf.seek(0, 2)
                                            file_len = bf.tell()

                                            HEAD, TAIL = 10_000, 30_000
                                            if file_len > HEAD + TAIL:
                                                bf.seek(0)
                                                head_str = bf.read(HEAD)
                                                bf.seek(file_len - TAIL)
                                                tail_str = bf.read()
                                                omitted = file_len - HEAD - TAIL
                                                pure_stdout = head_str + f"\n\n[... {omitted:,} bytes omitted (full output safely preserved at {bf.name}) ...]\n\n" + tail_str
                                            else:
                                                bf.seek(0)
                                                pure_stdout = bf.read()

                                            # Slice off the marker string and anything after it from the final result
                                            pure_stdout = pure_stdout.split(marker)[0]

                                            sess["result"]["stdout"] = pure_stdout
                                            sess["result"]["status"] = 0 if exit_code == 0 else 1
                                            
                                            # Close the file handle (leaves file on disk)
                                            sess["buffer_file"].close()
                                            sess["buffer_file"] = None
                                            
                                            sess["event"].set()
                                            decoded = pure_stdout.split(marker)[0][-4096:] if marker in pure_stdout else pure_stdout 
                                        except Exception as e:
                                            print(f"    [🐚⚠️] Marker parsing failed: {e}")
                                            sess["event"].set()

                            # Stream terminal output back (with stealth filtering)
                            if on_event:
                                stealth_out = decoded
                                import re
                                # Remove the short source command payload from the PTY echo
                                if " source " in stealth_out and "ctx_" in stealth_out:
                                    stealth_out = re.sub(r'.*source .*ctx_\d+\.sh.*[\r\n]*', '', stealth_out)
                                
                                if "__CORTEX_FIN_SH_" in stealth_out:
                                    # Remove any line that contains our internal marker from the output stream.
                                    # This covers both the fallback echo and the actual execution exit code emission.
                                    stealth_out = re.sub(r'.*__CORTEX_FIN_SH_.*[\r\n]*', '', stealth_out)

                                if stealth_out:
                                    # Phase 3: Client-Side Truncation (Stream Rate Limiting)
                                    # Limit real-time stream to 15KB/sec per session to prevent flooding the Hub over gRPC.
                                    # The full output is still safely written to the tempfile on disk.
                                    with self.lock:
                                        now = time.time()
                                        if now - sess.get("stream_window_start", 0) > 1.0:
                                            sess["stream_window_start"] = now
                                            sess["stream_bytes_sent"] = 0
                                            dropped = sess.get("stream_dropped_bytes", 0)
                                            if dropped > 0:
                                                drop_msg = f"\n[... {dropped:,} bytes truncated from live stream ...]\n"
                                                event = agent_pb2.SkillEvent(
                                                    session_id=session_id, task_id=sess.get("active_task") or "", terminal_out=drop_msg
                                                )
                                                on_event(agent_pb2.ClientTaskMessage(skill_event=event))
                                                sess["stream_dropped_bytes"] = 0

                                        if sess.get("stream_bytes_sent", 0) + len(stealth_out) > 15_000:
                                            sess["stream_dropped_bytes"] = sess.get("stream_dropped_bytes", 0) + len(stealth_out)
                                        else:
                                            sess["stream_bytes_sent"] = sess.get("stream_bytes_sent", 0) + len(stealth_out)
                                            event = agent_pb2.SkillEvent(
                                                session_id=session_id,
                                                task_id=sess.get("active_task") or "",
                                                terminal_out=stealth_out
                                            )
                                            on_event(agent_pb2.ClientTaskMessage(skill_event=event))

                                    # EDGE INTELLIGENCE: Proactively signal prompt detection
                                    # We only check for prompts if we are actively running a task and haven't found the marker yet.
                                    current_event = sess.get("event")
                                    if active_tid and current_event and not current_event.is_set():
                                        import re
                                        tail = sess["tail_buffer"][-100:] if len(sess["tail_buffer"]) > 100 else sess["tail_buffer"]
                                        for pattern in self.PROMPT_PATTERNS:
                                            if re.search(pattern, tail):
                                                # Send specific prompt signal
                                                # Use last 20 chars as the 'prompt' hint
                                                p_hint = tail[-20:].strip()
                                                prompt_event = agent_pb2.SkillEvent(
                                                    session_id=session_id,
                                                    task_id=active_tid,
                                                    prompt=p_hint
                                                )
                                                on_event(agent_pb2.ClientTaskMessage(skill_event=prompt_event))
                                                break
                    except (EOFError, OSError):
                        break
                    except Exception as catch_all:
                        print(f"    [🐚❌] Reader thread FATAL exception: {catch_all}")
                        break

                
                # Thread Cleanup
                print(f"    [🐚] Shell Session Terminated: {session_id}")
                with self.lock:
                    self.sessions.pop(session_id, None)

            t = threading.Thread(target=reader, daemon=True, name=f"ShellReader-{session_id}")
            t.start()
            sess["thread"] = t
            
            self.sessions[session_id] = sess
            return sess


    def handle_transparent_tty(self, task, on_complete, on_event=None):
        """Processes raw TTY/Resize events synchronously (bypasses threadpool/sandbox)."""
        cmd = task.payload_json
        session_id = task.session_id or "default-session"
        try:
            import json
            if cmd.startswith('{') and cmd.endswith('}'):
                raw_payload = json.loads(cmd)
                
                # 1. Raw Keystroke forward
                if isinstance(raw_payload, dict) and "tty" in raw_payload:
                    raw_bytes = raw_payload["tty"]
                    sess = self._ensure_session(session_id, None, on_event)
                    os.write(sess["fd"], raw_bytes.encode("utf-8"))
                    on_complete(task.task_id, {"stdout": "", "status": 0}, task.trace_id)
                    return True
                
                # 2. Window Resize
                if isinstance(raw_payload, dict) and raw_payload.get("action") == "resize":
                    cols = raw_payload.get("cols", 80)
                    rows = raw_payload.get("rows", 24)
                    sess = self._ensure_session(session_id, None, on_event)
                    import termios, struct, fcntl
                    s = struct.pack('HHHH', rows, cols, 0, 0)
                    fcntl.ioctl(sess["fd"], termios.TIOCSWINSZ, s)
                    print(f"    [🐚] Terminal Resized to {cols}x{rows}")
                    on_complete(task.task_id, {"stdout": f"resized to {cols}x{rows}", "status": 0}, task.trace_id)
                    return True
        except Exception as pe:
            print(f"    [🐚] Transparent TTY Fail: {pe}")
        return False

    def execute(self, task, sandbox, on_complete, on_event=None):
        """Dispatches command string to the persistent PTY shell and WAITS for completion."""
        session_id = task.session_id or "default-session"
        tid = task.task_id
        try:
            cmd = task.payload_json
            
            # --- Legacy Full-Command Execution (Sandboxed) ---
            allowed, status_msg = sandbox.verify(cmd)
            if not allowed:
                err_msg = f"\r\n[System] Command blocked: {status_msg}\r\n"
                if on_event:
                    event = agent_pb2.SkillEvent(
                        session_id=session_id, task_id=tid,
                        terminal_out=err_msg
                    )
                    on_event(agent_pb2.ClientTaskMessage(skill_event=event))
                
                return on_complete(tid, {"stderr": f"SANDBOX_VIOLATION: {status_msg}", "status": 2}, task.trace_id)

            # Resolve CWD jail
            cwd = None
            if self.sync_mgr and task.session_id:
                cwd = self.sync_mgr.get_session_dir(task.session_id)
            elif sandbox.policy.get("WORKING_DIR_JAIL"):
                cwd = sandbox.policy["WORKING_DIR_JAIL"]
                if not os.path.exists(cwd):
                    try: os.makedirs(cwd, exist_ok=True)
                    except: pass

            # Handle Session Persistent Process
            sess = self._ensure_session(session_id, cwd, on_event)
            
            # Check for RAW mode first (bypasses busy check for interactive control)
            is_raw = cmd.startswith("!RAW:")
            if is_raw:
                input_str = cmd[5:] + "\n"
                print(f"    [🐚⌨️] RAW Input Injection: {input_str.strip()}")
                os.write(sess["fd"], input_str.encode("utf-8"))
                return on_complete(tid, {"stdout": "INJECTED", "status": 0}, task.trace_id)

            # --- 0. Busy Check: Serialize access to the PTY for standard commands ---
            with self.lock:
                if sess.get("active_task"):
                    curr_tid = sess.get("active_task")
                    return on_complete(tid, {"stderr": f"[BUSY] Session {session_id} is already running task {curr_tid}", "status": 1}, task.trace_id)

            # --- Blocking Wait Logic ---
            # --- Blocking Wait Logic ---
            marker_id = int(time.time())
            marker = f"__CORTEX_FIN_SH_{marker_id}__"
            event = threading.Event()
            result_container = {"stdout": "", "status": 0} # 0 = Success (Protobuf Convention)
            
            # Register waiter in session state
            with self.lock:
                sess["active_task"] = tid
                sess["marker"] = marker
                sess["event"] = event
                # Create a persistent tempfile for stdout instead of RAM buffer
                sess["buffer_file"] = tempfile.NamedTemporaryFile("w+", encoding="utf-8", prefix=f"cortex_task_{tid}_", delete=False)
                sess["tail_buffer"] = ""
                sess["result"] = result_container
                sess["cancel_event"] = threading.Event()

            # Input injection: execute command via sourced script to hide plumbing from ZLE/Readline echo
            try:
                script_path = os.path.join(tempfile.gettempdir(), f"ctx_{marker_id}.sh")
                script_content = f"{cmd}\n__ctx_exit=$?\necho \"__CORTEX_FIN_SH_{marker_id}__\" $__ctx_exit\nrm -f {script_path}\n"
                try:
                    with open(script_path, "w") as f:
                        f.write(script_content)
                    # Space prefix often hides it from shell history. Extremely short payload prevents ZLE line-wrap scrambling!
                    full_input = f" source {script_path}\n"
                except Exception as e:
                    print(f"    [🐚⚠️] Failed to write injection script dictating fallback: {e}")
                    # Fallback if writing file natively fails for some system reason
                    full_input = f"({cmd}) ; echo \"__CORTEX_FIN_SH_\"\"{marker_id}__\" $?\n"
                    
                os.write(sess["fd"], full_input.encode("utf-8"))

                # Wait for completion (triggered by reader) OR cancellation
                timeout = (task.timeout_ms / 1000.0) if task.timeout_ms > 0 else 60.0
                start_time = time.time()
                while time.time() - start_time < timeout:
                    # Check for completion (reader found marker)
                    if event.is_set():
                        return on_complete(tid, result_container, task.trace_id)
                    
                    # Check for cancellation (HUB sent cancel)
                    if sess["cancel_event"].is_set():
                        print(f"    [🐚🛑] Task {tid} cancelled on node.")
                        return on_complete(tid, {"stderr": "ABORTED", "status": 2}, task.trace_id)
                    
                    # Sleep slightly to avoid busy loop
                    time.sleep(0.1)

                # Timeout Case
                print(f"    [🐚⚠️] Task {tid} timed out on node.")
                with self.lock:
                    if sess.get("buffer_file"):
                        try:
                            sess["buffer_file"].seek(0, 2)
                            file_len = sess["buffer_file"].tell()
                            HEAD, TAIL = 10_000, 30_000
                            if file_len > HEAD + TAIL:
                                sess["buffer_file"].seek(0)
                                head_str = sess["buffer_file"].read(HEAD)
                                sess["buffer_file"].seek(file_len - TAIL)
                                tail_str = sess["buffer_file"].read()
                                omitted = file_len - HEAD - TAIL
                                partial_out = head_str + f"\n\n[... {omitted:,} bytes omitted (full timeout output saved to {sess['buffer_file'].name}) ...]\n\n" + tail_str
                            else:
                                sess["buffer_file"].seek(0)
                                partial_out = sess["buffer_file"].read()
                        except:
                            partial_out = ""
                    else:
                        partial_out = ""
                
                on_complete(tid, {"stdout": partial_out, "stderr": "TIMEOUT", "status": 2}, task.trace_id)
                
            finally:
                # Cleanup session task state
                with self.lock:
                    if sess.get("active_task") == tid:
                        if sess.get("buffer_file"):
                            try:
                                sess["buffer_file"].close()
                            except: pass
                            sess["buffer_file"] = None
                        sess["active_task"] = None
                        sess["marker"] = None
                        sess["event"] = None
                        sess["result"] = None
                        sess["cancel_event"] = None

        except Exception as e:
            print(f"    [🐚❌] Execute Error for {tid}: {e}")
            on_complete(tid, {"stderr": str(e), "status": 2}, task.trace_id)

    def cancel(self, task_id: str):
        """Cancels an active task — for persistent shell, this sends a SIGINT (Ctrl+C)."""
        with self.lock:
            for sid, sess in self.sessions.items():
                if sess.get("active_task") == task_id:
                    print(f"[🛑] Sending SIGINT (Ctrl+C) to shell session (Task {task_id}): {sid}")
                    # Write \x03 (Ctrl+C) to the master FD
                    os.write(sess["fd"], b"\x03")
                    # Break the wait loop in execute thread
                    if sess.get("cancel_event"):
                        sess["cancel_event"].set()
        return True


    def shutdown(self):
        """Cleanup: Terminates all persistent shells."""
        with self.lock:
            for sid, sess in list(self.sessions.items()):
                print(f"[🛑] Cleaning up persistent shell: {sid}")
                try: os.close(sess["fd"])
                except: pass
                # kill pid
                try: os.kill(sess["pid"], 9)
                except: pass
            self.sessions.clear()
