backend/app/agents/visualization.py

"""
Visualization Agent — Generates chart configs and auto-insights from query results.
Last agent in the pipeline. Determines optimal chart type and generates follow-ups.
"""

import structlog
from collections import Counter

from app.agents.state import AgentState

logger = structlog.get_logger()

# Chart type selection thresholds
MAX_PIE_CATEGORIES = 8
MIN_LINE_POINTS = 3


def visualization_node(state: AgentState) -> dict:
    """
    Analyze query results and generate visualization config + insights.
    """
    results = state.get("query_results", [])
    columns = state.get("column_names", [])
    user_query = state.get("user_query", "")
    _sql = state.get("sanitized_sql", "") or state.get("generated_sql", "")
    trace_id = state.get("trace_id", "unknown")

    logger.info("agent_started", agent="visualization", trace_id=trace_id)

    if not results:
        return {
            "chart_config": None,
            "chart_type": None,
            "insights": ["No data returned from the query."],
            "follow_up_questions": _generate_followups_empty(user_query),
        }

    # ── Classify columns ─────────────────────────────────
    numeric_cols = []
    text_cols = []
    date_cols = []

    for col in columns:
        sample_val = results[0].get(col)
        if sample_val is None:
            # Check other rows
            for row in results[:5]:
                if row.get(col) is not None:
                    sample_val = row[col]
                    break

        col_lower = col.lower()
        if any(d in col_lower for d in ["date", "time", "created", "updated", "day", "month", "year"]):
            date_cols.append(col)
        elif isinstance(sample_val, (int, float)):
            numeric_cols.append(col)
        else:
            # Try to parse as number
            try:
                if sample_val is not None:
                    float(sample_val)
                    numeric_cols.append(col)
                else:
                    text_cols.append(col)
            except (ValueError, TypeError):
                text_cols.append(col)

    # ── Determine chart type ─────────────────────────────
    chart_config = None
    chart_type = None
    row_count = len(results)

    if numeric_cols and (text_cols or date_cols):
        label_col = date_cols[0] if date_cols else text_cols[0]
        value_col = numeric_cols[0]

        labels = [str(row.get(label_col, "")) for row in results]
        values = []
        for row in results:
            v = row.get(value_col, 0)
            try:
                values.append(float(v) if v is not None else 0)
            except (ValueError, TypeError):
                values.append(0)

        # Choose chart type
        if date_cols and row_count >= MIN_LINE_POINTS:
            chart_type = "line"
        elif row_count <= MAX_PIE_CATEGORIES:
            chart_type = "doughnut"
        else:
            chart_type = "bar"

        colors = [
            "#38bdf8", "#a855f7", "#ec4899", "#22c55e", "#eab308",
            "#f97316", "#14b8a6", "#6366f1", "#f43f5e", "#84cc16",
        ]

        chart_config = {
            "type": chart_type,
            "data": {
                "labels": labels[:50],  # Cap at 50 labels for readability
                "datasets": [{
                    "label": value_col.replace("_", " ").title(),
                    "data": values[:50],
                    "backgroundColor": colors[:len(labels)],
                    "borderColor": "#1e293b",
                    "borderWidth": 2,
                }],
            },
            "options": {
                "responsive": True,
                "maintainAspectRatio": False,
                "plugins": {
                    "legend": {"position": "bottom", "labels": {"color": "#94A3B8"}},
                },
            },
        }

    # ── Generate insights ────────────────────────────────
    insights = _generate_insights(results, numeric_cols, text_cols, date_cols, row_count)

    # ── Generate follow-up questions ─────────────────────
    follow_ups = _generate_followups(user_query, columns, results)

    logger.info(
        "visualization_complete",
        chart_type=chart_type,
        insights_count=len(insights),
        followups_count=len(follow_ups),
    )

    return {
        "chart_config": chart_config,
        "chart_type": chart_type,
        "insights": insights,
        "follow_up_questions": follow_ups,
    }


def _generate_insights(results, numeric_cols, text_cols, date_cols, row_count) -> list[str]:
    """Generate statistical insights from query results."""
    insights = []
    insights.append(f"📊 **{row_count}** records returned")

    for col in numeric_cols[:3]:  # Top 3 numeric columns
        values = []
        for row in results:
            try:
                v = float(row.get(col, 0))
                values.append(v)
            except (ValueError, TypeError):
                continue

        if values:
            avg_val = sum(values) / len(values)
            min_val = min(values)
            max_val = max(values)
            _total = sum(values)

            col_label = col.replace("_", " ").title()
            insights.append(f"**{col_label}**: avg {avg_val:,.2f} | min {min_val:,.2f} | max {max_val:,.2f}")

            if max_val > avg_val * 3 and len(values) > 2:
                insights.append(f"⚠️ Outlier detected in **{col_label}**: max ({max_val:,.2f}) is {max_val/avg_val:.1f}x the average")

            # Trend detection for date-sorted data
            if len(values) >= 3:
                first_half = sum(values[:len(values)//2])
                second_half = sum(values[len(values)//2:])
                if second_half > first_half * 1.2:
                    insights.append(f"📈 Upward trend detected in **{col_label}**")
                elif first_half > second_half * 1.2:
                    insights.append(f"📉 Downward trend detected in **{col_label}**")

    # Text column distribution
    for col in text_cols[:1]:
        values = [str(row.get(col, "")) for row in results]
        counter = Counter(values)
        if len(counter) > 1:
            top_val, top_count = counter.most_common(1)[0]
            pct = (top_count / len(values)) * 100
            if pct < 100:
                insights.append(f"🏷️ Most common **{col.replace('_', ' ')}**: '{top_val}' ({pct:.0f}%)")

    return insights


def _generate_followups(query: str, columns: list, results: list) -> list[str]:
    """Generate context-aware follow-up suggestions."""
    followups = []
    query_lower = query.lower()

    if "top" in query_lower or "best" in query_lower:
        followups.append("Show the bottom performers instead")
    
    if "salary" in query_lower or "amount" in query_lower or "revenue" in query_lower:
        followups.append("Show the distribution by department")
        followups.append("Compare with last year's data")

    if len(results) > 10:
        followups.append("Show only the top 5 results")

    if any(c.lower() in ["department", "region", "category"] for c in columns):
        followups.append("Break down by category")

    # Always offer these
    followups.extend([
        "Visualize this as a chart",
        "Export these results",
    ])

    return followups[:5]  # Cap at 5 suggestions


def _generate_followups_empty(query: str) -> list[str]:
    """Follow-up suggestions when no results are returned."""
    return [
        "Show all available data from this table",
        "List the tables in the database",
        "Try a broader search criteria",
    ]