"""开登录后使用者的问数据（D-c）：能发起什么样的运行、查询结果在工具边界遮、运行的 WebSocket 按运行判、数据源名单。

- 图：只许输入、模型、Agent、报告、成果、知识检索、数据整形、条件分支、口径卡、合并查询；Agent 只许数据源工具和只读的
  内置工具（app/auth/chat_policy.py）。调用工具、沙箱代码、自定义工具、MCP、网页、HTTP、文件一律 403；
- 使用者发起的问数据运行一定带 MASK_RESULTS_OPTION：真的执行一次（演示模型调数据源查询），事件、快照、实时流里都没有
  遮罩列的原值；
- WebSocket：没登录的握手前关掉；使用者只能连自己能读、查询结果在工具边界遮过的运行（握手时判，1008 拒）；
- GET /api/datasources/names：开了问数据的使用者能读，只有名字、说明、种类，只列启用着的。

数据源是现造的 SQLite 文件，姓名、手机号都是虚构的。
"""
from __future__ import annotations

import asyncio
import json
import sqlite3
import uuid

import pytest
from langchain_core.messages import AIMessage, ToolMessage
from sqlalchemy import delete, select

from app.auth.chat_policy import chat_graph_problem
from app.data import introspect
from app.db.base import SessionLocal
from app.db.models import Artifact, DataSource, Run, RunEvent
from app.engine.runner import run_manager
from app.tools.datasource import MASK_RESULTS_OPTION
from tests.fakes import auth_env

PHONE = "13900000001"
MASKED = "已遮罩"


def node(nid, ntype, **config):
    return {"id": nid, "type": ntype, "position": {"x": 0, "y": 0}, "data": {"label": nid, "config": config}}


def chain(*nodes):
    return {"nodes": list(nodes), "edges": [{"source": a["id"], "target": b["id"]} for a, b in zip(nodes, nodes[1:])]}


def chat_graph(*tools: str, middle: dict | None = None) -> dict:
    return chain(node("start", "input", fields=[{"name": "question"}]),
                 middle or node("ask", "agent", prompt="{{ input.question }}", tools=list(tools), max_steps=4),
                 node("out", "output", fields=[{"name": "回答", "value": "{{ nodes.ask.text }}"}]))


# ---------------------------------------------------------------- 图的白名单（纯函数）


@pytest.mark.parametrize("graph", [
    chat_graph("db_query__crm", "db_schema__crm", "calculator", "current_time", "json_query", "kb_search",
               "memory_recall"),
    chain(node("start", "input"), node("kb", "retrieve", query="x"), node("t", "transform", expression="1"),
          node("out", "output")),
    {"nodes": "坏的"},          # 不成形的图交给 launch_run 照常报校验问题
])
def test_graphs_a_user_may_run_in_chat(graph):
    assert chat_graph_problem(graph) is None


@pytest.mark.parametrize("graph, said", [
    (chat_graph(middle=node("ask", "code", code="print(1)")), "沙箱代码"),
    (chat_graph(middle=node("ask", "tool", tool="calculator")), "调用工具"),
    (chat_graph(middle=node("ask", "supervisor", agents=[])), "多 Agent 协作"),
    (chat_graph(middle=node("ask", "subgraph", workflow_id="w")), "子工作流"),
    (chat_graph(middle=node("ask", "human", mode="approve")), "人工审批"),
    (chat_graph(middle=node("ask", "memory", action="write")), "长期记忆"),
    (chat_graph("http_request"), "http_request"),
    (chat_graph("web_fetch"), "web_fetch"),
    (chat_graph("python_exec"), "python_exec"),
    (chat_graph("file_read"), "file_read"),
    (chat_graph("memory_write"), "memory_write"),
    (chat_graph("mcp:fs/read_file"), "mcp:fs/read_file"),
    (chat_graph("my_custom_http"), "my_custom_http"),
    (chat_graph("{{ vars.tool }}"), "模板"),
])
def test_graphs_a_user_may_not_run_in_chat(graph, said):
    problem = chat_graph_problem(graph)
    assert problem and problem.startswith("没有权限") and said in problem, problem


