"""Claude SDK 持久会话 — per-session 架构

照搬 claude-mem 的并发模型：每个 Claude Code 窗口(session_id)对应一个独立的
持久 claude 进程，通过 stream-json NDJSON 协议双向通信。

数据流:
  首次 observe → spawn claude 进程(~5s, --setting-sources 跳过 hooks) → 进程常驻
  后续 observe → 复用进程(~3-5s API 延迟)
  session 空闲 → cleanup_idle() 清理进程

协议格式(实测确认):
  stdin message: {"type":"user","message":{"role":"user","content":[{"type":"text","text":"..."}]}}
  stdout init: {"type":"system","subtype":"init",...}  — 首条 stdin 消息后才发送
  stdout response: {"type":"assistant","message":{"content":[{"type":"text","text":"..."}],...}}

注意:
  --output-format stream-json 在 -p 模式下必须配合 --verbose
  system:init 需要 stdin 触发（进程不会自动发送）
  stop_reason 始终为 null，每条 assistant 消息就是一个完整回合
  --verbose 模式会 echo user 消息和 hook 事件到 stdout
"""

from __future__ import annotations

import json
import logging
import os
import select
import shutil
import signal
import subprocess
import threading
import time
from pathlib import Path
from typing import TYPE_CHECKING

from ...env_utils import make_clean_env

if TYPE_CHECKING:
    from ...config import Config

logger = logging.getLogger(__name__)

_INDEX1_HOME = Path(os.environ["INDEX1_HOME"]) if os.environ.get("INDEX1_HOME") else Path.home() / ".claude-index1"
_SESSION_LOG = _INDEX1_HOME / "claude_session.log"
_CHILD_PIDS_FILE = _INDEX1_HOME / ".observer_child_pids"

# stream-json 协议中标识 LLM 响应完成的消息类型
_ASSISTANT_TYPE = "assistant"
_SYSTEM_INIT_SUBTYPE = "init"

# 禁用的工具列表 — Observer 只需要 LLM 推理，不需要任何工具
_DISALLOWED_TOOLS = (
    "Bash,Read,Write,Edit,Grep,Glob,WebFetch,WebSearch,"
    "Task,NotebookEdit,AskUserQuestion,TodoWrite"
)


def _add_child_pid(pid: int) -> None:
    """追踪 claude 子进程 PID（异常退出后可清理孤儿）"""
    try:
        _INDEX1_HOME.mkdir(parents=True, exist_ok=True)
        pids = _read_child_pids()
        pids.add(pid)
        _CHILD_PIDS_FILE.write_text("\n".join(str(p) for p in sorted(pids)) + "\n")
    except OSError:
        pass


def _remove_child_pid(pid: int) -> None:
    """移除已停止的子进程 PID"""
    try:
        pids = _read_child_pids()
        pids.discard(pid)
        if pids:
            _CHILD_PIDS_FILE.write_text("\n".join(str(p) for p in sorted(pids)) + "\n")
        else:
            _CHILD_PIDS_FILE.unlink(missing_ok=True)
    except OSError:
        pass


def _read_child_pids() -> set[int]:
    """读取追踪文件中的 PID 集合"""
    try:
        if _CHILD_PIDS_FILE.exists():
            text = _CHILD_PIDS_FILE.read_text().strip()
            return {int(line) for line in text.splitlines() if line.strip().isdigit()}
    except (OSError, ValueError):
        pass
    return set()


