"""
Jina 流式集成修复版测试：
- 修复 in_think 追踪（基于 regex 而非逐字符）
- 打印原始 annotation 格式（诊断 citations=0）
- 验证 @tool LangChain 兼容性
"""

import json, os, re, sys, time
import requests

JINA_API_KEY = os.environ["JINA_API_KEY"]
SEP = "=" * 60


# ─── Step 1: 修复后的 SSE 流式测试 ───────────────────────────
def test_deepsearch_fixed():
    print(f"\n{SEP}")
    print("STEP 1 — DeepSearch SSE 修复版（in_think 追踪 + annotation 诊断）")
    print(SEP)

    query = "What is Jina AI DeepSearch and how does it differ from RAG? Keep it brief."
    print(f"Query: {query}\n")

    headers = {
        "Authorization": f"Bearer {JINA_API_KEY}",
        "Content-Type": "application/json",
        "Accept": "text/event-stream",
    }
    payload = {
        "model": "jina-deepsearch-v1",
        "messages": [{"role": "user", "content": query}],
        "stream": True,
        "reasoning_effort": "low",
        "no_direct_answer": True,
    }

    full_content = ""
    raw_annotations = []
    chunk_count = 0
    t0 = time.time()

    print("── 流式接收中 ──")
    with requests.post(
        "https://deepsearch.jina.ai/v1/chat/completions",
        headers=headers, json=payload, stream=True, timeout=180,
    ) as resp:
        resp.raise_for_status()
        content_type = resp.headers.get("content-type", "?")
        print(f"HTTP {resp.status_code}  Content-Type: {content_type}")

        first_chunk_logged = False
        for raw_line in resp.iter_lines():
            if not raw_line:
                continue
            line = raw_line.decode("utf-8") if isinstance(raw_line, bytes) else raw_line
            if not line.startswith("data: "):
                continue
            data_str = line[6:].strip()
            if data_str == "[DONE]":
                break
            try:
                chunk = json.loads(data_str)
            except json.JSONDecodeError:
                continue

            chunk_count += 1

            # 打印第一个 chunk 的完整结构（诊断）
            if not first_chunk_logged:
                print(f"\n[Chunk #1 raw keys]: {list(chunk.keys())}")
                choice0 = (chunk.get("choices") or [{}])[0]
                print(f"[Choice keys]: {list(choice0.keys())}")
                delta0 = choice0.get("delta", {})
                print(f"[Delta keys]: {list(delta0.keys())}")
                first_chunk_logged = True

            choice = (chunk.get("choices") or [{}])[0]
            delta = choice.get("delta", {})
            piece = delta.get("content") or ""
            if piece:
                full_content += piece
                sys.stdout.write("." if len(full_content) % 200 != 0 else f"[{len(full_content)}]")
                sys.stdout.flush()

            # 收集 annotations —— 尝试多个字段名
            for ann in delta.get("annotations") or []:
                raw_annotations.append(ann)
            # 也检查顶层
            for ann in chunk.get("annotations") or []:
                raw_annotations.append(ann)

    elapsed = time.time() - t0
    print(f"\n\n── 统计 ──")
    print(f"  Chunks: {chunk_count}   Elapsed: {elapsed:.1f}s")
    print(f"  Raw content length: {len(full_content)} chars")
    print(f"  Raw annotations: {len(raw_annotations)}")

    # 修复后的 think 追踪（用 regex 计算）
    think_matches = re.findall(r"<think>(.*?)</think>", full_content, re.DOTALL)
    think_chars = sum(len(m) for m in think_matches)
    answer = re.sub(r"<think>.*?</think>", "", full_content, flags=re.DOTALL).strip()
    print(f"  Think blocks: {len(think_matches)}  Think chars: {think_chars}")
    print(f"  Answer chars: {len(answer)}")

    print(f"\n── 最终答案（前 600 字符）──")
    print(answer[:600])

    # 诊断 annotations
    if raw_annotations:
        print(f"\n── Annotations 样例 ──")
        print(json.dumps(raw_annotations[0], indent=2, ensure_ascii=False)[:500])
    else:
        print("\n── Annotations: 0（low effort 可能不返回引用，见 Step 2）")

    return answer