@pytest.mark.parametrize("graph", [
    chat_graph(middle=node("ask", "agent", prompt="x", tools=["memory_recall"], memory_scope="builders")),
    chat_graph(middle=node("kb", "retrieve", query="x", collection="hr")),
    {**chat_graph("kb_search"), "defaults": {"collection": "hr"}},
    {**chat_graph("memory_recall"), "defaults": {"memory_scope": "  "}},
])
def test_users_may_not_pick_memory_or_knowledge_scope(graph):
    """记忆、知识库的内容不经遮罩：节点上、图级 defaults 里都不许指定范围，用管理员设好的默认值。"""
    problem = chat_graph_problem(graph)
    assert problem and "记忆范围和知识库" in problem


def test_tools_from_graph_defaults_count_too():
    """Agent 上没写工具时运行时用图级默认值里的：白名单同样按它查（engine/launch.tool_bindings）。"""
    graph = chain(node("start", "input"), node("ask", "agent", prompt="x"), node("out", "output"))
    graph["defaults"] = {"tools": ["web_search"]}
    problem = chat_graph_problem(graph)
    assert problem and "web_search" in problem


# ---------------------------------------------------------------- 接口


async def _masked_source(tmp_path) -> str:
    path = tmp_path / f"{uuid.uuid4().hex[:8]}.db"
    db = sqlite3.connect(path)
    db.executescript("CREATE TABLE customers (id INTEGER PRIMARY KEY, name TEXT, phone TEXT, city TEXT);"
                     f"INSERT INTO customers VALUES (1, '张三', '{PHONE}', '杭州'), (2, '李四', '13900000002', '宁波');")
    db.commit()
    db.close()
    name = f"crm_{uuid.uuid4().hex[:6]}"
    async with SessionLocal() as session:
        row = DataSource(name=name, kind="sqlite", database=str(path), readonly=True, enabled=True,
                         description="客户", options={"mask_columns": "phone"})
        session.add(row)
        await session.flush()
        row.schema_cache = await introspect.introspect(row)
        session.add(DataSource(name=f"off_{uuid.uuid4().hex[:6]}", kind="sqlite", database=str(path), enabled=False))
        await session.commit()
    return name


@pytest.fixture
async def world(tmp_path):
    await auth_env.reset()
    await run_manager.setup()
    admin = await auth_env.turn_on_local()
    await auth_env.make_user("uma", "user")
    await auth_env.make_user("bob", "builder")
    source = await _masked_source(tmp_path)
    uma = await auth_env.login("uma")
    bob = await auth_env.login("bob")
    try:
        yield {"admin": admin, "uma": uma, "bob": bob, "source": source}
    finally:
        await run_manager.shutdown()
        async with SessionLocal() as s:
            await s.execute(delete(Run).where(Run.started_by.in_(["uma", "bob"])))
            await s.execute(delete(DataSource))
            await s.commit()
        await auth_env.reset()


async def _allow_chat(world) -> None:
    r = await world["admin"].put("/api/auth/settings", json={"user_can_chat": True})
    assert r.status_code == 200, r.text


async def _finish(run_id: str) -> Run:
    for _ in range(400):
        async with SessionLocal() as session:
            row = await session.get(Run, run_id)
        if row.status in ("succeeded", "failed", "interrupted", "cancelled"):
            await run_manager.wait_idle(run_id)
            return row
        await asyncio.sleep(0.05)
    raise AssertionError("运行没有结束")


def _script(monkeypatch, tool: str) -> None:
    """演示模型：第一轮查遮罩列，看到结果后作答。"""
    from app.providers import mock_model

    def _decide(self, messages):
        if any(isinstance(m, ToolMessage) for m in messages):
            return AIMessage(content="杭州有一位客户。")
        return AIMessage(content="", tool_calls=[{"name": tool, "args": {"sql": "SELECT name, phone, city FROM customers"},
                                                  "id": "call_1"}])

    monkeypatch.setattr(mock_model.MockChatModel, "_decide", _decide)


