#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# file name is agent_executor.py

import os
import re
import json
import logging
import asyncio
import requests
from html.parser import HTMLParser
from pathlib import Path
import time
from toolUtil import TOOLS, execute_tool, CFG

try:
    import websockets
except ImportError:
    websockets = None

# ---------- 配置（支持环境变量） ----------
WORK_DIR_ENV = os.getenv("AGENT_WORK_DIR", "/work")
WS_HOST = os.getenv("AGENT_WS_HOST", "127.0.0.1")
WS_PORT = int(os.getenv("AGENT_WS_PORT", 8765))

LLAMA_HOST = os.getenv("AGENT_LLAMA_HOST", "http://127.0.0.1:9931")
LLAMA_API_KEY = os.getenv("AGENT_LLAMA_API_KEY", "sk-no-key-required")
LLAMA_MODEL = os.getenv("AGENT_LLAMA_MODEL", "local-model")
LLAMA_TIMEOUT = int(os.getenv("AGENT_LLAMA_TIMEOUT", 600))

# ---------- 工作目录 ----------
AGENT_DIR = Path(__file__).parent
WORK_DIR = Path(WORK_DIR_ENV).resolve()
WORK_DIR.mkdir(exist_ok=True, parents=True)

# ---------- 日志配置（轮转） ----------
logger = logging.getLogger('AgentExecutor')
logger.setLevel(logging.INFO)
if not logger.handlers:
    console = logging.StreamHandler()
    console.setLevel(logging.INFO)
    logger.addHandler(console)

# ---------- 工具配置（无沙箱） ----------
CFG.root = Path(WORK_DIR).resolve()
CFG.root.mkdir(parents=True, exist_ok=True)
CFG.audit_log = CFG.root / "audit.log"


def load_web_prompt():
    """读取 prompt.md 作为系统提示词，供客户端 AI 使用。"""
    prompt_file = AGENT_DIR / "prompt-web.md"
    if prompt_file.exists():
        text = prompt_file.read_text(encoding='utf-8').strip()
        if text:
            return text
    return ""


def load_local_system_prompt():
    """给本地 LLM 用的 system prompt。优先读 prompt_local_llm.md，否则用内置。"""
    f = AGENT_DIR / "prompt_local_llm.md"
    if f.exists():
        text = f.read_text(encoding='utf-8').strip()
        if text:
            return text
    return ""


# ---------- 响应封装 ----------
def make_response(status, data=None, message=None):
    resp = {"status": status}
    if data is not None:
        resp["data"] = data
    if message is not None:
        resp["message"] = message
    return resp


# ---------- 消息分段：把网页 AI 的回复按 @小牛 / @用户 切开 ----------
_BLOCK_TAGS = {
    "p", "div", "br", "li", "ul", "ol", "tr", "td", "th", "table",
    "h1", "h2", "h3", "h4", "h5", "h6", "blockquote", "pre", "hr",
}


class _HTMLTextExtractor(HTMLParser):
    """把 AI 回复的 HTML 转成带段落换行的纯文本。"""

    def __init__(self):
        super().__init__(convert_charrefs=True)
        self._parts = []
        self._skip = 0

    def _newline(self):
        if self._parts and not self._parts[-1].endswith("\n"):
            self._parts.append("\n")

    def handle_starttag(self, tag, attrs):
        if tag in ("script", "style"):
            self._skip += 1
        elif tag == "br":
            self._parts.append("\n")
        elif tag == "li":
            self._newline()
            self._parts.append("- ")
        elif tag in _BLOCK_TAGS:
            self._newline()

    def handle_endtag(self, tag):
        if tag in ("script", "style"):
            self._skip = max(0, self._skip - 1)
        elif tag in _BLOCK_TAGS:
            self._newline()

    def handle_data(self, data):
        if not self._skip:
            self._parts.append(data)

    def text(self):
        text = "".join(self._parts)
        text = re.sub(r"[ \t]+\n", "\n", text)
        text = re.sub(r"\n{3,}", "\n\n", text)
        return text.strip()


def html_to_text(html: str) -> str:
    parser = _HTMLTextExtractor()
    parser.feed(html)
    parser.close()
    return parser.text()


_OCCURRENCE_RE = re.compile(r"@(小牛|用户)")
_LINE_PREFIX_RE = re.compile(r"[ \t]*(?:[-*]\s+)?\Z")


def _at_line_start(text: str, pos: int) -> bool:
    """pos 处的标记是不是段落开头（允许前面有空格或列表符号）。"""
    line_start = text.rfind("\n", 0, pos) + 1
    return _LINE_PREFIX_RE.match(text[line_start:pos]) is not None


