"""
Jina AI Tools for DeepAgents
=============================
Tool 1: jina_search  — s.jina.ai  语义网络搜索（返回 top-N 结构化结果）
Tool 2: jina_reader  — r.jina.ai  单 URL 全文提取（返回 Markdown）
Tool 3: jina_deepsearch — deepsearch.jina.ai  深度研究（流式，返回带引用的文章）

SubAgent: jina-writer  — deepagents 兼容的文章撰写子 agent
"""

import asyncio
import json
import os
import re
from typing import Annotated

import httpx
import requests
from langchain_core.tools import tool
from langgraph.config import get_stream_writer

JINA_API_KEY = os.getenv("JINA_API_KEY", "")
_SEARCH_URL = "https://s.jina.ai/"
_READER_URL = "https://r.jina.ai/"
_DEEPSEARCH_URL = "https://deepsearch.jina.ai/v1/chat/completions"


# ──────────────────────────────────────────────────────────────────────────────
# Tool 1 — jina_search
# ──────────────────────────────────────────────────────────────────────────────

@tool
def jina_search(
    query: Annotated[str, "Search query — describe what you want to find"],
    num_results: Annotated[int, "Number of results to return (1-10, default 5)"] = 5,
) -> str:
    """Search the web using Jina AI (s.jina.ai).

    Returns top results as LLM-friendly markdown with title, URL, and content.
    Use this for: current events, factual lookup, finding primary sources.
    Each result is capped at 2000 chars to stay within context limits.
    """
    headers = {
        "Authorization": f"Bearer {JINA_API_KEY}",
        "Accept": "application/json",
        "X-With-Links-Summary": "true",
        "X-Return-Format": "markdown",
    }
    payload = {
        "q": query,
        "num": max(1, min(10, num_results)),
    }

    try:
        resp = requests.post(
            _SEARCH_URL,
            headers=headers,
            json=payload,
            timeout=30,
        )
        resp.raise_for_status()
        data = resp.json()
    except requests.RequestException as e:
        return f"[jina_search error] {e}"

    items = data.get("data", [])
    if not items:
        return "No results found."

    parts = []
    for i, item in enumerate(items, 1):
        title = item.get("title", "Untitled")
        url = item.get("url", "")
        content = (item.get("content") or item.get("description") or "").strip()
        content = content[:2000] + ("…" if len(content) > 2000 else "")
        parts.append(f"[{i}] **{title}**\nURL: {url}\n\n{content}")

    return "\n\n---\n\n".join(parts)


# ──────────────────────────────────────────────────────────────────────────────
# Tool 2 — jina_reader
# ──────────────────────────────────────────────────────────────────────────────

@tool
def jina_reader(
    url: Annotated[str, "Full URL of the page to read (must start with http/https)"],
    max_chars: Annotated[int, "Max characters to return (default 6000)"] = 6000,
) -> str:
    """Read and extract clean Markdown content from a URL using Jina AI (r.jina.ai).

    Use this for: reading a specific article, report, or document in full;
    following up on a search result; extracting structured content from a page.
    Returns the page title + full markdown body, capped at max_chars.
    """
    if not url.startswith("http"):
        return "[jina_reader error] URL must start with http:// or https://"

    headers = {
        "Authorization": f"Bearer {JINA_API_KEY}",
        "Accept": "application/json",
        "X-Return-Format": "markdown",
        "X-With-Links-Summary": "false",
    }

    try:
        resp = requests.get(
            f"{_READER_URL}{url}",
            headers=headers,
            timeout=30,
        )
        resp.raise_for_status()
        data = resp.json()
    except requests.RequestException as e:
        return f"[jina_reader error] {e}"

    payload = data.get("data", {})
    title = payload.get("title", url)
    content = (payload.get("content") or "").strip()

    if not content:
        return f"[jina_reader] No content extracted from {url}"

    truncated = content[:max_chars] + ("…[truncated]" if len(content) > max_chars else "")
    return f"# {title}\nSource: {url}\n\n{truncated}"


# ──────────────────────────────────────────────────────────────────────────────
# Tool 3 — jina_deepsearch
# ──────────────────────────────────────────────────────────────────────────────

def _strip_thinking(text: str) -> str:
    """Remove <think>...</think> reasoning blocks from DeepSearch output."""
    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()


def _dedupe_citations(citations: list[dict]) -> list[dict]:
    seen, unique = set(), []
    for c in citations:
        url = c.get("url", "")
        if url and url not in seen:
            seen.add(url)
            unique.append(c)
    return unique