async def test_a_user_can_only_launch_whitelisted_chat_graphs(world):
    await _allow_chat(world)
    uma = world["uma"]
    r = await uma.post("/api/runs", json={"origin": "chat", "graph": chat_graph("http_request"), "input": {}})
    assert r.status_code == 403 and "http_request" in r.json()["detail"]
    r = await uma.post("/api/runs", json={"origin": "chat", "graph": chat_graph(middle=node("ask", "code"))})
    assert r.status_code == 403 and "沙箱代码" in r.json()["detail"]
    r = await uma.post("/api/runs", json={"origin": "chat", "run_class": "formal", "workflow_id": "x"})
    assert r.status_code == 403, "正式运行只经应用入口发起"
    # 搭建者照旧不受这份白名单约束（发起前检查另说：这里只看不是 403）
    r = await world["bob"].post("/api/runs", json={"origin": "chat", "graph": chat_graph(middle=node("ask", "code"))})
    assert r.status_code != 403


async def test_a_users_chat_run_is_masked_end_to_end(world, monkeypatch):
    """真的执行一次：使用者发起、Agent 查了遮罩列。运行选项带着遮罩开关，事件、查询快照、实时流里都没有原值。"""
    await _allow_chat(world)
    tool = f"db_query__{world['source']}"
    _script(monkeypatch, tool)
    r = await world["uma"].post("/api/runs", json={"origin": "chat", "graph": chat_graph(tool),
                                                   "input": {"question": "杭州有几位客户"}})
    assert r.status_code == 201, r.text
    run = await _finish(r.json()["id"])
    assert run.status == "succeeded", run.error
    assert run.options.get(MASK_RESULTS_OPTION) is True and run.started_by == "uma" and run.origin_kind == "chat"
    async with SessionLocal() as s:
        events = list((await s.execute(select(RunEvent).where(RunEvent.run_id == run.id))).scalars())
        snaps = list((await s.execute(select(Artifact).where(Artifact.run_id == run.id,
                                                             Artifact.kind == "query_snapshot"))).scalars())
    dumped = json.dumps([e.data for e in events], ensure_ascii=False)
    assert "tool.end" in {e.type for e in events}
    assert PHONE not in dumped and "张三" in dumped, "遮罩列换掉，别的列照旧"
    assert snaps, "查询照常落快照"
    from app.core.artifact_store import load

    for snap in snaps:
        assert PHONE not in json.dumps(load(snap.id), ensure_ascii=False)

    # 使用者能读自己这次运行的记录、事件，也能连实时流；实时流里同样没有原值
    assert (await world["uma"].get(f"/api/runs/{run.id}/events")).status_code == 200
    frames = await _ws(f"/api/runs/{run.id}/stream", _cookie(world["uma"]))
    assert frames[0]["type"] == "websocket.accept"
    texts = [f.get("text") or "" for f in frames if f["type"] == "websocket.send"]
    assert texts and '"stream.end"' in texts[-1] and PHONE not in "".join(texts)