def split_for_niux(text: str):
    """只保留发给 @小牛 的段落，@用户 的内容不喂给本地模型。

    规则：从段落开头的 @小牛 开始，到下一个 @用户 或下一个 @小牛 段落为止。
    没有段落级标记、但正文里提到 @小牛 时，退化成“从第一个 @小牛 截到下一个 @用户”。
    """
    occur = [(m.start(), m.group(1)) for m in _OCCURRENCE_RE.finditer(text)]
    if not occur:
        return None

    starts = [p for p, who in occur if who == "小牛" and _at_line_start(text, p)]
    if not starts:
        starts = [p for p, who in occur if who == "小牛"]
    if not starts:
        return None

    starts = sorted(set(starts))
    user_positions = [p for p, who in occur if who == "用户"]

    segments = []
    for i, start in enumerate(starts):
        next_start = starts[i + 1] if i + 1 < len(starts) else len(text)
        next_user = next((p for p in user_positions if p > start), len(text))
        segment = text[start:min(next_start, next_user)].strip()
        # 切在 @用户 前时可能留下下一段的列表符号
        segment = re.sub(r"\n[ \t]*[-*][ \t]*$", "", segment).strip()
        if len(segment) > 3:   # 只剩一个 "@小牛" 就算了
            segments.append(segment)

    return "\n\n".join(segments) if segments else None


# ---------- 模型调用 ----------
def chat_completion(messages,
                    tools=None,
                    host=None,
                    api_key=None,
                    model=None,
                    temperature=0.7,
                    stream=False,
                    timeout=None):
    host = host or LLAMA_HOST
    api_key = api_key or LLAMA_API_KEY
    model = model or LLAMA_MODEL
    timeout = timeout or LLAMA_TIMEOUT

    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    payload = {
        "model": model,
        "messages": messages,
        "temperature": temperature,
        "stream": bool(stream),
    }
    if tools:
        payload["tools"] = tools
        payload["tool_choice"] = "auto"

    url = f"{host}/v1/chat/completions"
    max_attempts = 3
    r = None
    for attempt in range(1, max_attempts + 1):
        try:
            r = requests.post(
                url,
                json=payload,
                headers=headers,
                timeout=timeout,
            )
        except (requests.ConnectionError, requests.Timeout) as e:
            logger.warning(f"LLM request failed (attempt {attempt}/{max_attempts}): {e}")
            if attempt < max_attempts:
                time.sleep(2 ** (attempt - 1))
                continue
            raise RuntimeError(f"LLM connection failed after {max_attempts} attempts: {e}")

        if r.status_code >= 500:
            logger.warning(f"LLM returned {r.status_code} (attempt {attempt}/{max_attempts})")
            if attempt < max_attempts:
                time.sleep(2 ** (attempt - 1))
                continue
            raise RuntimeError(f"LLM {r.status_code}: {r.text[:500]}")

        if r.status_code >= 400:
            body = r.text[:2000]
            logger.error(f"LLM returned {r.status_code}: {body}")
            raise RuntimeError(f"LLM {r.status_code}: {body[:500]}")

        break

    if stream:
        return r
    return r.json()


# ---------- Agent 循环 ----------
def run_agent_once(chat_completion_fn, messages):
    """
    单轮工具模式：
      1. 调模型，拿到 tool_calls
      2. 执行所有工具
      3. 直接返回工具执行结果，不再回传给模型
      4. 若模型没有请求工具，则直接返回它的文本回复
    """
    resp = chat_completion_fn(
        messages=messages,
        tools=TOOLS,
        stream=False,
    )

    try:
        msg = resp["choices"][0]["message"]
    except (KeyError, IndexError) as e:
        logger.error(f"模型响应结构异常: {e}, resp={resp}")
        raise RuntimeError(f"模型响应结构异常: {e}")

    content_ = msg.get("content") or ""
    if content_:
        logger.info(f"🤖 : {content_[:200]}")

    tool_calls = msg.get("tool_calls") or []

    # ---- 情况 A：模型没有请求工具，直接返回文本 ----
    if not tool_calls:
        logger.info(f"模型未请求工具，直接返回文本 {len(content_)} 字符")
        return {
            "type": "text",
            "content": content_,
        }

    # ---- 情况 B：模型请求了工具，逐个执行，收集结果 ----
    logger.info(f"模型请求 {len(tool_calls)} 个工具调用，开始执行")

    results = []
    for tc in tool_calls:
        fn = tc.get("function") or {}
        name = fn.get("name")
        args = fn.get("arguments")

        logger.info(f"  → 执行工具 {name}, 参数 {str(args)[:200]}")

        try:
            results.append(execute_tool(name, args))
        except Exception as e:
            logger.error(f"工具 {name} 执行失败: {e}", exc_info=True)
            results.append(f"{name} 失败：{e}")

    logger.info(f"工具全部执行完毕，共 {len(results)} 个")

    body = "\n".join(str(r) for r in results)
    content = (content_.strip() + "\n" + body).strip() if content_.strip() else body

    return {
        "type": "tools",
        "content": content,
    }


