from __future__ import annotations

from datetime import date

import pytest

from deerflow.integrations.weknora.local_index.advanced import (
    advanced_search,
    build_statistics,
    enrich_chunk_references,
    expand_related_terms,
    extract_subject_terms,
    parse_date_range,
)
from deerflow.persistence.llmwiki_index import MemoryLlmWikiIndexStore


def _page(*, title: str, slug: str, content: str, published_at: str, source_refs=None) -> dict:
    return {
        "id": slug,
        "knowledge_base_mapping_id": "kb",
        "slug": slug,
        "title": title,
        "summary": "",
        "content_md": f"---\npublished_at: {published_at}\n---\n{content}",
        "aliases": [],
        "page_type": "article",
        "status": "published",
        "is_remote_deleted": False,
        "source_refs": source_refs or [],
    }


def test_parses_chinese_date_ranges_and_subject() -> None:
    start, end = parse_date_range(
        "帮我总结9月1号到9月15号的报告，有多少是关于赖清德的？",
        today=date(2026, 9, 20),
    )
    assert start == date(2026, 9, 1)
    assert end == date(2026, 9, 15)
    assert extract_subject_terms("有多少是关于赖清德的？") == ["赖清德"]


def test_expands_first_degree_wiki_relations() -> None:
    page = _page(
        title="赖清德",
        slug="people/lai",
        published_at="2026-09-03",
        content="赖清德与[[民进党]]、[[台湾行政机构|行政机构]]有关。",
    )
    expanded = expand_related_terms(["赖清德"], [page], limit=5)
    assert [item["term"] for item in expanded] == ["民进党", "行政机构"]
    assert all(item["from_page"] == "赖清德" for item in expanded)


def test_statistics_prefers_source_documents_and_applies_exact_dates() -> None:
    pages = [
        _page(
            title="赖清德相关新闻汇总",
            slug="news/lai",
            published_at="2026-09-05",
            content="赖清德出席活动。",
            source_refs=[
                {"document_id": "n1", "title": "新闻一", "published_at": "2026-09-02", "url": "https://example.test/1"},
                {"document_id": "n2", "title": "新闻二", "published_at": "2026-09-18", "url": "https://example.test/2"},
                {"document_id": "n3", "title": "新闻三"},
            ],
        )
    ]
    payload = build_statistics(
        pages,
        [{"id": "kb", "name": "台湾新闻"}],
        subject_terms=["赖清德"],
        date_start=date(2026, 9, 1),
        date_end=date(2026, 9, 15),
        requested=True,
    )
    assert payload["unit"] == "source_document"
    assert payload["matched_count"] == 1
    assert payload["items"][0]["title"] == "新闻一"
    assert payload["undated_excluded_count"] == 1


class _ChunkClient:
    async def get_chunk(self, chunk_id: str):
        return {
            "id": chunk_id,
            "knowledge_base_id": "remote-kb",
            "knowledge_id": "news-document-1",
            "knowledge_title": "赖清德活动新闻",
            "content": "发布日期：2026年9月6日\n赖清德出席活动。",
            "metadata": {"url": "https://example.test/news"},
        }


@pytest.mark.asyncio
async def test_resolves_only_wiki_cited_chunks_to_news_sources() -> None:
    page = _page(
        title="赖清德",
        slug="people/lai",
        published_at="2026-09-03",
        content="赖清德活动",
        source_refs=["chunk-1"],
    )
    pages, warnings = await enrich_chunk_references(
        [page],
        [{"id": "kb", "weknora_id": "remote-kb"}],
        client=_ChunkClient(),
    )
    assert warnings == []
    resolved = next(item for item in pages[0]["source_refs"] if isinstance(item, dict))
    assert resolved["document_id"] == "news-document-1"
    assert resolved["title"] == "赖清德活动新闻"


class _Retrieval:
    async def search(self, query, mappings, *, top_k_pages, include_drafts):
        return {
            "mode": "hybrid",
            "results": [
                {
                    "id": "p1",
                    "wiki_page_id": "p1",
                    "wiki_slug": "news/lai",
                    "knowledge_base_id": "kb",
                    "knowledge_base_name": "台湾新闻",
                    "title": "赖清德相关新闻",
                    "content": "赖清德与民进党",
                    "snippet": "赖清德与民进党",
                    "score": 0.8,
                    "url": "/preview",
                    "matched_sections": [],
                }
            ],
            "warnings": [],
            "diagnostics": {"query_seen": query},
        }


@pytest.mark.asyncio
async def test_advanced_search_returns_plan_counts_and_references() -> None:
    store = MemoryLlmWikiIndexStore()
    page = _page(
        title="赖清德相关新闻",
        slug="news/lai",
        published_at="2026-09-03",
        content="赖清德与[[民进党]]",
    )
    page.update({"slug_hash": "hash", "content_hash": "content"})
    await store.replace_page_vectors(page=page, vectors=[], fingerprint="fp", sync_id="sync")
    payload = await advanced_search(
        query="9月1号到9月15号有多少关于赖清德的新闻？",
        mappings=[{"id": "kb", "name": "台湾新闻", "description": "每日新闻", "retrieval_settings": None}],
        retrieval_service=_Retrieval(),
        index_store=store,
        top_k=20,
    )
    assert payload["query_plan"]["subject_terms"] == ["赖清德"]
    assert payload["query_plan"]["expanded_terms"][0]["term"] == "民进党"
    assert payload["statistics"]["matched_count"] == 1
    assert payload["references"][0]["reference_id"] == "wiki:kb:news/lai"