def _cleanup_orphan_processes() -> int:
    """清理上一轮 Worker 留下的孤儿 claude 进程。

    Worker 启动时调用。读取 PID 追踪文件，kill 仍存活的进程。

    三层防御:
    - L1: kill 前验证进程名含 "claude"（防 PID 复用误杀）
    - L2: SIGTERM → 2s → SIGKILL（防 SIGTERM 被忽略的僵尸进程）
    - L3: 日志记录每个 PID 的处理结果

    Returns:
        清理的孤儿进程数
    """
    pids = _read_child_pids()
    if not pids:
        return 0

    killed = 0
    for pid in pids:
        try:
            os.kill(pid, 0)  # 检查存活
        except (OSError, ProcessLookupError):
            continue  # 进程已不存在

        # L1 事前校验：验证进程名含 "claude"，防止 PID 复用后误杀无关进程
        try:
            result = subprocess.run(
                ["ps", "-p", str(pid), "-o", "args="],
                capture_output=True, text=True, timeout=2.0,
            )
            proc_args = result.stdout.strip()
            if "claude" not in proc_args.lower():
                logger.warning(
                    "孤儿 PID %d 不是 claude 进程 (args=%s)，跳过",
                    pid, proc_args[:100],
                )
                continue
        except (OSError, subprocess.TimeoutExpired):
            logger.debug("无法验证 PID %d 的进程名，保守跳过", pid)
            continue

        # L2 事中监控：SIGTERM → 2s → SIGKILL
        try:
            os.kill(pid, signal.SIGTERM)
        except (OSError, ProcessLookupError):
            continue

        # 等待进程退出（最多 2s）
        _dead = False
        for _ in range(4):
            time.sleep(0.5)
            try:
                os.kill(pid, 0)
            except (OSError, ProcessLookupError):
                _dead = True
                break

        if not _dead:
            try:
                os.kill(pid, signal.SIGKILL)
                logger.warning("孤儿 PID %d SIGTERM 无效，已 SIGKILL", pid)
            except (OSError, ProcessLookupError):
                pass

        killed += 1
        logger.info("清理孤儿 claude 进程: pid=%d", pid)

    # 清空追踪文件
    try:
        _CHILD_PIDS_FILE.unlink(missing_ok=True)
    except OSError:
        pass

    if killed:
        logger.info("共清理 %d 个孤儿 claude 进程", killed)
    return killed


