app/tools/graph_tool.py

from __future__ import annotations

import json
from typing import Any, Literal, Optional

from langchain_core.tools import tool

from app.db.neo4j import run_read_query

QueryType = Literal[
    "entity_connections",
    "entity_mentions",
    "shared_entities_between_articles",
    "most_connected_entities",
    "article_verdict",
    "claims_for_article",
    "analysis_for_article",
]


async def _entity_connections(entity_name, limit):
    query = """
        MATCH (e:Entity {name: $entity_name})-[r]-(other)
        RETURN other, type(r) AS relationship, labels(other) AS other_labels
        LIMIT $limit
    """
    rows = await run_read_query(query, entity_name=entity_name, limit=limit)

    refs = []
    for row in rows:
        other = row.get("other", {}) or {}
        if "entity_id" in other:
            refs.append({"type": "entity", "id": other["entity_id"]})
        elif other.get("article_id") is not None:
            refs.append({"type": "article", "id": other["article_id"]})
    return rows, refs


async def _entity_mentions(entity_name, limit):
    query = """
        MATCH (a:Article)-[r:MENTIONS]->(e:Entity {name: $entity_name})
        RETURN a.article_id AS article_id, r.confidence AS confidence
        ORDER BY r.confidence DESC
        LIMIT $limit
    """
    rows = await run_read_query(query, entity_name=entity_name, limit=limit)
    refs = [{"type": "article", "id": r["article_id"]} for r in rows if r.get("article_id") is not None]
    return rows, refs


async def _shared_entities_between_articles(article_id, article_id_2, limit):
    query = """
        MATCH (a1:Article {article_id: $article_id})-[:MENTIONS]->(e:Entity)
              <-[:MENTIONS]-(a2:Article {article_id: $article_id_2})
        RETURN e.name AS name, e.type AS type, e.entity_id AS entity_id
        LIMIT $limit
    """
    rows = await run_read_query(query, article_id=article_id, article_id_2=article_id_2, limit=limit)
    refs = [{"type": "article", "id": article_id}, {"type": "article", "id": article_id_2}]
    refs += [{"type": "entity", "id": r["entity_id"]} for r in rows if r.get("entity_id")]
    return rows, refs


async def _most_connected_entities(limit):
    query = """
        MATCH (e:Entity)<-[:MENTIONS]-(a:Article)
        RETURN e.name AS name, e.type AS type, e.entity_id AS entity_id,
               count(a) AS mention_count
        ORDER BY mention_count DESC
        LIMIT $limit
    """
    rows = await run_read_query(query, limit=limit)
    refs = [{"type": "entity", "id": r["entity_id"]} for r in rows if r.get("entity_id")]
    return rows, refs


async def _article_verdict(article_id):
    query = """
        MATCH (a:Article {article_id: $article_id})-[:HAS_FINAL_VERDICT]->(v:Verdict)
        OPTIONAL MATCH (v)-[:BASED_ON]->(e:Evidence)
        RETURN v.verdict AS verdict, v.confidence AS confidence,
               v.explanation AS explanation, collect(e) AS evidence
    """
    rows = await run_read_query(query, article_id=article_id)
    return rows, [{"type": "article", "id": article_id}]


async def _claims_for_article(article_id):
    query = """
        MATCH (a:Article {article_id: $article_id})-[:PROPOSES]->(c:Claim)
              -[:EVALUATED_BY]->(v:Verdict)
        RETURN c.text AS claim_text, c.claim_idx AS claim_idx,
               v.verdict AS verdict, v.confidence AS confidence,
               v.explanation AS explanation
        ORDER BY c.claim_idx
    """
    rows = await run_read_query(query, article_id=article_id)
    return rows, [{"type": "article", "id": article_id}]


async def _analysis_for_article(article_id):
    query = """
        MATCH (a:Article {article_id: $article_id})-[r]->(an:Analysis)
        RETURN type(r) AS relationship, an.analysis_type AS analysis_type,
               an.value AS value, an.score AS score
    """
    rows = await run_read_query(query, article_id=article_id)
    return rows, [{"type": "article", "id": article_id}]


@tool
async def graph_query_tool(
    query_type: QueryType,
    entity_name: Optional[str] = None,
    article_id: Optional[int] = None,
    article_id_2: Optional[int] = None,
    limit: int = 20,
) -> str:
    """
    Run a pre-defined, read-only Cypher query over the Neo4j entity/article
    relationship graph.

    query_type: which relationship query to run —
        - entity_connections: everything connected to entity_name, any
          relationship direction (mentions + co-mentions)
        - entity_mentions: which articles mention entity_name
        - shared_entities_between_articles: entities mentioned in both
          article_id and article_id_2
        - most_connected_entities: global ranking of most-mentioned entities
          (no entity_name needed)
        - article_verdict: final fake-news verdict + evidence for article_id
        - claims_for_article: per-claim verdict breakdown for article_id
        - analysis_for_article: full sentiment/topic/dialect/propaganda/
          hate-speech profile for article_id

    entity_name: match by Entity.name (NOT internal entity_id). Required for
        entity_connections, entity_mentions, shared_entities_between_articles.
    article_id / article_id_2: required for article_verdict, claims_for_article,
        analysis_for_article, and shared_entities_between_articles (needs both).
    limit: max rows to return, ignored for single-article lookups.

    Returns a JSON string: {"rows": [...], "source_refs": [{"type": ...,
    "id": ...}, ...]}
    """
    try:
        if query_type == "entity_connections":
            if not entity_name:
                raise ValueError("entity_connections requires entity_name")
            rows, refs = await _entity_connections(entity_name, limit)
        elif query_type == "entity_mentions":
            if not entity_name:
                raise ValueError("entity_mentions requires entity_name")
            rows, refs = await _entity_mentions(entity_name, limit)
        elif query_type == "shared_entities_between_articles":
            if article_id is None or article_id_2 is None:
                raise ValueError("shared_entities_between_articles requires article_id and article_id_2")
            rows, refs = await _shared_entities_between_articles(article_id, article_id_2, limit)
        elif query_type == "most_connected_entities":
            rows, refs = await _most_connected_entities(limit)
        elif query_type == "article_verdict":
            if article_id is None:
                raise ValueError("article_verdict requires article_id")
            rows, refs = await _article_verdict(article_id)
        elif query_type == "claims_for_article":
            if article_id is None:
                raise ValueError("claims_for_article requires article_id")
            rows, refs = await _claims_for_article(article_id)
        elif query_type == "analysis_for_article":
            if article_id is None:
                raise ValueError("analysis_for_article requires article_id")
            rows, refs = await _analysis_for_article(article_id)
        else:
            return json.dumps({"rows": [], "source_refs": [], "error": f"Unknown query_type: {query_type}"})
    except ValueError as exc:
        return json.dumps({"rows": [], "source_refs": [], "error": str(exc)})

    return json.dumps({"rows": rows, "source_refs": refs}, default=str)