app/tools/sql_tool.py

from __future__ import annotations

import json
from typing import Any, Literal, Optional

from langchain_core.tools import tool
from psycopg.rows import dict_row

from app.db.postgres import get_readonly_connection

Metric = Literal[
    "fake_news_count",
    "articles_by_topic",
    "sentiment_breakdown",
    "propaganda_count",
    "hate_speech_count",
    "dialect_breakdown",
    "article_lookup",
]


async def _fetch(sql: str, params: list[Any]) -> list[dict[str, Any]]:
    async with get_readonly_connection() as conn:
        async with conn.cursor(row_factory=dict_row) as cur:
            await cur.execute(sql, params)
            return await cur.fetchall()


def _date_filter(where: list[str], params: list[Any], date_col: str,
                 date_from: Optional[str], date_to: Optional[str]) -> None:
    if date_from:
        where.append(f"{date_col} >= %s")
        params.append(date_from)
    if date_to:
        where.append(f"{date_col} <= %s")
        params.append(date_to)


async def _fake_news_count(date_from, date_to, article_id, limit):
    where: list[str] = []
    params: list[Any] = []
    _date_filter(where, params, "analyzed_at", date_from, date_to)
    if article_id is not None:
        where.append("article_id = %s")
        params.append(article_id)
    where_sql = f"WHERE {' AND '.join(where)}" if where else ""

    breakdown_sql = f"""
        SELECT verdict, COUNT(*) AS count
        FROM article_fake_news
        {where_sql}
        GROUP BY verdict
        ORDER BY count DESC
    """
    rows = await _fetch(breakdown_sql, params)

    refs_sql = f"""
        SELECT article_id
        FROM article_fake_news
        {where_sql}
        ORDER BY analyzed_at DESC
        LIMIT %s
    """
    refs = await _fetch(refs_sql, params + [limit])
    return breakdown_sql, rows, refs


async def _articles_by_topic(date_from, date_to, topic, article_id, limit):
    where: list[str] = []
    params: list[Any] = []
    _date_filter(where, params, "analyzed_at", date_from, date_to)
    if topic:
        where.append("(primary_topic = %s OR secondary_topic = %s)")
        params.extend([topic, topic])
    if article_id is not None:
        where.append("article_id = %s")
        params.append(article_id)
    where_sql = f"WHERE {' AND '.join(where)}" if where else ""

    breakdown_sql = f"""
        SELECT primary_topic, COUNT(*) AS count
        FROM article_topic
        {where_sql}
        GROUP BY primary_topic
        ORDER BY count DESC
    """
    rows = await _fetch(breakdown_sql, params)

    refs_sql = f"""
        SELECT article_id
        FROM article_topic
        {where_sql}
        ORDER BY analyzed_at DESC
        LIMIT %s
    """
    refs = await _fetch(refs_sql, params + [limit])
    return breakdown_sql, rows, refs


async def _sentiment_breakdown(date_from, date_to, sentiment, article_id, limit):
    where: list[str] = []
    params: list[Any] = []
    _date_filter(where, params, "analyzed_at", date_from, date_to)
    if sentiment:
        where.append("primary_sentiment = %s")
        params.append(sentiment)
    if article_id is not None:
        where.append("article_id = %s")
        params.append(article_id)
    where_sql = f"WHERE {' AND '.join(where)}" if where else ""

    breakdown_sql = f"""
        SELECT primary_sentiment, COUNT(*) AS count
        FROM article_sentiment
        {where_sql}
        GROUP BY primary_sentiment
        ORDER BY count DESC
    """
    rows = await _fetch(breakdown_sql, params)

    refs_sql = f"""
        SELECT article_id
        FROM article_sentiment
        {where_sql}
        ORDER BY analyzed_at DESC
        LIMIT %s
    """
    refs = await _fetch(refs_sql, params + [limit])
    return breakdown_sql, rows, refs


async def _propaganda_count(date_from, date_to, article_id, limit):
    where: list[str] = []
    params: list[Any] = []
    _date_filter(where, params, "analyzed_at", date_from, date_to)
    if article_id is not None:
        where.append("article_id = %s")
        params.append(article_id)
    where_sql = f"WHERE {' AND '.join(where)}" if where else ""
    propaganda_where = where + ["is_propaganda = true"]

    count_sql = f"""
        SELECT COUNT(*) AS count
        FROM article_propaganda
        WHERE {' AND '.join(propaganda_where)}
    """
    rows = await _fetch(count_sql, params)

    technique_sql = f"""
        SELECT technique, COUNT(*) AS count
        FROM article_propaganda, unnest(techniques) AS technique
        WHERE {' AND '.join(propaganda_where)}
        GROUP BY technique
        ORDER BY count DESC
        LIMIT 10
    """
    technique_rows = await _fetch(technique_sql, params)
    rows[0]["technique_breakdown"] = technique_rows

    refs_sql = f"""
        SELECT article_id
        FROM article_propaganda
        WHERE {' AND '.join(propaganda_where)}
        ORDER BY analyzed_at DESC
        LIMIT %s
    """
    refs = await _fetch(refs_sql, params + [limit])
    return count_sql, rows, refs