async def test_a_users_chat_run_cannot_write_a_writable_source(world, monkeypatch, tmp_path):
    """可写源 + Agent 自己写了 approval=never：写语句照样被守卫拒掉、不执行（只读不交给审批），表里的行一条没少。
    搭建者发起的同一张图照旧能写——只读只给 user 角色的问数据运行。"""
    from app.providers import mock_model
    from app.tools.datasource import READ_ONLY_OPTION

    path = tmp_path / "ledger.db"
    db = sqlite3.connect(path)
    db.executescript("CREATE TABLE ledger (id INTEGER PRIMARY KEY, amount INTEGER);"
                     "INSERT INTO ledger VALUES (1, 10), (2, 20);")
    db.commit()
    db.close()
    name = f"ledger_{uuid.uuid4().hex[:6]}"
    async with SessionLocal() as session:
        row = DataSource(name=name, kind="sqlite", database=str(path), readonly=False, enabled=True)
        session.add(row)
        await session.flush()
        row.schema_cache = await introspect.introspect(row)
        await session.commit()
    replies: list[str] = []

    def _decide(self, messages):
        if seen := [m for m in messages if isinstance(m, ToolMessage)]:
            replies.append(str(seen[-1].content))
            return AIMessage(content="好了。")
        return AIMessage(content="", tool_calls=[{"name": f"db_query__{name}", "args": {"sql": "DELETE FROM ledger"},
                                                  "id": "call_1"}])

    monkeypatch.setattr(mock_model.MockChatModel, "_decide", _decide)
    graph = chat_graph(middle=node("ask", "agent", prompt="{{ input.question }}", tools=[f"db_query__{name}"],
                                   max_steps=4, approval="never"))
    await _allow_chat(world)
    r = await world["uma"].post("/api/runs", json={"origin": "chat", "graph": graph, "input": {"question": "清空"}})
    assert r.status_code == 201, r.text
    run = await _finish(r.json()["id"])
    assert run.status == "succeeded", run.error
    assert run.options.get(READ_ONLY_OPTION) is True
    assert replies and replies[-1].startswith("SQL 被拒绝"), replies
    assert sqlite3.connect(path).execute("SELECT COUNT(*) FROM ledger").fetchone()[0] == 2

    replies.clear()
    r = await world["bob"].post("/api/runs", json={"origin": "chat", "graph": graph, "input": {"question": "清空"}})
    assert r.status_code == 201, r.text
    run = await _finish(r.json()["id"])
    assert run.status == "succeeded", run.error
    assert not (run.options or {}).get(READ_ONLY_OPTION)
    assert sqlite3.connect(path).execute("SELECT COUNT(*) FROM ledger").fetchone()[0] == 0, "搭建者照旧能写"


async def test_a_users_chat_run_keeps_to_its_own_memory_scope(world):
    """请求体里带记忆范围、知识库的拒；不带的，记忆是这位使用者自己的范围，知识库是设置里的默认值。"""
    await _allow_chat(world)
    uma = world["uma"]
    graph = chat_graph("memory_recall")
    for extra in ({"memory_scope": "default"}, {"collection": "hr"}):
        r = await uma.post("/api/runs", json={"origin": "chat", "graph": graph, "input": {}, **extra})
        assert r.status_code == 403 and "记忆范围和知识库" in r.json()["detail"]
    r = await uma.post("/api/runs", json={"origin": "chat", "graph": graph, "input": {"question": "x"}})
    assert r.status_code == 201, r.text
    body = r.json()
    assert body["memory_scope"] == "chat:uma" and body["collection"] == "default"
    await _finish(body["id"])


# ---------------------------------------------------------------- 问数据的生成流


class _Scripted:
    """第一轮按节点手册的标准写法搭（tool 节点查库），收到自查请求时改成 Agent 查库。记下 system 和每一轮的请求。"""

    def __init__(self, source: str) -> None:
        self.source = source
        self.systems: list[str] = []
        self.humans: list[str] = []

    def _lines(self, repair: bool) -> list[str]:
        if repair:
            ops = [{"op": "remove_node", "id": "fetch"},
                   {"op": "add_node", "node": {"id": "ask", "type": "agent", "label": "查库",
                                               "config": {"prompt": "{{ input.question }}",
                                                          "tools": [f"db_query__{self.source}"]}}},
                   {"op": "add_edge", "edge": {"source": "start", "target": "ask"}},
                   {"op": "add_edge", "edge": {"source": "ask", "target": "out"}},
                   {"op": "done", "explanation": "改成 Agent 查库"}]
        else:
            ops = [{"op": "plan", "summary": "查客户"},
                   {"op": "add_node", "node": {"id": "start", "type": "input", "label": "输入",
                                               "config": {"fields": [{"name": "question"}]}}},
                   {"op": "add_node", "node": {"id": "fetch", "type": "tool", "label": "查库",
                                               "config": {"tool": f"db_query__{self.source}",
                                                          "args": {"sql": "SELECT city FROM customers"}}}},
                   {"op": "add_node", "node": {"id": "out", "type": "output", "label": "成果",
                                               "config": {"fields": [{"name": "回答", "value": "{{ last_message }}"}]}}},
                   {"op": "add_edge", "edge": {"source": "start", "target": "fetch"}},
                   {"op": "add_edge", "edge": {"source": "fetch", "target": "out"}},
                   {"op": "done", "explanation": "查好了", "run": True}]
        return [json.dumps(o, ensure_ascii=False) for o in ops]

    async def astream(self, messages):
        from langchain_core.messages import AIMessageChunk

        self.systems.append(next(text for role, text in messages if role == "system"))
        human = next(text for role, text in messages if role == "human")
        self.humans.append(human)
        for line in self._lines("上线前自查" in human):
            yield AIMessageChunk(content=line + "\n")