@tool
def jina_deepsearch(
    query: Annotated[str, "Research question or article topic to investigate"],
    reasoning_effort: Annotated[
        str,
        "Depth of research: 'low' (fast, ~30s), 'medium' (balanced, ~1min), 'high' (thorough, ~3min)",
    ] = "medium",
) -> str:
    """Conduct deep web research using Jina AI DeepSearch.

    Autonomously plans search strategies, reads multiple sources, reasons over
    evidence, and returns a well-cited research synthesis or article.

    Use this for:
    - Writing comprehensive articles or reports
    - Multi-source research requiring synthesis
    - Questions needing up-to-date information with citations

    Returns the full research result with inline citations and a Sources section.
    NOTE: This tool is slow (30s – 3min). Use reasoning_effort='low' for quick lookups.
    """
    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": reasoning_effort,
        "no_direct_answer": True,  # always perform web search, never answer from memory
    }

    full_content = ""
    citations: list[dict] = []

    try:
        with requests.post(
            _DEEPSEARCH_URL,
            headers=headers,
            json=payload,
            stream=True,
            timeout=300,  # 5-minute timeout for high-effort research
        ) 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)
                    choice = (chunk.get("choices") or [{}])[0]
                    delta = choice.get("delta", {})

                    # accumulate content
                    content_piece = delta.get("content") or ""
                    if content_piece:
                        full_content += content_piece

                    # collect citations — arrive in final chunk, check delta + chunk level
                    # ann structure: {"type":"url_citation","url_citation":{...}} OR flat dict
                    for ann in (delta.get("annotations") or chunk.get("annotations") or []):
                        if not isinstance(ann, dict):
                            continue
                        if ann.get("type") == "url_citation":
                            uc = ann.get("url_citation") or ann
                            if uc.get("url"):
                                citations.append(uc)

                except (json.JSONDecodeError, KeyError):
                    continue

    except requests.RequestException as e:
        return f"[jina_deepsearch error] {e}"

    if not full_content:
        return "[jina_deepsearch] No content returned."

    # Strip internal reasoning blocks
    answer = _strip_thinking(full_content)

    # Append sources section if not already present
    unique_citations = _dedupe_citations(citations)
    has_sources = re.search(r"^## (Sources|References|参考)", answer, re.MULTILINE)
    if unique_citations and not has_sources:
        answer += "\n\n## Sources\n"
        for i, c in enumerate(unique_citations, 1):
            title = c.get("title") or c.get("url", "")
            url = c.get("url", "")
            answer += f"[{i}] {title} — {url}\n"

    return answer


# ──────────────────────────────────────────────────────────────────────────────
# SubAgent: jina-writer
# deepagents-compatible subagent for article writing via Jina DeepSearch
# ──────────────────────────────────────────────────────────────────────────────

JINA_WRITER_INSTRUCTIONS = f"""You are an expert research writer powered by Jina AI.

## Role
Given a writing request (article, report, analysis, summary), you:
1. Research the topic thoroughly using jina_deepsearch
2. Optionally deepen specific angles using jina_search + jina_reader
3. Synthesize all findings into a well-structured, citation-backed piece

## Tools Available
- `jina_deepsearch(query, reasoning_effort)` — primary: deep multi-source research with citations
- `jina_search(query, num_results)` — supplement: targeted keyword searches
- `jina_reader(url)` — supplement: read a specific source in full

## Writing Process

### Step 1 — Understand the request
Identify: topic scope, target length, required sections, target audience, tone (technical / general).

### Step 2 — Research
- Call `jina_deepsearch` with the core topic
  - Use `reasoning_effort="high"` for long-form articles or reports (>1000 words)
  - Use `reasoning_effort="medium"` for standard pieces
  - Use `reasoning_effort="low"` for quick summaries
- If specific sub-topics need more depth, call `jina_search` then `jina_reader` on key URLs

### Step 3 — Write the article
Synthesize all research into a coherent piece. Structure:
```
# [Title]

## Introduction / Executive Summary

## [Section 1]
[content with inline citations [1][2]]

## [Section 2 … N]

## Conclusion / Key Takeaways

## References
[1] Title — URL
[2] ...
```

## Citation Rules
- Every factual claim must carry an inline citation: [1], [2], etc.
- Consolidate all URLs into a References section at the end
- Do not fabricate sources — only cite what jina_deepsearch or jina_reader actually returned

## Output
Return the complete article in Markdown. Do NOT return a raw dump of DeepSearch output —
write a proper article that synthesizes the research.
"""

jina_writer_subagent = {
    "name": "jina-writer",
    "description": (
        "Writes comprehensive, citation-backed articles and research reports "
        "using Jina AI DeepSearch for up-to-date web research. "
        "Use for: long-form articles, market research, topic overviews, "
        "technical summaries requiring current sources. "
        "Pass: article topic, desired length/sections, and target audience."
    ),
    "system_prompt": JINA_WRITER_INSTRUCTIONS,
    "tools": [jina_deepsearch, jina_search, jina_reader],
}


