import asyncio
import json
from agentui.pipeline.executor import PipelineExecutor

def run_checks():
    async def main():
        print("\n=== JP FUNCTIONS: SetVars expr → jp/from_json/jp_text ===")
        # 1) from_json + jp: извлечь число по пути
        p1 = {
            "id": "p_jp_1",
            "name": "JP basic",
            "loop_mode": "dag",
            "nodes": [
                {
                    "id": "n1",
                    "type": "SetVars",
                    "config": {
                        "variables": [
                            {
                                "id": "v1",
                                "name": "VAL",
                                "mode": "expr",
                                # JSON-строка → объект → jp по пути "a.b.1.x" -> 2
                                "value": "jp(from_json('{\"a\":{\"b\":[{\"x\":1},{\"x\":2}]}}'), 'a.b.1.x')"
                            }
                        ]
                    },
                    "in": {}
                },
                {
                    "id": "n2",
                    "type": "Return",
                    "config": {
                        "target_format": "openai",
                        # Должно быть число 2, выводим как текст
                        "text_template": "{{ VAL }}"
                    },
                    "in": { "depends": "n1.done" }
                }
            ]
        }
        ctx = {
            "model": "gpt-x",
            "vendor_format": "openai",
            "params": {},
            "chat": {"last_user": "hi"},
            "OUT": {}
        }
        out1 = await PipelineExecutor(p1).run(ctx)
        msg1 = out1["result"]["choices"][0]["message"]["content"]
        print("OUT1:", msg1)
        assert msg1 == "2"

        # 2) jp_text: собрать строки из массива
        p2 = {
            "id": "p_jp_2",
            "name": "JP text join",
            "loop_mode": "dag",
            "nodes": [
                {
                    "id": "n1",
                    "type": "SetVars",
                    "config": {
                        "variables": [
                            {
                                "id": "v1",
                                "name": "TXT",
                                "mode": "expr",
                                "value": "jp_text(from_json('{\"items\":[{\"t\":\"A\"},{\"t\":\"B\"},{\"t\":\"C\"}]}'), 'items.*.t', ' | ')"
                            }
                        ]
                    },
                    "in": {}
                },
                {
                    "id": "n2",
                    "type": "Return",
                    "config": {
                        "target_format": "openai",
                        "text_template": "[[TXT]]"
                    },
                    "in": { "depends": "n1.done" }
                }
            ]
        }
        out2 = await PipelineExecutor(p2).run(ctx)
        msg2 = out2["result"]["choices"][0]["message"]["content"]
        print("OUT2:", msg2)
        assert msg2 == "A | B | C"

        print("JP functions tests: OK")

    asyncio.run(main())

if __name__ == "__main__":
    run_checks()
    print("Tests completed")