Atomizer/atomizer-dashboard/backend/api/routes/terminal.py

"""
Terminal WebSocket for Claude Code CLI

Provides a PTY-based terminal that runs Claude Code in the dashboard.
Uses pywinpty on Windows for proper interactive terminal support.
"""

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from typing import Optional
import asyncio
import subprocess
import sys
import os
import signal
import json

router = APIRouter()

# Store active terminal sessions
_terminal_sessions: dict = {}

# Path to Atomizer root (for loading prompts)
ATOMIZER_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
)))


def get_session_prompt(study_name: str = None) -> str:
    """
    Generate the initial prompt for a Claude session.

    This injects the Protocol Operating System context and study-specific info.
    """
    prompt_lines = [
        "# Atomizer Session Context",
        "",
        "You are assisting with **Atomizer** - an LLM-first FEA optimization framework.",
        "",
        "## Bootstrap (READ FIRST)",
        "",
        "Read these files to understand how to help:",
        "- `.claude/skills/00_BOOTSTRAP.md` - Task classification and routing",
        "- `.claude/skills/01_CHEATSHEET.md` - Quick reference (I want X → Use Y)",
        "- `.claude/skills/02_CONTEXT_LOADER.md` - What to load per task",
        "",
        "## Protocol System",
        "",
        "| Layer | Location | Purpose |",
        "|-------|----------|---------|",
        "| Operations | `docs/protocols/operations/OP_*.md` | How-to guides |",
        "| System | `docs/protocols/system/SYS_*.md` | Core specs |",
        "| Extensions | `docs/protocols/extensions/EXT_*.md` | Adding features |",
        "",
    ]

    if study_name:
        prompt_lines.extend([
            f"## Current Study: `{study_name}`",
            "",
            f"**Directory**: `studies/{study_name}/`",
            "",
            "Key files:",
            f"- `studies/{study_name}/1_setup/optimization_config.json` - Configuration",
            f"- `studies/{study_name}/2_results/study.db` - Optuna database",
            f"- `studies/{study_name}/README.md` - Study documentation",
            "",
            "Quick status check:",
            "```bash",
            f"python -c \"import optuna; s=optuna.load_study('{study_name}', 'sqlite:///studies/{study_name}/2_results/study.db'); print(f'Trials: {{len(s.trials)}}, Best: {{s.best_value}}')\"",
            "```",
            "",
        ])
    else:
        prompt_lines.extend([
            "## No Study Selected",
            "",
            "No specific study context. You can:",
            "- List studies: `ls studies/`",
            "- Create new study: Ask user what they want to optimize",
            "- Load context: Read `.claude/skills/core/study-creation-core.md`",
            "",
        ])

    prompt_lines.extend([
        "## Key Principles",
        "",
        "1. **Read bootstrap first** - Follow task routing from 00_BOOTSTRAP.md",
        "2. **Use centralized extractors** - Check `optimization_engine/extractors/`",
        "3. **Never modify master models** - Work on copies",
        "4. **Python env**: Always use `conda activate atomizer`",
        "",
        "---",
        "*Session launched from Atomizer Dashboard*",
    ])

    return "\n".join(prompt_lines)

# Check if winpty is available (for Windows)
try:
    from winpty import PtyProcess
    HAS_WINPTY = True
except ImportError:
    HAS_WINPTY = False