async def _generate(client, monkeypatch, model: _Scripted) -> list[dict]:
    import app.api.copilot as copilot

    async def _model(*_a, **_k):
        return model, "mock-fast"

    monkeypatch.setattr(copilot, "get_chat_model", _model)
    events: list[dict] = []
    async with client.stream("POST", "/api/copilot/generate-stream", json={"instruction": "杭州有几位客户"}) as r:
        assert r.status_code == 200
        async for line in r.aiter_lines():
            if line.startswith("data: "):
                events.append(json.loads(line[6:]))
    return events


async def test_generation_for_a_user_keeps_to_the_chat_policy(world, monkeypatch):
    """使用者问数据时，提示词先说清只能 Agent 查库；照标准写法搭出的 tool 节点图按 chat_policy 交回去改，改好才交付、
    自动运行。搭建者照旧按节点手册来，tool 节点不算问题。"""
    await _allow_chat(world)
    model = _Scripted(world["source"])
    events = await _generate(world["uma"], monkeypatch, model)
    assert "这位用户是「使用者」" in model.systems[0] and "- http_request" not in model.systems[0]
    assert len(model.humans) == 2 and "调用工具" in model.humans[1], "tool 节点按 chat_policy 交回去改"
    check = [e for e in events if e.get("op") == "check"][-1]
    final = next(e for e in events if e.get("op") == "final")
    assert check["status"] == "passed" and final["autorun"] is True
    assert {n["type"] for n in final["graph"]["nodes"]} == {"input", "agent", "output"}
    assert chat_graph_problem(final["graph"]) is None

    model = _Scripted(world["source"])
    events = await _generate(world["bob"], monkeypatch, model)
    assert "这位用户是「使用者」" not in model.systems[0] and len(model.humans) == 1
    final = next(e for e in events if e.get("op") == "final")
    assert "tool" in {n["type"] for n in final["graph"]["nodes"]}


async def test_generation_for_a_user_reports_what_it_could_not_fix(world, monkeypatch):
    """改不好的照实交付：自查失败、final 里列出 chat.not_allowed，不自动运行。"""
    await _allow_chat(world)
    model = _Scripted(world["source"])
    model._lines = lambda repair: _Scripted(world["source"])._lines(False) if not repair else [
        json.dumps({"op": "done", "explanation": "没改"}, ensure_ascii=False)]
    events = await _generate(world["uma"], monkeypatch, model)
    check = [e for e in events if e.get("op") == "check"][-1]
    final = next(e for e in events if e.get("op") == "final")
    assert check["status"] == "failed" and any(i["code"] == "chat.not_allowed" for i in check["issues"])
    assert final["autorun"] is False and any(i["code"] == "chat.not_allowed" for i in final["issues"])


# ---------------------------------------------------------------- WebSocket


def _cookie(client) -> str:
    return f"agentlab_session={client.cookies.get('agentlab_session')}"


