"""目标模式（Goal Mode）的工作区级状态管理。

每个工作区/项目最多存在一个活动目标。状态落盘到 `{data_dir}/goal_state.json`，
以便在切换对话、对话压缩（conversation_id 变化）、进程重启后仍能维持。

设计要点：
- 不依赖 web_terminal，只接受一个 data_dir，便于单测与复用。
- token 基线对齐工作区级累计 `{data_dir}/token_totals.json`（input/output/total）。
- review_history 保存历轮 {main_output, review_reply}，用于构建交叉结构审核文本。
"""

from __future__ import annotations

import json
import time
from copy import deepcopy
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

GOAL_STATE_FILENAME = "goal_state.json"

# 状态机
STATUS_RUNNING = "running"
STATUS_DONE = "done"
STATUS_STOPPED = "stopped"

# 停止原因
REASON_IDLE_NO_TOOL = "idle_no_tool"
REASON_MAX_TURNS = "max_turns"
REASON_MAX_TOKENS = "max_tokens"
REASON_USER_CANCEL = "user_cancel"

# 审核模式
REVIEW_MODE_READONLY = "readonly"
REVIEW_MODE_ACTIVE = "active"

PathLike = Union[str, Path]


def _empty_state() -> Dict[str, Any]:
    return {
        "active": False,
        "status": STATUS_STOPPED,
        "goal": "",
        "review_mode": REVIEW_MODE_READONLY,
        "max_turns": 5,
        "max_tokens": None,
        "start_conversation_id": None,
        "started_at": None,
        "turn_count": 0,
        "token_baseline": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0},
        "tool_call_count_baseline": 0,
        "tool_calls_used": 0,
        "review_history": [],
        "final_summary": None,
        "stopped_reason": None,
    }