# ─── Step 2: medium effort 测试引用 ──────────────────────────
def test_deepsearch_citations():
    print(f"\n{SEP}")
    print("STEP 2 — DeepSearch citations 验证 (reasoning_effort=medium)")
    print(SEP)

    query = "Jina AI Reader API key features and pricing 2025"
    headers = {
        "Authorization": f"Bearer {JINA_API_KEY}",
        "Content-Type": "application/json",
        "Accept": "text/event-stream",
    }
    payload = {
        "model": "jina-deepsearch-v1",
        "messages": [{"role": "user", "content": query}],
        "stream": True,
        "reasoning_effort": "medium",
        "no_direct_answer": True,
    }

    full_content = ""
    citations = []
    chunk_count = 0
    t0 = time.time()
    annotation_chunk_idx = None

    with requests.post(
        "https://deepsearch.jina.ai/v1/chat/completions",
        headers=headers, json=payload, stream=True, timeout=180,
    ) as resp:
        resp.raise_for_status()
        for raw_line in resp.iter_lines():
            if not raw_line:
                continue
            line = raw_line.decode("utf-8") if isinstance(raw_line, bytes) else raw_line
            if not line.startswith("data: "):
                continue
            data_str = line[6:].strip()
            if data_str == "[DONE]":
                break
            try:
                chunk = json.loads(data_str)
            except json.JSONDecodeError:
                continue
            chunk_count += 1
            choice = (chunk.get("choices") or [{}])[0]
            delta = choice.get("delta", {})
            piece = delta.get("content") or ""
            if piece:
                full_content += piece
                sys.stdout.write(".")
                sys.stdout.flush()

            anns = delta.get("annotations") or chunk.get("annotations") or []
            if anns and annotation_chunk_idx is None:
                annotation_chunk_idx = chunk_count
            for ann in anns:
                if isinstance(ann, dict):
                    uc = ann.get("url_citation") or ann
                    url = uc.get("url", "")
                    if url and url not in [c.get("url") for c in citations]:
                        citations.append(uc)

    elapsed = time.time() - t0
    answer = re.sub(r"<think>.*?</think>", "", full_content, flags=re.DOTALL).strip()
    print(f"\n\nChunks: {chunk_count}  Elapsed: {elapsed:.1f}s  Citations: {len(citations)}")
    print(f"Annotations first appeared at chunk #{annotation_chunk_idx}")
    print(f"\n── Answer (first 500 chars) ──\n{answer[:500]}")
    if citations:
        print(f"\n── Citations ──")
        for i, c in enumerate(citations[:5], 1):
            print(f"  [{i}] {c.get('title','?')[:55]} — {c.get('url','?')[:65]}")
    return citations


# ─── Step 3: @tool LangChain 兼容性 ─────────────────────────
def test_tool_invoke():
    print(f"\n{SEP}")
    print("STEP 3 — @tool LangChain .invoke() 兼容性")
    print(SEP)

    sys.path.insert(0, "/home/yzq/tmp/tech/deepagents")
    from jina_tools import jina_search, jina_reader, jina_deepsearch

    # jina_search
    print("► jina_search.invoke()…")
    t0 = time.time()
    r = jina_search.invoke({"query": "DeepSearch vs RAG LLM 2025", "num_results": 2})
    print(f"  ✓ {time.time()-t0:.1f}s — {len(r)} chars")
    print(f"  {r[:200].replace(chr(10),' ')}\n")

    # jina_reader
    print("► jina_reader.invoke()…")
    t0 = time.time()
    r = jina_reader.invoke({"url": "https://jina.ai/deepsearch", "max_chars": 500})
    print(f"  ✓ {time.time()-t0:.1f}s — {len(r)} chars")
    print(f"  {r[:200].replace(chr(10),' ')}\n")

    # jina_deepsearch
    print("► jina_deepsearch.invoke() [reasoning_effort=low]…")
    t0 = time.time()
    r = jina_deepsearch.invoke({
        "query": "Key differences between Jina Reader and Firecrawl",
        "reasoning_effort": "low",
    })
    print(f"  ✓ {time.time()-t0:.1f}s — {len(r)} chars")
    print(f"\n── Tool output (first 700 chars) ──")
    print(r[:700])


if __name__ == "__main__":
    test_deepsearch_fixed()
    test_deepsearch_citations()
    test_tool_invoke()
    print(f"\n{SEP}\nAll tests done.\n{SEP}")