class ClaudeSession:
    """单个持久 Claude CLI 会话 — stream-json NDJSON 协议

    每个 Claude Code 窗口 session 对应一个 ClaudeSession 实例，
    内部维持一个 claude 子进程，通过 stdin/stdout 双向通信。

    线程安全：同 session 内通过 _lock 序列化请求。
    """

    def __init__(self, config: Config, session_id: str):
        self.session_id = session_id
        self._config = config
        self._proc: subprocess.Popen | None = None
        self._stderr_file = None  # open file handle, closed in stop()
        self._lock = threading.Lock()
        self._last_activity = time.monotonic()
        self._ready = False

    def _diag(self, msg: str) -> None:
        """写入诊断日志到 session log 文件（排障用）。

        直接写文件而非 logger —— hook 子进程可能没配置 logging handler。
        """
        try:
            if self._stderr_file and not self._stderr_file.closed:
                ts = time.strftime("%H:%M:%S")
                self._stderr_file.write(f"[{ts}] [{self.session_id}] {msg}\n")
                self._stderr_file.flush()
        except OSError:
            pass

    def start(self) -> bool:
        """启动 claude 持久进程，等待 init 消息确认就绪。

        Returns:
            True — 进程就绪，可以发送消息
            False — 启动失败
        """
        claude_bin = shutil.which("claude")
        if not claude_bin:
            logger.warning("claude CLI 不在 PATH 中，无法启动持久会话")
            return False

        from ...config import OBSERVER_MODELS
        model_raw = getattr(self._config, "observer_claude_model", "")
        # 主动校验 — 非法模型名拒绝 spawn + 明确告警（不是被动记录等人翻日志）
        if not isinstance(model_raw, str):
            self._diag(f"ABORT: model is not str: {type(model_raw).__name__}({model_raw!r})")
            logger.error(
                "observer_claude_model 不是字符串 (%s)，拒绝 spawn — 疑似 mock 泄漏",
                type(model_raw).__name__,
            )
            return False
        model = model_raw or OBSERVER_MODELS["claude_cli"]

        env = make_clean_env(extra_vars={"INDEX1_HOOK_CHILD": "1"})

        cmd = [
            claude_bin, "-p",
            "--input-format", "stream-json",
            "--output-format", "stream-json",
            "--model", model,
            "--verbose",
            "--no-session-persistence",
            "--setting-sources", "",  # 跳过 hooks/MCP 加载（14.7s→4.5s）
            "--disallowedTools", _DISALLOWED_TOOLS,
            "--permission-mode", "default",
            "You are a coding session observer. Analyze tool calls and reply with XML observations.",
        ]

        try:
            _INDEX1_HOME.mkdir(parents=True, exist_ok=True)
            self._stderr_file = open(_SESSION_LOG, "a")
        except OSError:
            self._stderr_file = None

        self._diag(f"START cmd={' '.join(cmd[:6])}... model={model}")

        try:
            self._proc = subprocess.Popen(
                cmd,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=self._stderr_file or subprocess.DEVNULL,
                env=env,
                start_new_session=True,
            )
        except OSError as e:
            self._diag(f"START FAILED: {e}")
            logger.warning("claude 持久进程启动失败: %s", e)
            self._close_stderr()
            return False

        _add_child_pid(self._proc.pid)
        self._diag(f"SPAWNED pid={self._proc.pid}, waiting for init...")

        # 等待 system:init 消息确认就绪（最多 60s — 首次 spawn 含 Node.js 启动）
        if not self._wait_for_init(timeout=60.0):
            self._diag("INIT FAILED — stopping process")
            self.stop()
            return False

        self._ready = True
        self._last_activity = time.monotonic()
        self._diag(f"READY pid={self._proc.pid}")
        logger.info(
            "Claude 持久会话就绪: session=%s, pid=%d",
            self.session_id, self._proc.pid,
        )
        return True

    def _wait_for_init(self, timeout: float = 60.0) -> bool:
        """发送触发消息并等待 system:init + 初始 assistant 响应。

        stream-json 协议（实测确认）：
        1. 进程启动 → hooks 触发（hook_started / hook_response）
        2. **stdin 触发**：进程等待 stdin 输入后才发送 system:init
        3. system:init 到达 → 进程就绪
        4. 初始 prompt 的 assistant 响应到达 → 排空后可用

        必须主动发送 stdin 消息触发 init（被动等待永远收不到）。
        """
        if not self._proc or not self._proc.stdout or not self._proc.stdin:
            return False

        deadline = time.monotonic() + timeout
        fd = self._proc.stdout.fileno()

        # 设置非阻塞模式（macOS/Linux）— try/finally 确保所有路径恢复
        os.set_blocking(fd, False)
        try:
            return self._wait_for_init_inner(fd, deadline, timeout)
        finally:
            # 无论成功/失败/异常，都恢复阻塞模式
            try:
                os.set_blocking(fd, True)
            except OSError:
                pass

    def _wait_for_init_inner(
        self, fd: int, deadline: float, timeout: float,
    ) -> bool:
        """_wait_for_init 的内层 — 非阻塞读取逻辑（由外层 try/finally 保护 set_blocking）"""
        # Phase 1: 发送触发消息 — 启动 init 序列
        # stream-json 模式下，进程等待第一条 stdin 消息后才发送 system:init
        self._diag("INIT: sending trigger message on stdin")
        try:
            trigger = json.dumps({
                "type": "user",
                "message": {
                    "role": "user",
                    "content": [{"type": "text", "text": "ready"}],
                },
            }, ensure_ascii=False)
            self._proc.stdin.write((trigger + "\n").encode("utf-8"))
            self._proc.stdin.flush()
        except (BrokenPipeError, OSError) as e:
            self._diag(f"INIT: stdin write failed: {e}")
            logger.warning("写入 claude stdin 触发消息失败: %s", e)
            return False

        # Phase 2: 读取直到收到 system:init + 首个 assistant 响应
        buf = b""
        got_init = False
        msg_count = 0

        self._diag("INIT: waiting for system:init + assistant response")

        while time.monotonic() < deadline:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                break

            ready, _, _ = select.select([fd], [], [], min(remaining, 1.0))
            if not ready:
                # 检查进程是否已退出
                if self._proc.poll() is not None:
                    self._diag(f"INIT: process exited rc={self._proc.returncode} after {msg_count} msgs")
                    logger.warning(
                        "claude 持久进程在等待 init 期间退出 (rc=%s)",
                        self._proc.returncode,
                    )
                    return False
                continue

            try:
                chunk = os.read(fd, 65536)
            except OSError:
                break
            if not chunk:
                break

            buf += chunk

            # 逐行解析 NDJSON
            while b"\n" in buf:
                line, buf = buf.split(b"\n", 1)
                line = line.strip()
                if not line:
                    continue

                try:
                    msg = json.loads(line)
                except json.JSONDecodeError:
                    self._diag(f"INIT: bad JSON: {line[:100]}")
                    continue

                msg_type = msg.get("type", "")
                subtype = msg.get("subtype", "")
                msg_count += 1

                if msg_type == "system" and subtype == _SYSTEM_INIT_SUBTYPE:
                    got_init = True
                    self._diag(f"INIT: got system:init (after {msg_count} msgs)")
                    continue

                if msg_type == "system":
                    self._diag(f"INIT: system:{subtype}")
                    continue

                if got_init and msg_type == _ASSISTANT_TYPE:
                    # 初始 prompt 的响应已收到，进程完全就绪
                    elapsed = timeout - (deadline - time.monotonic())
                    self._diag(f"INIT: got assistant response ({msg_count} msgs, {elapsed:.1f}s)")
                    return True

        # 超时但收到了 init — 也算就绪（初始响应可能较慢）
        if got_init:
            self._diag(f"INIT: timeout but got init, treating as ready ({msg_count} msgs)")
            return True

        self._diag(f"INIT: TIMEOUT after {timeout:.0f}s, got_init={got_init}, {msg_count} msgs received")
        logger.warning("等待 claude 持久进程 init 超时 (%.0fs)", timeout)
        return False

    def generate(self, prompt: str, timeout: float = 60.0) -> str | None:
        """发送 prompt 并等待 LLM 响应。

        线程安全 — 同 session 内加锁序列化。

        Args:
            prompt: 发送给 LLM 的文本
            timeout: 等待响应的超时秒数

        Returns:
            LLM 响应文本，或 None（失败/超时）
        """
        with self._lock:
            return self._generate_locked(prompt, timeout)

    def _generate_locked(self, prompt: str, timeout: float) -> str | None:
        if not self._proc or not self._ready:
            self._diag(f"GENERATE: skip (proc={self._proc is not None}, ready={self._ready})")
            return None

        # 检查进程存活
        if self._proc.poll() is not None:
            self._diag(f"GENERATE: process dead rc={self._proc.returncode}")
            logger.warning(
                "claude 持久进程已退出 (rc=%s), session=%s",
                self._proc.returncode, self.session_id,
            )
            self._ready = False
            return None

        prompt_preview = prompt[:80].replace("\n", "\\n")
        self._diag(f"GENERATE: sending {len(prompt)} chars: {prompt_preview}...")

        # 发送 user message（NDJSON 格式）
        msg = json.dumps({
            "type": "user",
            "message": {
                "role": "user",
                "content": [{"type": "text", "text": prompt}],
            },
        }, ensure_ascii=False)

        try:
            self._proc.stdin.write((msg + "\n").encode("utf-8"))
            self._proc.stdin.flush()
        except (BrokenPipeError, OSError) as e:
            self._diag(f"GENERATE: stdin write failed: {e}")
            logger.warning("写入 claude stdin 失败: %s", e)
            self._ready = False
            return None

        # 读取响应 — 等待 assistant 消息
        t0 = time.monotonic()
        result = self._read_response(timeout)
        elapsed = time.monotonic() - t0
        if result:
            self._diag(f"GENERATE: got {len(result)} chars in {elapsed:.1f}s")
        else:
            self._diag(f"GENERATE: no response after {elapsed:.1f}s")
        self._last_activity = time.monotonic()
        return result

    def _read_response(self, timeout: float) -> str | None:
        """从 stdout 读取直到收到 assistant 消息。

        stream-json 协议（实测确认）：
        - 每条 assistant 消息包含完整的一个回合响应
        - stop_reason 始终为 null（不可用于判断完成）
        - 通过检测 assistant 消息是否有 text content 判断响应完成
        - 中间可能穿插 system:hook_* 和 user echo 消息（--verbose），跳过
        """
        if not self._proc or not self._proc.stdout:
            return None

        fd = self._proc.stdout.fileno()
        deadline = time.monotonic() + timeout
        buf = b""
        msg_count = 0

        while time.monotonic() < deadline:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                break

            ready, _, _ = select.select([fd], [], [], min(remaining, 1.0))
            if not ready:
                if self._proc.poll() is not None:
                    self._diag(f"RESPONSE: process exited rc={self._proc.returncode} after {msg_count} msgs")
                    self._ready = False
                    break
                continue

            try:
                chunk = os.read(fd, 65536)
            except OSError:
                break
            if not chunk:
                self._diag(f"RESPONSE: EOF after {msg_count} msgs")
                self._ready = False
                break

            buf += chunk

            while b"\n" in buf:
                line, buf = buf.split(b"\n", 1)
                line = line.strip()
                if not line:
                    continue

                try:
                    msg = json.loads(line)
                except json.JSONDecodeError:
                    self._diag(f"RESPONSE: bad JSON: {line[:100]}")
                    continue

                msg_type = msg.get("type", "")
                msg_count += 1

                if msg_type == _ASSISTANT_TYPE:
                    # 提取 content text
                    message = msg.get("message", {})
                    content_blocks = message.get("content", [])
                    text_parts: list[str] = []
                    for block in content_blocks:
                        if isinstance(block, dict) and block.get("type") == "text":
                            text = block.get("text", "")
                            if text:
                                text_parts.append(text)

                    # stream-json: 每条 assistant 消息就是完整回合
                    if text_parts:
                        return "".join(text_parts)

                if msg_type == "system":
                    subtype = msg.get("subtype", "")
                    self._diag(f"RESPONSE: system:{subtype}")

        self._diag(f"RESPONSE: TIMEOUT after {timeout:.0f}s, {msg_count} msgs received")
        logger.warning(
            "等待 claude 响应超时 (%.0fs), session=%s",
            timeout, self.session_id,
        )
        return None

    def _close_stderr(self) -> None:
        """关闭 stderr 日志文件句柄"""
        if self._stderr_file is not None:
            try:
                self._stderr_file.close()
            except OSError:
                pass
            self._stderr_file = None

    def stop(self) -> None:
        """停止持久进程 — SIGTERM → 2s → SIGKILL"""
        self._ready = False
        if self._proc is None:
            self._close_stderr()
            return

        pid = self._proc.pid
        self._diag(f"STOP pid={pid}")
        try:
            self._proc.terminate()
            try:
                self._proc.wait(timeout=2.0)
            except subprocess.TimeoutExpired:
                self._proc.kill()
                self._proc.wait(timeout=1.0)
        except (OSError, ProcessLookupError):
            pass
        finally:
            self._proc = None
            _remove_child_pid(pid)
            self._close_stderr()

        logger.info("Claude 持久会话已关闭: session=%s, pid=%d", self.session_id, pid)

    @property
    def is_alive(self) -> bool:
        """进程是否存活且就绪"""
        return (
            self._ready
            and self._proc is not None
            and self._proc.poll() is None
        )

    @property
    def idle_seconds(self) -> float:
        """距上次活动的秒数"""
        return time.monotonic() - self._last_activity