# ---------- 核心处理：一条客户端消息 ----------
async def handle_client_message(data):
    """
    处理来自 WebSocket 客户端的一条消息。
    返回 response dict（不含 id，调用方补）。
    """
    # ---- action == 'prompt' ----
    if data.get('action') == 'prompt':
        try:
            prompt = load_web_prompt()
            return make_response("success", data={"prompt": prompt})
        except Exception as e:
            logger.error(f"获取提示词失败: {e}")
            return make_response("error", message=str(e))

    # ---- 其他未知 action ----
    if 'action' in data:
        return make_response("error", message=f"不支持的 action: {data.get('action')}")

    # ---- 无 action：视为客户端回传的 AI 回复 ----
    content = data.get('content')
    if content is None:
        return make_response("error", message="缺少 content 字段")
    if not isinstance(content, str):
        return make_response("error", message=f"content 字段类型错误: {type(content).__name__}")
    if not content.strip():
        return make_response("skipped", message="content 为空，已跳过")

    # ---- 只取发给 @小牛 的段落，转成纯文本 ----
    payload = split_for_niux(html_to_text(content))
    if not payload:
        return make_response("skipped", message="没有发给小牛的内容")

    # ---- 组装 messages 并跑 Agent ----
    try:
        system_prompt = load_local_system_prompt()
    except Exception as e:
        logger.warning(f"加载 prompt 失败，使用默认: {e}")
        system_prompt = ""

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": payload},
    ]

    logger.info(f"小牛段落 {len(payload)} 字符（网页原文 {len(content)} 字符）")
    logger.info(f"😊 : {payload[:500]}")
    try:
        outcome = await asyncio.to_thread(run_agent_once, chat_completion, messages)

        if outcome["type"] == "text":
            text = (outcome["content"] or "").strip()
            # 小牛判断这条消息与它无关，或什么都没说：不回灌，避免死循环
            if not text or "不关我事" in text:
                return make_response("skipped", message="小牛未执行操作")
            return make_response("success", data="小牛：" + text)

        return make_response("success", data="小牛：" + outcome["content"])

    except Exception as e:
        logger.error(f"Agent 执行异常: {e}", exc_info=True)
        return make_response("error", message=f"Agent 执行失败: {e}")


# ---------- WebSocket 处理 ----------
async def websocket_handler(websocket):
    peer = getattr(websocket, "remote_address", None)
    logger.info(f"✅ 客户端已连接: {peer}")
    try:
        async for message in websocket:
            # 二进制帧兼容
            if isinstance(message, bytes):
                try:
                    message = message.decode('utf-8')
                except UnicodeDecodeError:
                    logger.error(f"无法用 UTF-8 解码二进制数据: {message!r}")
                    await websocket.send(json.dumps(
                        make_response("error", message="无法解码二进制消息"),
                        ensure_ascii=False))
                    continue

            # 解析 JSON
            try:
                data = json.loads(message)
            except json.JSONDecodeError:
                await websocket.send(json.dumps(
                    make_response("error", message="无效的 JSON"),
                    ensure_ascii=False))
                continue

            if not isinstance(data, dict):
                await websocket.send(json.dumps(
                    make_response("error", message="JSON 消息必须是对象"),
                    ensure_ascii=False))
                continue

            request_id = data.get('id')

            try:
                result = await handle_client_message(data)
            except Exception as e:
                logger.error(f"处理单条消息异常: {e}", exc_info=True)
                result = make_response("error", message=f"内部错误: {e}")

            if result is None:
                continue
            if request_id is not None:
                result['id'] = request_id

            try:
                await websocket.send(json.dumps(result, ensure_ascii=False))
            except Exception as e:
                logger.error(f"发送响应失败: {e}")

    except websockets.exceptions.ConnectionClosed:
        logger.info("WebSocket 连接关闭")
    except Exception as e:
        logger.error(f"WebSocket 处理异常: {e}", exc_info=True)


# ---------- 启动 WebSocket 服务 ----------
if __name__ == '__main__':
    if websockets is None:
        print("❌ 未安装 websockets 库，请执行: pip install websockets")
        exit(1)

    print("🚀 Agent WebSocket 服务启动中")
    print(f"📁 工作目录 (默认相对根): {WORK_DIR}")
    print(f"🔌 WebSocket 监听地址: {WS_HOST}:{WS_PORT}")
    print(f"🧠 模型接口: {LLAMA_HOST}  model={LLAMA_MODEL}")
    web_prompt_file = AGENT_DIR / "prompt.md"
    local_prompt_file = AGENT_DIR / "prompt_local_llm.md"
    print(f"📝 网页 AI 提示词: {web_prompt_file if web_prompt_file.exists() else '内置默认'}")
    print(f"🧠 本地 LLM 提示词: {local_prompt_file if local_prompt_file.exists() else '内置默认'}")
    print(f"🧰 工具范围: 文件/命令/网络均不限制，路径默认为工作目录")


    async def start_ws():
        async with websockets.serve(websocket_handler, WS_HOST, WS_PORT):
            await asyncio.Future()  # 永久运行

    print("🚀 Agent WebSocket 服务启动完毕")

    try:
        asyncio.run(start_ws())
    except KeyboardInterrupt:
        print("服务已停止")