"""
Bundled test datasets for evaluation.

Provides small, synthetic datasets that demonstrate HIEROMEM capabilities
without external dependencies.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
import networkx as nx


@dataclass
class TestDocument:
    """A document in a test dataset."""
    id: str
    content: str
    graph: Optional[nx.DiGraph] = None
    metadata: dict = None

    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}


@dataclass
class TestQuery:
    """A query in a test dataset."""
    query: str
    relevant_doc_ids: list[str]
    answer: Optional[str] = None


@dataclass
class TestDataset:
    """A complete test dataset."""
    name: str
    description: str
    documents: list[TestDocument]
    queries: list[TestQuery]


def create_entity_chain_dataset() -> TestDataset:
    """
    Create a dataset where graph traversal is essential.

    Chain: Alice -> TechCorp -> Seattle
    Query asks about Alice's city (requires 2-hop reasoning).
    """
    docs = []

    # Doc 1: Alice works at TechCorp
    G1 = nx.DiGraph()
    G1.add_edge("Alice", "TechCorp", relation="works_at")
    docs.append(TestDocument(
        id="alice_doc",
        content="Alice is a senior engineer at TechCorp. She has been working there for 5 years.",
        graph=G1,
        metadata={"entity": "Alice"}
    ))

    # Doc 2: TechCorp located in Seattle
    G2 = nx.DiGraph()
    G2.add_edge("TechCorp", "Seattle", relation="located_in")
    docs.append(TestDocument(
        id="seattle_doc",
        content="TechCorp headquarters is in Seattle, Washington. The campus overlooks the bay.",
        graph=G2,
        metadata={"entity": "TechCorp", "has_answer": True}
    ))

    # Add distractors
    distractors = [
        "The weather in Miami is warm and sunny.",
        "Python programming is popular for data science.",
        "Coffee production has increased globally.",
        "Electric vehicles are becoming more common.",
        "The stock market showed volatility this week.",
    ]
    for i, text in enumerate(distractors):
        docs.append(TestDocument(
            id=f"distractor_{i}",
            content=text,
            graph=nx.DiGraph(),
        ))

    queries = [
        TestQuery(
            query="In which city does Alice work?",
            relevant_doc_ids=["alice_doc", "seattle_doc"],
            answer="Seattle"
        ),
    ]

    return TestDataset(
        name="entity_chain",
        description="2-hop entity chain requiring graph traversal",
        documents=docs,
        queries=queries,
    )


def create_multi_hop_dataset() -> TestDataset:
    """
    Create a multi-hop QA dataset.

    Multiple chains:
    - Bob -> Acme Inc -> New York
    - Carol -> Acme Inc -> New York (shared entity)
    - David -> GlobalTech -> London
    """
    docs = []

    # Bob works at Acme
    G1 = nx.DiGraph()
    G1.add_edge("Bob", "Acme Inc", relation="works_at")
    docs.append(TestDocument(
        id="bob_doc",
        content="Bob is the CTO of Acme Inc. He leads the technical team.",
        graph=G1,
    ))

    # Carol also works at Acme
    G2 = nx.DiGraph()
    G2.add_edge("Carol", "Acme Inc", relation="works_at")
    docs.append(TestDocument(
        id="carol_doc",
        content="Carol joined Acme Inc last year as a product manager.",
        graph=G2,
    ))

    # Acme in New York
    G3 = nx.DiGraph()
    G3.add_edge("Acme Inc", "New York", relation="located_in")
    docs.append(TestDocument(
        id="acme_location_doc",
        content="Acme Inc has its main office in downtown New York City.",
        graph=G3,
        metadata={"has_answer": True}
    ))

    # David at GlobalTech
    G4 = nx.DiGraph()
    G4.add_edge("David", "GlobalTech", relation="works_at")
    docs.append(TestDocument(
        id="david_doc",
        content="David is the CEO of GlobalTech, a fintech startup.",
        graph=G4,
    ))

    # GlobalTech in London
    G5 = nx.DiGraph()
    G5.add_edge("GlobalTech", "London", relation="located_in")
    docs.append(TestDocument(
        id="globaltech_location_doc",
        content="GlobalTech is headquartered in London's financial district.",
        graph=G5,
    ))

    queries = [
        TestQuery(
            query="Where does Bob work?",
            relevant_doc_ids=["bob_doc"],
            answer="Acme Inc"
        ),
        TestQuery(
            query="Which city is Bob's company located in?",
            relevant_doc_ids=["bob_doc", "acme_location_doc"],
            answer="New York"
        ),
        TestQuery(
            query="Do Carol and Bob work at the same company?",
            relevant_doc_ids=["bob_doc", "carol_doc"],
            answer="Yes"
        ),
    ]

    return TestDataset(
        name="multi_hop",
        description="Multi-hop reasoning with shared entities",
        documents=docs,
        queries=queries,
    )


def create_semantic_gap_dataset() -> TestDataset:
    """
    Create a dataset where semantic similarity is LOW but graph connection exists.

    This is where graph retrieval should significantly outperform dense.
    Query about "Alice's city" won't match "Seattle campus" semantically.
    """
    docs = []

    # Doc 1: Alice at TechCorp (connection point)
    G1 = nx.DiGraph()
    G1.add_edge("Alice", "TechCorp", relation="works_at")
    docs.append(TestDocument(
        id="alice_job",
        content="Alice is the VP of Engineering at TechCorp. She oversees the ML platform team.",
        graph=G1,
    ))

    # Doc 2: TechCorp campus (semantically different from query about "city")
    G2 = nx.DiGraph()
    G2.add_edge("TechCorp", "Seattle", relation="headquartered_in")
    docs.append(TestDocument(
        id="seattle_campus",
        content="The TechCorp campus features modern architecture with bay views. The building has LEED certification.",
        graph=G2,
        metadata={"has_answer": True}
    ))

    # Semantic distractors (high similarity to "city" query but irrelevant)
    distractors = [
        ("city_guide", "New York City is known for its skyline. Many tourists visit the city annually."),
        ("urban_living", "Living in a city offers convenience and cultural experiences. Urban planning matters."),
        ("metro_report", "City populations have grown 15% worldwide. Metropolitan areas dominate GDP."),
        ("travel_blog", "I visited three cities last month. Each city had unique food culture."),
        ("real_estate", "City apartments are expensive. Downtown rental prices keep rising."),
    ]
    for doc_id, text in distractors:
        docs.append(TestDocument(id=doc_id, content=text, graph=nx.DiGraph()))

    queries = [
        TestQuery(
            query="What city is Alice's workplace located in?",
            relevant_doc_ids=["alice_job", "seattle_campus"],
            answer="Seattle"
        ),
    ]

    return TestDataset(
        name="semantic_gap",
        description="Low semantic similarity, requires graph traversal",
        documents=docs,
        queries=queries,
    )


def create_long_tail_dataset() -> TestDataset:
    """
    Create a dataset with rare/specific entities.

    Tests ability to link specific entities vs relying on common terms.
    """
    docs = []

    # Rare entity chain
    G1 = nx.DiGraph()
    G1.add_edge("Dr. Xanthippe Morrow", "Quantum Dynamics Lab", relation="directs")
    docs.append(TestDocument(
        id="xanthippe_doc",
        content=(
            "Dr. Xanthippe Morrow is a pioneering researcher in quantum coherence. "
            "She directs the Quantum Dynamics Lab."
        ),
        graph=G1,
    ))

    G2 = nx.DiGraph()
    G2.add_edge("Quantum Dynamics Lab", "MIT", relation="part_of")
    docs.append(TestDocument(
        id="qdl_doc",
        content="The Quantum Dynamics Lab is part of MIT's physics department. It has 12 researchers.",
        graph=G2,
        metadata={"has_answer": True}
    ))

    # Common distractors about physics/labs
    distractors = [
        ("physics_news", "Physics research continues to advance. Many labs study quantum mechanics today."),
        ("mit_general", "MIT is a leading research university. It was founded in 1861."),
        ("quantum_101", "Quantum physics describes particle behavior. Researchers study wave functions."),
    ]
    for doc_id, text in distractors:
        docs.append(TestDocument(id=doc_id, content=text, graph=nx.DiGraph()))

    queries = [
        TestQuery(
            query="Which university is Dr. Xanthippe Morrow's lab affiliated with?",
            relevant_doc_ids=["xanthippe_doc", "qdl_doc"],
            answer="MIT"
        ),
    ]

    return TestDataset(
        name="long_tail",
        description="Rare/specific entities requiring exact matching",
        documents=docs,
        queries=queries,
    )


def get_dataset(name: str) -> TestDataset:
    """Get a bundled dataset by name."""
    datasets = {
        "entity_chain": create_entity_chain_dataset,
        "multi_hop": create_multi_hop_dataset,
        "semantic_gap": create_semantic_gap_dataset,
        "long_tail": create_long_tail_dataset,
    }
    if name not in datasets:
        available = ", ".join(datasets.keys())
        raise ValueError(f"Unknown dataset: {name}. Available: {available}")
    return datasets[name]()


def list_datasets() -> list[str]:
    """List available bundled datasets."""
    return ["entity_chain", "multi_hop", "semantic_gap", "long_tail"]