class TerminalSession:
    """Manages a Claude Code terminal session."""

    def __init__(self, session_id: str, working_dir: str):
        self.session_id = session_id
        self.working_dir = working_dir
        self.process = None
        self.websocket: Optional[WebSocket] = None
        self._read_task: Optional[asyncio.Task] = None
        self._running = False
        self._use_winpty = sys.platform == "win32" and HAS_WINPTY

    async def start(self, websocket: WebSocket):
        """Start the Claude Code process."""
        self.websocket = websocket
        self._running = True

        try:
            if self._use_winpty:
                # Use winpty for proper PTY on Windows
                # Spawn claude directly - winpty handles the interactive terminal
                import shutil
                claude_path = shutil.which("claude") or "claude"

                # Ensure HOME and USERPROFILE are properly set for Claude CLI auth
                env = {**os.environ}
                env["FORCE_COLOR"] = "1"
                env["TERM"] = "xterm-256color"
                # Claude CLI looks for credentials in HOME/.claude or USERPROFILE/.claude
                # Ensure these are set correctly
                if "USERPROFILE" in env and "HOME" not in env:
                    env["HOME"] = env["USERPROFILE"]

                self.process = PtyProcess.spawn(
                    claude_path,
                    cwd=self.working_dir,
                    env=env
                )
            elif sys.platform == "win32":
                # Fallback: Windows without winpty - use subprocess
                # Run claude with --dangerously-skip-permissions for non-interactive mode
                self.process = subprocess.Popen(
                    ["cmd.exe", "/k", claude_cmd],
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    cwd=self.working_dir,
                    bufsize=0,
                    creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
                    env={**os.environ, "FORCE_COLOR": "1", "TERM": "xterm-256color"}
                )
            else:
                # On Unix, use pty
                import pty
                master_fd, slave_fd = pty.openpty()
                self.process = subprocess.Popen(
                    [claude_cmd],
                    stdin=slave_fd,
                    stdout=slave_fd,
                    stderr=slave_fd,
                    cwd=self.working_dir,
                    env={**os.environ, "TERM": "xterm-256color"}
                )
                os.close(slave_fd)
                self._master_fd = master_fd

            # Start reading output
            self._read_task = asyncio.create_task(self._read_output())

            await self.websocket.send_json({
                "type": "started",
                "message": f"Claude Code started in {self.working_dir}"
            })

        except FileNotFoundError:
            await self.websocket.send_json({
                "type": "error",
                "message": "Claude Code CLI not found. Please install it with: npm install -g @anthropic-ai/claude-code"
            })
            self._running = False
        except Exception as e:
            await self.websocket.send_json({
                "type": "error",
                "message": f"Failed to start Claude Code: {str(e)}"
            })
            self._running = False

    async def _read_output(self):
        """Read output from the process and send to WebSocket."""
        try:
            while self._running:
                if self._use_winpty:
                    # Read from winpty
                    if self.process and self.process.isalive():
                        loop = asyncio.get_event_loop()
                        try:
                            data = await loop.run_in_executor(
                                None,
                                lambda: self.process.read(4096)
                            )
                            if data:
                                await self.websocket.send_json({
                                    "type": "output",
                                    "data": data
                                })
                        except Exception:
                            break
                    else:
                        break
                elif sys.platform == "win32":
                    # Windows subprocess pipe mode
                    if self.process and self.process.poll() is None:
                        if self.process.stdout:
                            loop = asyncio.get_event_loop()
                            try:
                                # Use non-blocking read with a timeout
                                import msvcrt
                                import ctypes

                                # Read available data
                                data = await loop.run_in_executor(
                                    None,
                                    lambda: self.process.stdout.read(1)
                                )
                                if data:
                                    # Read more if available
                                    more_data = b""
                                    try:
                                        # Try to read more without blocking
                                        while True:
                                            extra = self.process.stdout.read(1)
                                            if extra:
                                                more_data += extra
                                            else:
                                                break
                                    except:
                                        pass

                                    full_data = data + more_data
                                    await self.websocket.send_json({
                                        "type": "output",
                                        "data": full_data.decode("utf-8", errors="replace")
                                    })
                            except Exception:
                                break
                    else:
                        break
                else:
                    # Unix PTY
                    loop = asyncio.get_event_loop()
                    try:
                        data = await loop.run_in_executor(
                            None,
                            lambda: os.read(self._master_fd, 4096)
                        )
                        if data:
                            await self.websocket.send_json({
                                "type": "output",
                                "data": data.decode("utf-8", errors="replace")
                            })
                    except OSError:
                        break

                await asyncio.sleep(0.01)

            # Process ended
            if self.websocket:
                exit_code = -1
                if self._use_winpty:
                    exit_code = self.process.exitstatus if self.process else -1
                elif self.process:
                    exit_code = self.process.poll() if self.process.poll() is not None else -1

                await self.websocket.send_json({
                    "type": "exit",
                    "code": exit_code
                })

        except Exception as e:
            if self.websocket:
                try:
                    await self.websocket.send_json({
                        "type": "error",
                        "message": str(e)
                    })
                except:
                    pass

    async def write(self, data: str):
        """Write input to the process."""
        if not self.process or not self._running:
            return

        try:
            if self._use_winpty:
                self.process.write(data)
            elif sys.platform == "win32":
                if self.process.stdin:
                    self.process.stdin.write(data.encode())
                    self.process.stdin.flush()
            else:
                os.write(self._master_fd, data.encode())
        except Exception as e:
            if self.websocket:
                await self.websocket.send_json({
                    "type": "error",
                    "message": f"Write error: {str(e)}"
                })

    async def resize(self, cols: int, rows: int):
        """Resize the terminal."""
        if self._use_winpty and self.process:
            try:
                self.process.setwinsize(rows, cols)
            except:
                pass
        elif sys.platform != "win32" and hasattr(self, '_master_fd'):
            import struct
            import fcntl
            import termios
            winsize = struct.pack("HHHH", rows, cols, 0, 0)
            fcntl.ioctl(self._master_fd, termios.TIOCSWINSZ, winsize)

    async def stop(self):
        """Stop the terminal session."""
        self._running = False

        if self._read_task:
            self._read_task.cancel()
            try:
                await self._read_task
            except asyncio.CancelledError:
                pass

        if self.process:
            try:
                if self._use_winpty:
                    self.process.terminate()
                elif sys.platform == "win32":
                    self.process.terminate()
                else:
                    os.kill(self.process.pid, signal.SIGTERM)
                    self.process.wait(timeout=2)
            except:
                try:
                    if hasattr(self.process, 'kill'):
                        self.process.kill()
                except:
                    pass

        if sys.platform != "win32" and hasattr(self, '_master_fd'):
            try:
                os.close(self._master_fd)
            except:
                pass