class GoalStateManager:
    """工作区级目标状态。一个实例对应一个工作区的 goal_state.json。"""

    def __init__(self, data_dir: PathLike):
        self.data_dir = Path(data_dir).expanduser()
        self.state: Dict[str, Any] = self.load()

    # ------------------------------------------------------------------ 路径/持久化

    def _path(self) -> Path:
        return self.data_dir / GOAL_STATE_FILENAME

    @classmethod
    def load_from(cls, data_dir: PathLike) -> "GoalStateManager":
        """便捷构造：等价于 GoalStateManager(data_dir)。"""
        return cls(data_dir)

    def load(self) -> Dict[str, Any]:
        path = self._path()
        if not path.exists():
            self.state = _empty_state()
            return self.state
        try:
            with open(path, "r", encoding="utf-8") as fh:
                raw = json.load(fh) or {}
            merged = _empty_state()
            if isinstance(raw, dict):
                merged.update(raw)
            # 结构兜底
            if not isinstance(merged.get("review_history"), list):
                merged["review_history"] = []
            if not isinstance(merged.get("token_baseline"), dict):
                merged["token_baseline"] = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
            self.state = merged
        except (OSError, json.JSONDecodeError, ValueError):
            self.state = _empty_state()
        return self.state

    def save(self) -> None:
        path = self._path()
        path.parent.mkdir(parents=True, exist_ok=True)
        tmp = path.with_suffix(".json.tmp")
        with open(tmp, "w", encoding="utf-8") as fh:
            json.dump(self.state, fh, ensure_ascii=False, indent=2)
        tmp.replace(path)

    # ------------------------------------------------------------------ 生命周期

    def start(
        self,
        *,
        goal: str,
        review_mode: str,
        max_turns: Optional[int],
        max_tokens: Optional[int],
        conversation_id: Optional[str],
        token_baseline: Optional[Dict[str, int]] = None,
        tool_call_baseline: int = 0,
    ) -> Dict[str, Any]:
        rm = review_mode if review_mode in (REVIEW_MODE_READONLY, REVIEW_MODE_ACTIVE) else REVIEW_MODE_READONLY
        baseline = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
        if isinstance(token_baseline, dict):
            for k in baseline:
                try:
                    baseline[k] = max(0, int(token_baseline.get(k) or 0))
                except (TypeError, ValueError):
                    baseline[k] = 0
        self.state = {
            "active": True,
            "status": STATUS_RUNNING,
            "goal": str(goal or "").strip(),
            "review_mode": rm,
            "max_turns": int(max_turns) if max_turns else None,
            "max_tokens": int(max_tokens) if max_tokens else None,
            "start_conversation_id": conversation_id,
            "started_at": time.time(),
            "turn_count": 0,
            "token_baseline": baseline,
            "tool_call_count_baseline": max(0, int(tool_call_baseline or 0)),
            "tool_calls_used": 0,
            "review_history": [],
            "final_summary": None,
            "stopped_reason": None,
        }
        self.save()
        return deepcopy(self.state)

    def is_active(self) -> bool:
        return bool(self.state.get("active")) and self.state.get("status") == STATUS_RUNNING

    def get_goal(self) -> str:
        return str(self.state.get("goal") or "")

    def get_review_mode(self) -> str:
        rm = self.state.get("review_mode")
        return rm if rm in (REVIEW_MODE_READONLY, REVIEW_MODE_ACTIVE) else REVIEW_MODE_READONLY

    def get_turn(self) -> int:
        try:
            return int(self.state.get("turn_count") or 0)
        except (TypeError, ValueError):
            return 0

    def increment_turn(self) -> int:
        self.state["turn_count"] = self.get_turn() + 1
        self.save()
        return self.state["turn_count"]

    def append_review(self, main_output: str, review_reply: str) -> None:
        history = self.state.get("review_history")
        if not isinstance(history, list):
            history = []
        history.append(
            {
                "main_output": str(main_output or ""),
                "review_reply": str(review_reply or ""),
            }
        )
        self.state["review_history"] = history
        self.save()

    def mark_done(self, summary: str) -> Dict[str, Any]:
        self.state["active"] = False
        self.state["status"] = STATUS_DONE
        self.state["final_summary"] = str(summary or "")
        self.state["stopped_reason"] = None
        self.save()
        return deepcopy(self.state)

    def mark_stopped(self, reason: str) -> Dict[str, Any]:
        self.state["active"] = False
        self.state["status"] = STATUS_STOPPED
        self.state["stopped_reason"] = reason
        self.save()
        return deepcopy(self.state)

    def clear(self) -> None:
        """完全清除目标状态（用户取消）。"""
        self.state = _empty_state()
        self.save()

    # ------------------------------------------------------------------ 边界判定

    def reached_max_turns(self) -> bool:
        mt = self.state.get("max_turns")
        if not mt:
            return False
        return self.get_turn() >= int(mt)

    def reached_max_tokens(self, current_total_tokens: int) -> bool:
        """current_total_tokens 为工作区当前累计 total_tokens。"""
        mt = self.state.get("max_tokens")
        if not mt:
            return False
        used = self.tokens_used(current_total_tokens)
        return used >= int(mt)

    def tokens_used(self, current_total_tokens: int) -> int:
        baseline = self.state.get("token_baseline") or {}
        try:
            base_total = int(baseline.get("total_tokens") or 0)
        except (TypeError, ValueError):
            base_total = 0
        return max(0, int(current_total_tokens or 0) - base_total)

    # ------------------------------------------------------------------ 审核输入构建

    def build_review_payload_text(self, current_main_output: str) -> str:
        """生成交叉结构的单条 user 文本：目标 + 历轮(主模型输出/审核反馈) + 当前轮待审。"""
        goal = self.get_goal()
        lines: List[str] = ["【本次目标】", goal or "（目标为空）", ""]
        history = self.state.get("review_history") or []
        for idx, entry in enumerate(history, start=1):
            if not isinstance(entry, dict):
                continue
            lines.append(f"【第{idx}轮】")
            lines.append(f"主执行模型输出：{str(entry.get('main_output') or '').strip()}")
            lines.append(f"你的反馈：{str(entry.get('review_reply') or '').strip()}")
            lines.append("")
        lines.append("【当前轮·待审核】")
        lines.append(f"主执行模型输出：{str(current_main_output or '').strip()}")
        return "\n".join(lines)

    # ------------------------------------------------------------------ 前端快照

    def progress_snapshot(self, *, current_total_tokens: int = 0, current_tool_calls: int = 0) -> Dict[str, Any]:
        started_at = self.state.get("started_at")
        duration = None
        if started_at:
            try:
                duration = max(0.0, time.time() - float(started_at))
            except (TypeError, ValueError):
                duration = None
        try:
            tool_baseline = int(self.state.get("tool_call_count_baseline") or 0)
        except (TypeError, ValueError):
            tool_baseline = 0
        try:
            stored_tool_calls = max(0, int(self.state.get("tool_calls_used") or 0))
        except (TypeError, ValueError):
            stored_tool_calls = 0
        current_tool_calls_used = max(0, int(current_tool_calls or 0) - tool_baseline)
        tool_calls_used = max(stored_tool_calls, current_tool_calls_used)
        if tool_calls_used != stored_tool_calls:
            self.state["tool_calls_used"] = tool_calls_used
            try:
                self.save()
            except Exception:
                pass
        return {
            "goal": self.get_goal(),
            "status": self.state.get("status"),
            "turn_count": self.get_turn(),
            "tokens_used": self.tokens_used(current_total_tokens),
            "tool_calls": tool_calls_used,
            "duration_seconds": duration,
            "review_mode": self.get_review_mode(),
            "max_turns": self.state.get("max_turns"),
            "max_tokens": self.state.get("max_tokens"),
            "final_summary": self.state.get("final_summary"),
            "stopped_reason": self.state.get("stopped_reason"),
        }