class SessionManager:
    """管理多个 ClaudeSession — 每个 Claude Code 窗口一个持久进程

    照搬 claude-mem 架构：
    - get_or_create(): 按需 spawn，首次 13s，后续复用
    - cleanup_idle(): 定期清理空闲 session
    - stop_all(): Worker shutdown 时全量清理

    对 mem 进程堆积问题的改进：
    - cleanup 间隔 30s（mem 是 2min）
    - stop_all() 在 shutdown 时强制清理（mem 缺少 Stop hook）

    spawn 冷却保护：
    - 连续失败后暂停 spawn 60s，避免 N 个 consumer × 60s timeout 资源浪费
    """

    _SPAWN_COOLDOWN = 60.0  # spawn 失败后冷却秒数

    def __init__(self, config: Config):
        self._config = config
        self._sessions: dict[str, ClaudeSession] = {}
        self._lock = threading.Lock()
        self._total_spawned = 0
        self._last_spawn_failure: float = 0.0  # monotonic timestamp
        self._spawning: set[str] = set()  # 正在 spawn 的 session_id — 防止 N 个 consumer 同时阻塞

        # Worker 启动时清理上一轮留下的孤儿 claude 进程
        _cleanup_orphan_processes()

    def get_or_create(self, session_id: str) -> ClaudeSession | None:
        """获取或创建 session 的持久 claude 进程。

        Args:
            session_id: Claude Code 窗口的 session ID

        Returns:
            ClaudeSession — 就绪的会话
            None — 启动失败（含冷却期内跳过）
        """
        with self._lock:
            session = self._sessions.get(session_id)
            if session and session.is_alive:
                return session

            # 清理已死的 session
            if session:
                session.stop()
                del self._sessions[session_id]

            # spawn 冷却检查 — 连续失败后暂停 spawn 避免资源浪费
            if self._last_spawn_failure > 0:
                elapsed = time.monotonic() - self._last_spawn_failure
                if elapsed < self._SPAWN_COOLDOWN:
                    logger.debug(
                        "Claude spawn 冷却中 (剩余 %.0fs)",
                        self._SPAWN_COOLDOWN - elapsed,
                    )
                    return None

            # 同 session 正在 spawn — 另一个 consumer 已在等，本线程直接跳过
            if session_id in self._spawning:
                logger.debug("Claude session %s 正在 spawn，跳过", session_id)
                return None
            self._spawning.add(session_id)

        # spawn 新进程（不在锁内，避免阻塞其他 session）
        try:
            new_session = ClaudeSession(self._config, session_id)
            if not new_session.start():
                with self._lock:
                    self._last_spawn_failure = time.monotonic()
                return None
        finally:
            with self._lock:
                self._spawning.discard(session_id)

        with self._lock:
            # 双重检查 — 另一个线程可能已经创建了
            existing = self._sessions.get(session_id)
            if existing and existing.is_alive:
                new_session.stop()
                return existing

            self._sessions[session_id] = new_session
            self._total_spawned += 1
            self._last_spawn_failure = 0.0  # spawn 成功，重置冷却
            return new_session

    def cleanup_idle(self, max_idle: float = 300.0) -> int:
        """清理空闲超过 max_idle 秒的 session 进程。

        Returns:
            清理的 session 数量
        """
        to_cleanup: list[tuple[str, ClaudeSession]] = []

        with self._lock:
            for sid, session in list(self._sessions.items()):
                if not session.is_alive or session.idle_seconds > max_idle:
                    to_cleanup.append((sid, session))
                    del self._sessions[sid]

        for sid, session in to_cleanup:
            logger.info(
                "清理空闲 Claude 会话: session=%s, idle=%.0fs",
                sid, session.idle_seconds,
            )
            session.stop()

        return len(to_cleanup)

    def cleanup_session(self, session_id: str) -> bool:
        """主动清理指定 session（SessionEnd 时调用）。

        Returns:
            True — 找到并清理了
            False — session 不存在
        """
        with self._lock:
            session = self._sessions.pop(session_id, None)

        if session:
            session.stop()
            return True
        return False

    def stop_all(self) -> None:
        """停止所有持久进程 — Worker shutdown 时调用"""
        with self._lock:
            sessions = list(self._sessions.values())
            self._sessions.clear()

        for session in sessions:
            session.stop()

        if sessions:
            logger.info("已关闭 %d 个 Claude 持久会话", len(sessions))

    @property
    def stats(self) -> dict:
        """统计信息 — /health 端点用"""
        with self._lock:
            active = sum(1 for s in self._sessions.values() if s.is_alive)
            return {
                "active_sessions": active,
                "total_sessions": len(self._sessions),
                "total_spawned": self._total_spawned,
            }