# ──────────────────────────────────────────────────────────────────────────────
# Tool 4 — jina_deepsearch_stream  (async, real-time streaming via get_stream_writer)
# ──────────────────────────────────────────────────────────────────────────────

@tool
async def jina_deepsearch_stream(
    query: Annotated[str, "Research question or article topic to investigate"],
    reasoning_effort: Annotated[
        str,
        "Depth: 'low' (~60s), 'medium' (~3min), 'high' (~5min)",
    ] = "low",
) -> str:
    """Deep web research with REAL-TIME streaming output.

    Uses Jina AI DeepSearch and pushes answer tokens to the stream as they
    arrive via get_stream_writer(). The <think> reasoning block is suppressed;
    only the final answer content streams out.

    Use this instead of jina_deepsearch when the caller needs live output.
    """
    write = get_stream_writer()

    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": reasoning_effort,
        "no_direct_answer": True,
    }

    full_content = ""
    citations: list[dict] = []

    # State machine: suppress <think>…</think> block, stream everything after
    thinking_done = False
    think_buf = ""

    try:
        async with httpx.AsyncClient(timeout=300) as client:
            async with client.stream(
                "POST", _DEEPSEARCH_URL, headers=headers, json=payload
            ) as resp:
                resp.raise_for_status()
                async for raw_line in resp.aiter_lines():
                    if not raw_line.startswith("data: "):
                        continue
                    data_str = raw_line[6:].strip()
                    if data_str == "[DONE]":
                        break
                    try:
                        chunk = json.loads(data_str)
                        choice = (chunk.get("choices") or [{}])[0]
                        delta = choice.get("delta", {})
                        piece = delta.get("content") or ""

                        if piece:
                            full_content += piece

                            if not thinking_done:
                                think_buf += piece
                                if "</think>" in think_buf:
                                    thinking_done = True
                                    after = think_buf.split("</think>", 1)[1]
                                    if after:
                                        # 双路输出：stdout 实时可见 + LangGraph custom stream
                                        import sys
                                        sys.stdout.write(after)
                                        sys.stdout.flush()
                                        write(after)
                            else:
                                import sys
                                sys.stdout.write(piece)
                                sys.stdout.flush()
                                write(piece)

                        for ann in (delta.get("annotations") or chunk.get("annotations") or []):
                            if isinstance(ann, dict) and ann.get("type") == "url_citation":
                                uc = ann.get("url_citation") or ann
                                if uc.get("url"):
                                    citations.append(uc)

                    except (json.JSONDecodeError, KeyError, IndexError):
                        continue

    except Exception as e:
        return f"[jina_deepsearch_stream error] {e}"

    if not full_content:
        return "[jina_deepsearch_stream] No content returned."

    answer = _strip_thinking(full_content)
    unique = _dedupe_citations(citations)
    if unique and not re.search(r"^## (Sources|References)", answer, re.MULTILINE):
        answer += "\n\n## Sources\n"
        for i, c in enumerate(unique, 1):
            answer += f"[{i}] {c.get('title', c.get('url',''))} — {c.get('url','')}\n"

    return answer


# jina-writer subagent variant that uses the streaming tool
JINA_WRITER_STREAM_INSTRUCTIONS = JINA_WRITER_INSTRUCTIONS.replace(
    "jina_deepsearch(query, reasoning_effort)",
    "jina_deepsearch_stream(query, reasoning_effort)",
)

jina_writer_stream_subagent = {
    "name": "jina-writer",
    "description": jina_writer_subagent["description"],
    "system_prompt": JINA_WRITER_STREAM_INSTRUCTIONS,
    "tools": [jina_deepsearch_stream, jina_search, jina_reader],
}


# ──────────────────────────────────────────────────────────────────────────────
# Quick usage example
# ──────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    # Test jina_search
    print("=== jina_search ===")
    result = jina_search.invoke({"query": "Jina AI DeepSearch API 2025", "num_results": 3})
    print(result[:500])

    # Test jina_reader
    print("\n=== jina_reader ===")
    result = jina_reader.invoke({"url": "https://jina.ai/deepsearch", "max_chars": 1000})
    print(result[:500])

    # Test jina_deepsearch
    print("\n=== jina_deepsearch ===")
    result = jina_deepsearch.invoke({
        "query": "What are the key differences between RAG and DeepSearch for LLM applications?",
        "reasoning_effort": "low",
    })
    print(result[:1000])