async def _ws(path: str, cookie: str | None) -> list[dict]:
    """把一次 WebSocket 握手原样交给整个应用（中间件 + 路由级依赖 + 路由），记下发出去的帧。"""
    from app.main import app

    inbox: asyncio.Queue = asyncio.Queue()
    await inbox.put({"type": "websocket.connect"})
    sent: list[dict] = []

    async def receive():
        return await inbox.get()

    async def send(message):
        sent.append(message)
        if message["type"] == "websocket.close":
            await inbox.put({"type": "websocket.disconnect", "code": 1000})

    headers = [(b"host", b"test")] + ([(b"cookie", cookie.encode())] if cookie else [])
    scope = {"type": "websocket", "asgi": {"version": "3.0"}, "scheme": "ws", "http_version": "1.1", "path": path,
             "raw_path": path.encode(), "query_string": b"", "root_path": "", "headers": headers,
             "client": ("127.0.0.1", 50000), "server": ("test", 80), "subprotocols": [], "state": {}}
    await asyncio.wait_for(app(scope, receive, send), 10)
    return sent


async def _settled_run(**fields) -> str:
    async with SessionLocal() as s:
        run = Run(workflow_name="问数据", graph={}, input={}, status="succeeded", run_class="exploratory", **fields)
        s.add(run)
        await s.flush()
        s.add(RunEvent(run_id=run.id, seq=1, type="run.started", ts=1.0, data={}))
        await s.commit()
        return run.id


def _refused(frames: list[dict]) -> bool:
    return [f["type"] for f in frames] == ["websocket.close"] and frames[0].get("code") == 1008


async def test_websocket_is_judged_per_run_at_the_handshake(world):
    mine = await _settled_run(started_by="uma", origin_kind="chat", options={MASK_RESULTS_OPTION: True})
    others = await _settled_run(started_by="bob", origin_kind="chat", options={MASK_RESULTS_OPTION: True})
    raw = await _settled_run(started_by="uma", origin_kind="chat")      # 关登录时自报同名留下的，没在边界遮过
    uma, bob = _cookie(world["uma"]), _cookie(world["bob"])

    assert _refused(await _ws(f"/api/runs/{mine}/stream", None)), "没登录的握手前关掉"
    assert _refused(await _ws(f"/api/runs/{mine}/stream", uma)), "没开问数据时使用者连不了运行"
    await _allow_chat(world)
    frames = await _ws(f"/api/runs/{mine}/stream", uma)
    assert frames[0]["type"] == "websocket.accept" and '"stream.end"' in frames[-2]["text"]
    assert _refused(await _ws(f"/api/runs/{others}/stream", uma)), "别人的运行"
    assert _refused(await _ws(f"/api/runs/{raw}/stream", uma)), "查询结果没在工具边界遮过的运行"
    assert _refused(await _ws("/api/runs/nope/stream", uma))
    assert (await _ws(f"/api/runs/{raw}/stream", bob))[0]["type"] == "websocket.accept", "搭建者照旧"


# ---------------------------------------------------------------- 数据源名单


async def test_source_names_for_the_chat_page(world):
    uma, bob = world["uma"], world["bob"]
    assert (await uma.get("/api/datasources/names")).status_code == 403, "没开问数据时不给使用者"
    assert (await uma.get("/api/datasources")).status_code == 403
    await _allow_chat(world)
    body = (await uma.get("/api/datasources/names")).json()
    assert [i["name"] for i in body["items"]] == [world["source"]], "只列启用着的"
    assert set(body["items"][0]) == {"id", "name", "description", "kind"}, "连接信息不在名单里"
    assert body["items"][0]["description"] == "客户" and body["items"][0]["kind"] == "sqlite"
    assert (await uma.get("/api/datasources")).status_code == 403, "完整的数据源列表照旧只给搭建者"
    assert (await bob.get("/api/datasources/names")).status_code == 200


async def test_source_names_with_login_off():
    await auth_env.reset()
    async with auth_env.client() as c:
        r = await c.get("/api/datasources/names")
    assert r.status_code == 200 and "items" in r.json()