async def _hate_speech_count(date_from, date_to, category, article_id, limit):
    where: list[str] = []
    params: list[Any] = []
    _date_filter(where, params, "analyzed_at", date_from, date_to)
    if category:
        where.append("predicted_category = %s")
        params.append(category)
    if article_id is not None:
        where.append("article_id = %s")
        params.append(article_id)
    where_sql = f"WHERE {' AND '.join(where)}" if where else ""

    breakdown_sql = f"""
        SELECT predicted_category, COUNT(*) AS count
        FROM article_hate_speech
        {where_sql}
        GROUP BY predicted_category
        ORDER BY count DESC
    """
    rows = await _fetch(breakdown_sql, params)

    refs_sql = f"""
        SELECT article_id
        FROM article_hate_speech
        {where_sql}
        ORDER BY analyzed_at DESC
        LIMIT %s
    """
    refs = await _fetch(refs_sql, params + [limit])
    return breakdown_sql, rows, refs


async def _dialect_breakdown(date_from, date_to, dialect, article_id, limit):
    where: list[str] = []
    params: list[Any] = []
    _date_filter(where, params, "analyzed_at", date_from, date_to)
    if dialect:
        where.append("detected_dialect = %s")
        params.append(dialect)
    if article_id is not None:
        where.append("article_id = %s")
        params.append(article_id)
    where_sql = f"WHERE {' AND '.join(where)}" if where else ""

    breakdown_sql = f"""
        SELECT detected_dialect, COUNT(*) AS count
        FROM article_dialect
        {where_sql}
        GROUP BY detected_dialect
        ORDER BY count DESC
    """
    rows = await _fetch(breakdown_sql, params)

    refs_sql = f"""
        SELECT article_id
        FROM article_dialect
        {where_sql}
        ORDER BY analyzed_at DESC
        LIMIT %s
    """
    refs = await _fetch(refs_sql, params + [limit])
    return breakdown_sql, rows, refs


async def _article_lookup(date_from, date_to, article_id, keyword, limit):
    where: list[str] = []
    params: list[Any] = []
    _date_filter(where, params, "published_at", date_from, date_to)
    if article_id is not None:
        where.append("id = %s")
        params.append(article_id)
    if keyword:
        where.append("title ILIKE %s")
        params.append(f"%{keyword}%")
    where_sql = f"WHERE {' AND '.join(where)}" if where else ""

    sql = f"""
        SELECT id, title, url, published_at
        FROM news_articles
        {where_sql}
        ORDER BY published_at DESC
        LIMIT %s
    """
    rows = await _fetch(sql, params + [limit])
    refs = [{"article_id": r["id"]} for r in rows]
    return sql, rows, refs


@tool
async def sql_query_tool(
    metric: Metric,
    date_from: Optional[str] = None,
    date_to: Optional[str] = None,
    topic: Optional[str] = None,
    sentiment: Optional[str] = None,
    category: Optional[str] = None,
    dialect: Optional[str] = None,
    article_id: Optional[int] = None,
    keyword: Optional[str] = None,
    limit: int = 20,
) -> str:
    """
    Run a pre-defined, read-only query over analyzed news articles.

    metric: which statistic/lookup to run —
        - fake_news_count: breakdown of fake-news verdicts
          (SUPPORTED, REFUTED, PARTIALLY_TRUE, UNVERIFIABLE)
        - articles_by_topic: article counts grouped by primary topic
        - sentiment_breakdown: article counts grouped by sentiment
          (positive, negative, neutral)
        - propaganda_count: count of articles flagged as propaganda,
          plus a technique frequency breakdown
        - hate_speech_count: article counts grouped by hate-speech category
          (none, other, origin, gender, religion)
        - dialect_breakdown: article counts grouped by detected dialect
        - article_lookup: look up specific articles by id or title keyword

    date_from / date_to: ISO dates (YYYY-MM-DD). For all metrics except
        article_lookup, these filter on the analysis timestamp
        (`analyzed_at`) — i.e. when the ML model ran, not when the article
        was published. For article_lookup they filter on `published_at`.
    topic, sentiment, category, dialect: optional exact-match filters for
        the corresponding metric.
    article_id: restrict to a single article.
    keyword: for article_lookup, a substring to search for in the title.
    limit: max number of source article references / lookup rows to return.

    Returns a JSON string: {"rows": [...], "sql_used": "...",
    "source_refs": [{"article_id": ...}, ...]}
    """
    handlers = {
        "fake_news_count": lambda: _fake_news_count(date_from, date_to, article_id, limit),
        "articles_by_topic": lambda: _articles_by_topic(date_from, date_to, topic, article_id, limit),
        "sentiment_breakdown": lambda: _sentiment_breakdown(date_from, date_to, sentiment, article_id, limit),
        "propaganda_count": lambda: _propaganda_count(date_from, date_to, article_id, limit),
        "hate_speech_count": lambda: _hate_speech_count(date_from, date_to, category, article_id, limit),
        "dialect_breakdown": lambda: _dialect_breakdown(date_from, date_to, dialect, article_id, limit),
        "article_lookup": lambda: _article_lookup(date_from, date_to, article_id, keyword, limit),
    }

    handler = handlers.get(metric)
    if handler is None:
        return json.dumps({"rows": [], "sql_used": "", "source_refs": [],
                           "error": f"Unknown metric: {metric}"})

    sql_used, rows, refs = await handler()
    source_refs = [{"type": "article", "id": r["article_id"]} for r in refs]

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