@router.websocket("/claude")
async def claude_terminal(websocket: WebSocket, working_dir: str = None, study_id: str = None):
    """
    WebSocket endpoint for Claude Code terminal.

    Query params:
        working_dir: Directory to start Claude Code in (defaults to Atomizer root)
        study_id: Optional study ID to set context for Claude

    Client -> Server messages:
        {"type": "input", "data": "user input text"}
        {"type": "resize", "cols": 80, "rows": 24}
        {"type": "stop"}

    Server -> Client messages:
        {"type": "started", "message": "..."}
        {"type": "output", "data": "terminal output"}
        {"type": "exit", "code": 0}
        {"type": "error", "message": "..."}
        {"type": "context", "prompt": "..."}  # Initial context prompt
    """
    await websocket.accept()

    # Default to Atomizer root directory
    if not working_dir:
        working_dir = ATOMIZER_ROOT

    # Create session
    session_id = f"claude-{id(websocket)}"
    session = TerminalSession(session_id, working_dir)
    _terminal_sessions[session_id] = session

    try:
        # Send context prompt to frontend (for display/reference)
        context_prompt = get_session_prompt(study_id)
        await websocket.send_json({
            "type": "context",
            "prompt": context_prompt,
            "study_id": study_id
        })

        # Start Claude Code
        await session.start(websocket)

        # If study_id provided, send initial context to Claude after startup
        if study_id:
            # Wait a moment for Claude to initialize
            await asyncio.sleep(1.0)
            # Send the context as the first message
            initial_message = f"I'm working with the Atomizer study '{study_id}'. Please read .claude/skills/00_BOOTSTRAP.md first to understand the Protocol Operating System, then help me with this study.\n"
            await session.write(initial_message)

        # Handle incoming messages
        while session._running:
            try:
                message = await websocket.receive_json()

                if message.get("type") == "input":
                    await session.write(message.get("data", ""))
                elif message.get("type") == "resize":
                    await session.resize(
                        message.get("cols", 80),
                        message.get("rows", 24)
                    )
                elif message.get("type") == "stop":
                    break

            except WebSocketDisconnect:
                break
            except Exception as e:
                await websocket.send_json({
                    "type": "error",
                    "message": str(e)
                })

    finally:
        await session.stop()
        _terminal_sessions.pop(session_id, None)


@router.get("/status")
async def terminal_status():
    """Check if Claude Code CLI is available."""
    import shutil

    claude_path = shutil.which("claude")

    return {
        "available": claude_path is not None,
        "path": claude_path,
        "winpty_available": HAS_WINPTY,
        "message": "Claude Code CLI is available" if claude_path else "Claude Code CLI not found. Install with: npm install -g @anthropic-ai/claude-code"
    }


@router.get("/context")
async def get_context(study_id: str = None):
    """
    Get the context prompt for a Claude session without starting a terminal.

    Useful for displaying context in the UI or preparing prompts.

    Query params:
        study_id: Optional study ID to include study-specific context
    """
    prompt = get_session_prompt(study_id)

    return {
        "study_id": study_id,
        "prompt": prompt,
        "bootstrap_files": [
            ".claude/skills/00_BOOTSTRAP.md",
            ".claude/skills/01_CHEATSHEET.md",
            ".claude/skills/02_CONTEXT_LOADER.md",
        ],
        "study_files": [
            f"studies/{study_id}/1_setup/optimization_config.json",
            f"studies/{study_id}/2_results/study.db",
            f"studies/{study_id}/README.md",
        ] if study_id else []
    }