backend/services/response_formatter.py

"""
Response Formatter Service

Handles formatting of query results into citations, charts, GeoJSON layers, and raw data for the frontend.
Separates presentation logic from execution logic.
"""

import logging
from typing import List, Dict, Any, Optional
import uuid

logger = logging.getLogger(__name__)


class ResponseFormatter:
    # Property keys that carry styling/bookkeeping rather than data.
    SYSTEM_PROPERTY_KEYS = (
        "geom", "geometry", "style", "layer_name", "layer_id",
        "choropleth", "fillColor", "color",
    )

    @staticmethod
    def generate_citations(tables: List[str], sql: Optional[str] = None) -> List[str]:
        """
        Build source citations from the catalog's own attribution metadata.

        `tables` is the set of *candidate* tables loaded for the query, which is
        deliberately wider than what the query used. When `sql` is supplied, the
        citation list is narrowed to tables the SQL actually references — citing a
        dataset that did not contribute to the answer misattributes the result.

        Tables absent from the catalog (e.g. session layers created by the agent)
        are skipped, since they are derived rather than sourced.
        """
        from backend.core.data_catalog import get_data_catalog

        try:
            catalog = get_data_catalog()
        except Exception as e:
            logger.warning(f"Could not load catalog for citations: {e}")
            return []

        candidates = list(dict.fromkeys(tables or []))

        if sql:
            referenced = ResponseFormatter._tables_referenced_in_sql(sql, candidates)
            # Only narrow when we positively identified references; if the SQL
            # parsed to nothing recognizable, fall back to the candidate list
            # rather than dropping attribution entirely.
            if referenced:
                candidates = referenced

        citations: List[str] = []
        for table in candidates:
            meta = catalog.get_table_metadata(table)
            if not meta:
                continue
            attribution = meta.get("attribution")
            citations.append(f"{table} ({attribution})" if attribution else table)

        return citations

    @staticmethod
    def _tables_referenced_in_sql(sql: str, candidates: List[str]) -> List[str]:
        """Return the candidate tables that appear as identifiers in the SQL."""
        import re

        # Strip string literals and comments so a table name mentioned inside a
        # WHERE clause value or a comment is not mistaken for a real reference.
        cleaned = re.sub(r"'[^']*'", "''", sql)
        cleaned = re.sub(r"--[^\n]*", " ", cleaned)

        found = []
        for table in candidates:
            if re.search(rf'(?<![\w."]){re.escape(table)}(?![\w."])', cleaned, re.IGNORECASE):
                found.append(table)
        return found

    @staticmethod
    def generate_chart_data(sql: str, features: List[Dict], query: str = "", llm_config: Optional[Dict] = None) -> Optional[Dict[str, Any]]:
        """
        Generates Chart.js compatible data structure.
        Prioritizes llm_config if provided, otherwise falls back to heuristics.
        """
        if not features:
            return None

        # 0. Check LLM explicitly disabled chart
        if llm_config and llm_config.get("use_chart") is False:
            return None
            
        # Try to find string (label) and number (value) in properties
        try:
            chart_items = []
            x_key = "name"
            y_key = "value"
            x_label = "Feature"
            y_label = "Value"
            title = None
            chart_type = "bar"

            # 1. Analyze properties to find X (Label) and Y (Value)
            if features:
                sample_props = features[0].get("properties", {})
                
                # Exclude system keys
                valid_keys = [k for k in sample_props.keys() if k not in ResponseFormatter.SYSTEM_PROPERTY_KEYS]
                
                # Find Y (Value) - First numeric column
                for k in valid_keys:
                    if isinstance(sample_props[k], (int, float)) and not k.endswith("_id") and not k.endswith("_code"):
                        y_key = k
                        y_label = k.replace("_", " ").title()
                        if "sqkm" in k: y_label = "Area (km²)"
                        elif "pop" in k: y_label = "Population"
                        elif "count" in k: y_label = "Count"
                        break
                
                # Find X (Label) - Heuristic Strategy
                found_x = False
                
                # Priority 1: 'name' in key
                for k in valid_keys:
                    if isinstance(sample_props[k], str) and "name" in k.lower():
                        x_key = k
                        x_label = k.replace("_", " ").title().replace("Name", "").strip() or "Region"
                        found_x = True
                        break
                
                # Priority 2: Match LLM x_label (if provided)
                if not found_x and llm_config and llm_config.get("x_label"):
                     target = llm_config.get("x_label").lower()
                     for k in valid_keys:
                         if isinstance(sample_props[k], str) and (k.lower() in target or target in k.lower()):
                             x_key = k
                             x_label = llm_config.get("x_label")
                             found_x = True
                             break

                # Priority 3: Fallback to first string column
                if not found_x:
                    for k in valid_keys:
                        if isinstance(sample_props[k], str):
                            x_key = k
                            x_label = k.replace("_", " ").title()
                            found_x = True
                            break
            
            # 2. Build Data
            # First, check if we have multiple Y keys defined (LLM Config)
            y_keys = [y_key]
            is_multi_metric = False
            
            # Case 1: LLM explicitly provided y_label as a list
            if llm_config and isinstance(llm_config.get("y_label"), list):
                is_multi_metric = True
                possible_y_labels = llm_config.get("y_label")
                y_keys = []
                
                if features:
                    sample_props = features[0].get("properties", {})
                    valid_keys = [k for k in sample_props.keys() if k not in ResponseFormatter.SYSTEM_PROPERTY_KEYS and isinstance(sample_props[k], (int, float))]
                    
                    # For each label, try to find a matching key
                    for lbl in possible_y_labels:
                        lbl_lower = lbl.lower()
                        best_key = None
                        for k in valid_keys:
                            if k in lbl_lower or lbl_lower in k.lower():
                                best_key = k
                                break
                        if not best_key and valid_keys: best_key = valid_keys[0] # Fallback
                        if best_key: y_keys.append(best_key)
                        
                    # Deduplicate
                    y_keys = list(set(y_keys))
                    
            # Case 2: LLM said stacked=True OR type is stacked_bar/multi_bar but y_label is NOT a list - auto-detect numeric columns
            elif llm_config and (llm_config.get("stacked") is True or llm_config.get("type") in ["stacked_bar", "multi_bar"]):
                if features:
                    sample_props = features[0].get("properties", {})
                    # Find ALL numeric columns (exclude IDs and system columns)
                    valid_keys = [k for k in sample_props.keys() 
                                  if k not in ResponseFormatter.SYSTEM_PROPERTY_KEYS
                                  and isinstance(sample_props[k], (int, float))
                                  and not k.endswith("_id") and not k.endswith("_code")]
                    
                    if len(valid_keys) >= 2:
                        is_multi_metric = True
                        y_keys = valid_keys[:3]  # Use up to 3 numeric columns for stacked chart
                        logger.debug(f"Auto-detected multi-metric columns for stacked chart: {y_keys}")

            for f in features:
                props = f.get("properties", {})
                label = props.get(x_key)
                
                if label is not None:
                    item = {"name": str(label)}
                    
                    if is_multi_metric:
                        for i, key in enumerate(y_keys):
                            val = props.get(key)
                            if val is not None:
                                item[f"value{i+1}"] = val
                                item[f"key{i+1}"] = key # Store key name for legend
                    else:
                        val = props.get(y_keys[0])
                        if val is not None:
                            item["value"] = val
                            
                    if "value" in item or "value1" in item:
                        chart_items.append(item)
            
            if not chart_items:
                return None

            # 3. Determine Chart Configuration (LLM vs Heuristics)
            if llm_config and llm_config.get("use_chart") is True:
                # LLM Driven
                chart_type = llm_config.get("type", "bar")
                title = llm_config.get("title", f"{y_label} by {x_label}")
                if llm_config.get("x_label"): x_label = llm_config.get("x_label")
                if llm_config.get("y_label"): y_label = llm_config.get("y_label")
                
                # Special Handling for Histograms
                if chart_type == "histogram":
                    # We need to re-process the data to create bins from the NUMERIC values
                    # Extract all values
                    raw_values = [item["value"] for item in chart_items if isinstance(item["value"], (int, float))]
                    
                    if raw_values:
                        # Simple binning logic (Freedman-Diaconis or sqrt rule simplified)
                        import math
                        n = len(raw_values)
                        if n > 0:
                            min_val = min(raw_values)
                            max_val = max(raw_values)
                            # Sqrt rule for bin count, clamped between 5 and 20
                            num_bins = max(5, min(20, int(math.sqrt(n))))
                            
                            bin_width = (max_val - min_val) / num_bins
                            if bin_width == 0: bin_width = 1
                            
                            bins = [0] * num_bins
                            bin_labels = []
                            
                            # Create labels
                            for i in range(num_bins):
                                start = min_val + (i * bin_width)
                                end = start + bin_width
                                bin_labels.append(f"{start:.1f}-{end:.1f}")
                            
                            # Populate bins
                            for v in raw_values:
                                idx = min(int((v - min_val) / bin_width), num_bins - 1)
                                bins[idx] += 1
                                
                            # Reconstruct chart_items for histogram
                            chart_items = [{"name": lbl, "value": count} for lbl, count in zip(bin_labels, bins)]
                            
                            # Update axes defaults if not set by LLM
                            if not llm_config.get("y_label"): y_label = "Frequency"
                
                # Sort based on logic suitable for chart type
                elif chart_type == "line":
                    # Likely a time series or ordered category
                    try:
                        chart_items.sort(key=lambda x: x["name"])
                    except (KeyError, TypeError):
                        pass
                else:
                    # Default sort by value descending for bar/pie
                    try:
                        chart_items.sort(key=lambda x: x["value"], reverse=True)
                    except (KeyError, TypeError):
                        pass
            else:
                # Heuristic Driven (Legacy Fallback)
                is_time_series = any(t in x_key.lower() for t in ["year", "date", "time", "month", "day"])
                unique_items = len(chart_items)
                query_lower = query.lower()
                
                if "pie" in query_lower:
                    chart_type = "pie"
                    chart_items.sort(key=lambda x: x["value"], reverse=True)
                elif "donut" in query_lower:
                    chart_type = "donut"
                    chart_items.sort(key=lambda x: x["value"], reverse=True)
                elif "line" in query_lower or "trend" in query_lower:
                    chart_type = "line"
                    chart_items.sort(key=lambda x: x["name"])
                elif "bar" in query_lower:
                    chart_type = "bar"
                    chart_items.sort(key=lambda x: x["value"], reverse=True)
                elif is_time_series:
                    chart_type = "line"
                    chart_items.sort(key=lambda x: x["name"])
                elif unique_items <= 5:
                    chart_type = "pie"
                    chart_items.sort(key=lambda x: x["value"], reverse=True)
                else:
                    chart_type = "bar"
                    chart_items.sort(key=lambda x: x["value"], reverse=True)
                
                title = f"{y_label} by {x_label}"

            # Prepare series metadata for frontend
            series_config = None
            if is_multi_metric:
                 series_config = []
                 for i, key in enumerate(y_keys):
                     # Find label
                     label = key.title() # Default
                     if llm_config and isinstance(llm_config.get("y_label"), list) and i < len(llm_config["y_label"]):
                         label = llm_config["y_label"][i]
                     
                     series_config.append({
                         "key": f"value{i+1}",
                         "name": label,
                         "color": None # Frontend handles rotation
                     })

            # Check for stacked override (explicit stacked=True OR type is stacked_bar)
            is_stacked = (llm_config.get("stacked") is True or llm_config.get("type") == "stacked_bar") if llm_config else False
            
            # Normalize chart types for frontend
            if chart_type in ["stacked_bar", "multi_bar"]:
                chart_type = "bar"  # Frontend renders these as bar charts with special handling

            return {
                "type": chart_type,
                "title": title,
                "data": chart_items[:75], # Increased limit for scrollable charts
                "xKey": "name",
                "yKey": "value",
                "series": series_config, # Passing series config
                "stacked": is_stacked,
                "xAxisLabel": x_label,
                "yAxisLabel": y_label if isinstance(y_label, str) else "Values" # Normalize y label if list
            }

        except Exception as e:
            logger.warning(f"Error generating chart data: {e}", exc_info=True)
            return None

    @staticmethod
    def prepare_raw_data(features: List[Dict]) -> List[Dict]:
        """Cleans feature properties for display in the raw data table."""
        raw_data = []
        if not features:
            return raw_data
            
        for f in features:
            props = f.get("properties", {}).copy()
            # Serialize
            props = ResponseFormatter._serialize_properties(props)
            
            # Remove system/visual properties
            for key in ResponseFormatter.SYSTEM_PROPERTY_KEYS:
                props.pop(key, None)
            raw_data.append(props)
            
        return raw_data

    @staticmethod
    def format_geojson_layer(query: str, geojson: Dict[str, Any], features: List[Dict], layer_name: str, layer_emoji: str = "📍", point_style: Optional[str] = None, admin_levels: Optional[List[str]] = None, color_by: Optional[str] = None) -> tuple[Dict[str, Any], str, str]:
        """
        styles the GeoJSON layer and generates metadata (ID, Name, Choropleth).
        
        Args:
            point_style: "icon" for emoji markers, "circle" for simple colored circles, None for auto-detect
            color_by: Explicit column name to use for coloring (from LLM). Overrides auto-detection.
        """
        
        # 0. Serialize properties to avoid datetime errors
        if features:
            for f in features:
                if "properties" in f:
                    f["properties"] = ResponseFormatter._serialize_properties(f["properties"])

        # 2. Random/Distinct Colors
        # Palette of distinct colors (avoiding pure blue which is default)
        palette = [
            "#E63946", # Red
            "#F4A261", # Orange
            "#2A9D8F", # Teal
            "#E9C46A", # Yellow
            "#9C6644", # Brown
            "#D62828", # Dark Red
            "#8338EC", # Purple
            "#3A86FF", # Blue-ish (but distinct)
            "#FB5607", # Orange-Red
            "#FF006E", # Pink
        ]
        
        # Deterministic color based on query hash to keep it stable for same query
        color_idx = abs(hash(query)) % len(palette)
        layer_color = palette[color_idx]

        # Choropleth Logic
        # 1. Identify valid numeric column
        choropleth_col = None
        if features:
            sample = features[0].get("properties", {})
            valid_numerics = [
                k for k, v in sample.items()
                if isinstance(v, (int, float))
                and not isinstance(v, bool)
                and k not in ["layer_id", "style"]
                and not k.endswith("_code")
                and not k.endswith("_id")
                # Ordinal position markers (week_index, band_index) are labels for
                # a step, not a measurement. Colouring by one paints the map with
                # the calendar instead of the data.
                and not k.endswith("_index")
            ]

            # Measurement-like names first. Without `abundance`/`mean` here the
            # fallback below picks whichever numeric happens to come first in the
            # SELECT, which is arbitrary.
            priority_cols = [
                "abundance", "density", "population", "pop", "count", "num",
                "mean", "median", "total", "percent", "area_sqkm", "area",
            ]

            for p in priority_cols:
                matches = [c for c in valid_numerics if p in c]
                if matches:
                    choropleth_col = matches[0]
                    break
            
            # Fallback to first numeric
            if not choropleth_col and valid_numerics:
                choropleth_col = valid_numerics[0]

        # LLM-specified color_by override - use this column if provided
        if color_by and features:
            sample = features[0].get("properties", {})
            if color_by in sample:
                col_value = sample.get(color_by)
                
                # Check if the specified column is numeric or categorical
                if isinstance(col_value, (int, float)):
                    # Numeric column - use for choropleth
                    choropleth_col = color_by
                else:
                    # Categorical column - generate color mapping
                    tableau10 = ['#4e79a7', '#f28e2c', '#e15759', '#76b7b2', '#59a14f', '#edc949', '#af7aa1', '#ff9da7', '#9c755f', '#bab0ab']
                    # Get all unique values, treating None/empty as "(empty)"
                    all_values = [f["properties"].get(color_by) for f in features]
                    unique_categories = sorted(set(
                        "(empty)" if v is None or v == "" else str(v)
                        for v in all_values
                    ))
                    color_map = {cat: tableau10[i % len(tableau10)] for i, cat in enumerate(unique_categories)}
                    
                    geojson["properties"]["categoricalChoropleth"] = {
                        "column": color_by,
                        "colors": color_map,
                        "palette": "tableau10"
                    }
                    
                    # Skip auto-detection since we used LLM's choice
                    layer_id = str(uuid.uuid4())[:8]
                    full_name = layer_name
                    geojson["properties"]["layer_name"] = full_name
                    geojson["properties"]["layer_id"] = layer_id
                    
                    if point_style:
                        geojson["properties"]["pointStyle"] = point_style
                    
                    return (geojson, layer_id, full_name)

        # 2. Try to find categorical column FIRST (priority if matches known category columns)
        categorical_col = None
        categorical_is_priority = False
        if features:
            sample = features[0].get("properties", {})
            valid_strings = [
                k for k, v in sample.items()
                if isinstance(v, str)
                and k not in ["name", "geometry", "geom", "style", "layer_id", "layer_name"]
                and not k.endswith("_name")  # Usually not good for categorical (too many unique values)
                and not k.endswith("_id")
            ]
            
            # Likely category columns, most meaningful first. "season" leads because
            # seasonal breakdowns are the primary categorical axis in this data.
            priority_cols = ["season", "type", "category", "class", "agency", "status", "kind",
                           "amenity", "building",
                           "operator", "operador", "operada", "opera", "operated",
                           "owner", "propietario", "uso", "use", "sector", "estado", "fuente", "source",
                           "proveedor", "provider", "entidad", "entity", "institucion", "institution"]

            def usable_categories(col: str) -> bool:
                """A column is usable if it splits the data into 2-15 groups."""
                uniques = {
                    f["properties"].get(col)
                    for f in features
                    if f["properties"].get(col)
                }
                return 1 < len(uniques) <= 15

            # Walk the priority list and take the first column that actually
            # discriminates. Previously the first name match won outright and was
            # then discarded if it had a single value, losing a perfectly good
            # column further down the list (e.g. "type" is constant while "season"
            # is the real category).
            for p in priority_cols:
                matches = [c for c in valid_strings if p in c.lower() and usable_categories(c)]
                if matches:
                    categorical_col = matches[0]
                    categorical_is_priority = True
                    break

            # Fallback: any string column that discriminates usefully.
            if not categorical_col:
                categorical_col = next((c for c in valid_strings if usable_categories(c)), None)
        
        # 3. Decide: use categorical if priority match, otherwise prefer numeric
        if categorical_col and categorical_is_priority:
            # Generate color mapping using Tableau10 palette
            tableau10 = ['#4e79a7', '#f28e2c', '#e15759', '#76b7b2', '#59a14f', '#edc949', '#af7aa1', '#ff9da7', '#9c755f', '#bab0ab']
            unique_categories = sorted(set(
                f["properties"].get(categorical_col) 
                for f in features 
                if f["properties"].get(categorical_col)
            ))
            color_map = {cat: tableau10[i % len(tableau10)] for i, cat in enumerate(unique_categories)}
            
            geojson["properties"]["categoricalChoropleth"] = {
                "column": categorical_col,
                "colors": color_map,
                "palette": "tableau10"
            }
        elif choropleth_col:
            # Enable numeric choropleth if values actually vary
            values = [f["properties"].get(choropleth_col, 0) for f in features]
            if len(set(values)) > 1:
                geojson["properties"]["choropleth"] = {
                    "enabled": True, 
                    "palette": "viridis",
                    "column": choropleth_col,
                    "scale": "linear"
                }
        elif categorical_col:
            # Use categorical as fallback if no numeric
            tableau10 = ['#4e79a7', '#f28e2c', '#e15759', '#76b7b2', '#59a14f', '#edc949', '#af7aa1', '#ff9da7', '#9c755f', '#bab0ab']
            unique_categories = sorted(set(
                f["properties"].get(categorical_col) 
                for f in features 
                if f["properties"].get(categorical_col)
            ))
            color_map = {cat: tableau10[i % len(tableau10)] for i, cat in enumerate(unique_categories)}
            
            geojson["properties"]["categoricalChoropleth"] = {
                "column": categorical_col,
                "colors": color_map,
                "palette": "tableau10"
            }
        else:
            # Apply random color if NOT a choropleth (numeric or categorical)
            geojson["properties"]["style"] = {
                "color": layer_color,
                "fillColor": layer_color,
                "opacity": 0.8,
                "fillOpacity": 0.4
            }
        
        layer_id = str(uuid.uuid4())[:8]
        geojson["properties"]["layer_name"] = layer_name
        geojson["properties"]["layer_id"] = layer_id
        
        # Add Point Marker Configuration
        # Use pointStyle to determine whether to show icon or circle
        marker_icon = None
        marker_style = "circle"  # default
        
        if point_style == "icon":
            # Use emoji icon for categorical POI
            marker_icon = layer_emoji
            marker_style = "icon"
        elif point_style == "circle":
            # Use simple circle for large datasets or density viz
            marker_icon = None
            marker_style = "circle"
        else:
            # Auto-detect: default to icon for now (backward compatibility)
            marker_icon = layer_emoji
            marker_style = "icon"
        
        geojson["properties"]["pointMarker"] = {
            "icon": marker_icon,
            "style": marker_style,
            "color": layer_color,
            "size": 32
        }
        
        return geojson, layer_id, layer_name

    @staticmethod
    def generate_data_summary(features: List[Dict]) -> str:
        """
        Summarize a result set as short text for the LLM explanation prompt.

        Labels and metrics are discovered from whatever columns the query
        returned, so this works for any dataset in the catalog.
        """
        if not features:
            return "No features found matching the query."

        sample = features[0].get("properties", {})
        data_keys = [k for k in sample if k not in ResponseFormatter.SYSTEM_PROPERTY_KEYS]

        label_key = next(
            (k for k in data_keys if "name" in k.lower() and isinstance(sample[k], str)),
            next((k for k in data_keys if isinstance(sample[k], str)), None),
        )
        metric_key = next(
            (k for k in data_keys
             if isinstance(sample[k], (int, float)) and not isinstance(sample[k], bool)),
            None,
        )

        descriptions = []
        for f in features[:5]:
            props = f.get("properties", {})
            label = str(props.get(label_key)) if label_key and props.get(label_key) is not None else "Feature"
            value = props.get(metric_key) if metric_key else None
            if isinstance(value, (int, float)) and not isinstance(value, bool):
                descriptions.append(f"{label} ({metric_key}={value:,.2f})")
            else:
                descriptions.append(label)

        summary = f"Found {len(features)} features. Sample: {', '.join(descriptions)}"
        if metric_key:
            values = [
                f.get("properties", {}).get(metric_key) for f in features
            ]
            numeric = [v for v in values if isinstance(v, (int, float)) and not isinstance(v, bool)]
            if numeric:
                summary += (
                    f". {metric_key}: min={min(numeric):,.2f}, "
                    f"max={max(numeric):,.2f}, mean={sum(numeric) / len(numeric):,.2f}"
                )
        return summary

    @staticmethod
    def _serialize_properties(properties: Dict[str, Any]) -> Dict[str, Any]:
        """Recursively converts datetime/date objects to strings for JSON serialization."""
        from datetime import datetime, date
        
        serialized = {}
        for k, v in properties.items():
            if isinstance(v, (datetime, date)):
                serialized[k] = v.isoformat()
            elif isinstance(v, dict):
                serialized[k] = ResponseFormatter._serialize_properties(v)
            elif isinstance(v, list):
                serialized[k] = [
                    x.isoformat() if isinstance(x, (datetime, date)) else x 
                    for x in v
                ]
            else:
                serialized[k] = v
        return serialized