Blame for src/tools.py

"""MCP tools for querying territorial ecological indicators.

Design notes (the tools are called by LLMs):
- every tool returns a compact JSON string (no indentation, ASCII kept as UTF-8);
- long lists are capped and the payload says so (`truncated`, `total_count`);
- errors are returned as JSON with a `hint`, never raised to the client;
- parameter names are stable (existing clients keep working).
"""

import json
import logging
import time
import hashlib
from collections import defaultdict
from dataclasses import dataclass, field
from functools import wraps
from typing import Any, Callable

import httpx

from .api_client import get_client, CubeJsClientError
from .cache import fold_text, get_cache, initialize_cache, refresh_cache_if_needed
from .cube_resolver import get_resolver
from .models import (
    SourceMetadata,
    GEOGRAPHIC_LEVELS,
    GEO_DIMENSION_PATTERNS,
    REGION_CODES,
)

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s | %(levelname)s | %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger("mcp_tools")

# Result caps (tokens matter for LLM clients)
LIST_MAX = 200          # list_indicators: catalogue is ~170 items, never truncated in practice
SEARCH_MAX = 30         # search_indicators
QUERY_DEFAULT_LIMIT = 100
QUERY_MAX_LIMIT = 500   # query_indicator_data: hard cap (Cube.js limit)
TERRITORY_MAX = 20


# =============================================================================
# Session Tracker - Track usage patterns across calls
# =============================================================================
@dataclass
class SessionData:
    """Track data for a single session."""
    session_id: str
    start_time: float = field(default_factory=time.time)
    calls: list = field(default_factory=list)
    last_call_time: float = 0
    indicators_queried: set = field(default_factory=set)
    levels_queried: set = field(default_factory=set)

    def add_call(self, tool: str, params: dict, duration_ms: int,
                 result_count: int, response_size: int, status: str):
        """Record a tool call."""
        now = time.time()
        time_since_last = int((now - self.last_call_time) * 1000) if self.last_call_time else 0

        self.calls.append({
            "tool": tool,
            "params": params,
            "duration_ms": duration_ms,
            "result_count": result_count,
            "response_size": response_size,
            "status": status,
            "time_since_last_ms": time_since_last,
        })
        self.last_call_time = now

        # Track what's being queried
        if "indicator_id" in params:
            self.indicators_queried.add(params["indicator_id"])
        if "geographic_level" in params:
            self.levels_queried.add(params["geographic_level"])

    def get_sequence(self) -> str:
        """Get the sequence of tools called."""
        return "→".join(c["tool"].replace("_indicators", "").replace("_indicator", "")
                       for c in self.calls)

    def get_total_duration_ms(self) -> int:
        """Total time spent in API calls."""
        return sum(c["duration_ms"] for c in self.calls)


class UsageTracker:
    """Track MCP usage patterns across sessions."""

    # Session timeout in seconds (new session if no call for 5 minutes)
    SESSION_TIMEOUT = 300

    def __init__(self):
        self.sessions: dict[str, SessionData] = {}
        self.patterns: defaultdict[str, int] = defaultdict(int)  # sequence -> count
        self.tool_stats: defaultdict[str, dict] = defaultdict(
            lambda: {"calls": 0, "total_ms": 0, "errors": 0}
        )

    def get_or_create_session(self, session_hint: str = "default") -> SessionData:
        """Get existing session or create new one."""
        # Simple session management based on hint (could be IP, user-agent hash, etc.)
        session_id = hashlib.md5(session_hint.encode()).hexdigest()[:8]

        now = time.time()

        # Check if session exists and is not expired
        if session_id in self.sessions:
            session = self.sessions[session_id]
            if session.last_call_time and (now - session.last_call_time) > self.SESSION_TIMEOUT:
                # Session expired, log pattern and create new
                self._finalize_session(session)
                session = SessionData(session_id=session_id)
                self.sessions[session_id] = session
                logger.info(f"[SESSION] id={session_id} | new_session (previous expired)")
        else:
            session = SessionData(session_id=session_id)
            self.sessions[session_id] = session
            logger.info(f"[SESSION] id={session_id} | new_session")

        return session

    def _finalize_session(self, session: SessionData):
        """Log session summary when it ends."""
        if len(session.calls) > 1:
            sequence = session.get_sequence()
            self.patterns[sequence] += 1

            logger.info(
                f"[PATTERN] id={session.session_id} | "
                f"sequence={sequence} | "
                f"calls={len(session.calls)} | "
                f"total_ms={session.get_total_duration_ms()} | "
                f"indicators={list(session.indicators_queried)} | "
                f"levels={list(session.levels_queried)}"
            )

    def log_stats_summary(self):
        """Log accumulated statistics."""
        if self.patterns:
            top_patterns = sorted(self.patterns.items(), key=lambda x: -x[1])[:5]
            logger.info(f"[STATS] top_patterns={top_patterns}")


# Global tracker instance
_tracker = UsageTracker()


def log_tool_call(func: Callable) -> Callable:
    """Decorator to log MCP tool calls with rich metrics."""
    @wraps(func)
    async def wrapper(*args, **kwargs):
        tool_name = func.__name__
        start_time = time.time()

        # Get or create session
        session = _tracker.get_or_create_session()

        # Extract params (only non-empty)
        params = {k: v for k, v in kwargs.items() if v}

        # Build context info
        call_num = len(session.calls) + 1
        prev_tool = session.calls[-1]["tool"] if session.calls else None

        # Log the call with context
        context = f"call#{call_num}"
        if prev_tool:
            context += f" | prev={prev_tool}"
        logger.info(f"[CALL] {tool_name} | {context} | params={params}")

        try:
            result = await func(*args, **kwargs)
            elapsed_ms = int((time.time() - start_time) * 1000)
            response_size = len(result.encode('utf-8'))

            # Parse result to get metrics
            status = "ok"
            result_count = 0
            try:
                result_data = json.loads(result)
                if "error" in result_data:
                    status = "error"
                    logger.warning(
                        f"[ERROR] {tool_name} | {elapsed_ms}ms | "
                        f"error={result_data['error'][:100]}"
                    )
                else:
                    result_count = (
                        result_data.get("count") or
                        result_data.get("total_count") or
                        len(result_data.get("data", [])) or
                        (1 if "metadata" in result_data else 0)
                    )
                    logger.info(
                        f"[OK] {tool_name} | {elapsed_ms}ms | "
                        f"count={result_count} | size={response_size}B"
                    )
            except json.JSONDecodeError:
                logger.info(f"[OK] {tool_name} | {elapsed_ms}ms | size={response_size}B")

            # Record in session
            session.add_call(
                tool=tool_name,
                params=params,
                duration_ms=elapsed_ms,
                result_count=result_count,
                response_size=response_size,
                status=status,
            )

            # Update global stats
            _tracker.tool_stats[tool_name]["calls"] += 1
            _tracker.tool_stats[tool_name]["total_ms"] += elapsed_ms
            if status == "error":
                _tracker.tool_stats[tool_name]["errors"] += 1

            return result

        except Exception as e:
            elapsed_ms = int((time.time() - start_time) * 1000)
            logger.error(f"[EXCEPTION] {tool_name} | {elapsed_ms}ms | {type(e).__name__}: {e}")
            _tracker.tool_stats[tool_name]["errors"] += 1
            raise

    return wrapper


# =============================================================================
# Helpers
# =============================================================================
def _dump(payload: dict[str, Any]) -> str:
    """Compact JSON (fewer tokens for LLM clients)."""
    return json.dumps(payload, ensure_ascii=False, separators=(",", ":"))


def _error(message: str, **extra: Any) -> str:
    return _dump({"error": message, **extra})


def _parse_int(value: str | int | None, default: int, lo: int, hi: int) -> int:
    try:
        n = int(value)  # type: ignore[arg-type]
    except (TypeError, ValueError):
        return default
    return max(lo, min(n, hi))


def _to_number(value: Any) -> Any:
    """Cube.js returns numeric measures as strings; convert when possible."""
    if isinstance(value, (int, float)) or value is None:
        return value
    try:
        f = float(value)
    except (TypeError, ValueError):
        return value
    return int(f) if f.is_integer() else round(f, 4)


def _split_csv(value: str | None) -> list[str]:
    return [v.strip() for v in (value or "").split(",") if v.strip()]


def _list_item(ind: Any) -> dict[str, Any]:
    return {
        "id": ind.id,
        "libelle": ind.libelle,
        "unite": ind.unite,
        "mailles": _split_csv(ind.mailles_disponibles),
        "thematique": ind.thematique_fnv,
    }


async def _ensure_cache_initialized() -> str | None:
    """Ensure the cache is initialized; return an error JSON string on failure."""
    cache = get_cache()
    try:
        if not cache.is_initialized:
            await initialize_cache()
        else:
            await refresh_cache_if_needed()
    except ValueError as e:  # missing token
        return _error(f"Server misconfigured: {e}")
    except (CubeJsClientError, httpx.HTTPError) as e:
        if cache.is_initialized:
            return None  # stale cache is better than no answer
        return _error(f"Indicators API unavailable: {e}", hint="Retry in a few seconds.")
    return None


# =============================================================================
# MCP tools
# =============================================================================
@log_tool_call
async def list_indicators(
    thematique: str = "",
    maille: str = "",
) -> str:
    """List the catalogue of French territorial ecological indicators (~170), optionally filtered.

    Prefer search_indicators for a topic; use this tool to browse a theme or to know
    what exists at a given geographic level.

    Args:
        thematique: Optional France Nation Verte theme, partial and accent-insensitive match, e.g. "déplacer" (mobility), "loger" (housing), "produire", "nourrir", "consommer", "ecosystèmes".
        maille: Optional geographic level the indicator must be available at: "region", "departement", "epci" or "commune".

    Returns:
        JSON: count, indicators [{id, libelle, unite, mailles, thematique}], filters_applied, themes (distinct values, to refine).
    """
    err = await _ensure_cache_initialized()
    if err:
        return err
    cache = get_cache()

    theme_filter = thematique.strip() if thematique else None
    maille_filter = maille.strip().lower() if maille else None

    if maille_filter and maille_filter not in GEOGRAPHIC_LEVELS:
        return _error(f"Invalid geographic level: {maille}", valid_levels=GEOGRAPHIC_LEVELS)

    indicators = cache.list_indicators(thematique=theme_filter, maille=maille_filter)
    themes = sorted({i.thematique_fnv for i in cache.list_indicators() if i.thematique_fnv})

    payload: dict[str, Any] = {
        "count": len(indicators),
        "indicators": [_list_item(i) for i in indicators[:LIST_MAX]],
        "filters_applied": {"thematique": theme_filter, "maille": maille_filter},
        "themes": themes,
    }
    if len(indicators) > LIST_MAX:
        payload["truncated"] = True
    if not indicators:
        payload["hint"] = "No indicator matches; try another theme (see `themes`) or search_indicators."
    return _dump(payload)


@log_tool_call
async def search_indicators(query: str) -> str:
    """Find indicators by keywords (French), e.g. "émissions gaz effet de serre", "surface bio", "artificialisation".

    Accent-, case- and plural-insensitive; all words must match (in the name, description
    or theme). Best matches first. Then call get_indicator_details(id) or query_indicator_data.

    Args:
        query: Keywords separated by spaces, e.g. "consommation espace", "stationnement vélo", "logements vacants".

    Returns:
        JSON: query, total_count, indicators [{id, libelle, unite, mailles, thematique}] (max 30), truncated flag.
    """
    err = await _ensure_cache_initialized()
    if err:
        return err
    cache = get_cache()

    search_query = query.strip() if query else ""
    indicators = cache.search_indicators(search_query) if search_query else cache.list_indicators()

    payload: dict[str, Any] = {
        "query": search_query,
        "total_count": len(indicators),
        "indicators": [_list_item(i) for i in indicators[:SEARCH_MAX]],
    }
    if len(indicators) > SEARCH_MAX:
        payload["truncated"] = True
        payload["hint"] = "Refine with more specific keywords."
    if not indicators:
        payload["hint"] = "Try fewer or broader words (e.g. one noun), or list_indicators(thematique=...)."
    return _dump(payload)


@log_tool_call
async def get_indicator_details(indicator_id: str) -> str:
    """Full metadata of one indicator: definition, calculation method, available years and geographic levels, data completeness, sources.

    Args:
        indicator_id: Numeric indicator id as returned by search_indicators / list_indicators (e.g. "611").

    Returns:
        JSON: metadata {id, libelle, unite, description, methode_calcul, annees_disponibles [..], mailles_disponibles [..], couverture_geographique, completion {level: ratio 0-1}}, sources [..], available_cubes {level: cube}.
    """
    err = await _ensure_cache_initialized()
    if err:
        return err

    try:
        ind_id = int(str(indicator_id).strip())
    except ValueError:
        return _error(f"Invalid indicator ID: {indicator_id}. Must be a number.")

    cache = get_cache()
    indicator = cache.get_indicator(ind_id)
    if indicator is None:
        return _error(
            f"Indicator {ind_id} not found.",
            hint="Use search_indicators(query) to find the right id.",
        )

    resolver = get_resolver()
    available_cubes = resolver.get_cubes_for_indicator(ind_id)

    # Sources: request only dimensions that the API actually exposes
    client = get_client()
    sources: list[dict[str, Any]] = []
    sources_warning: str | None = None
    try:
        rows = await client.load_sources_metadata(
            indicator_id=ind_id,
            available_dimensions=resolver.get_cube_dimensions("indicateur_x_source_metadata"),
        )
        for row in rows:
            src = SourceMetadata.from_api_response(row).model_dump(exclude_none=True)
            src.pop("id_indicateur", None)
            sources.append(src)
    except (CubeJsClientError, httpx.HTTPError) as e:
        sources_warning = f"Could not fetch sources: {e}"

    metadata = {
        "id": indicator.id,
        "libelle": indicator.libelle,
        "unite": indicator.unite,
        "description": indicator.description,
        "methode_calcul": indicator.methode_calcul,
        "annees_disponibles": _split_csv(indicator.annees_disponibles),
        "mailles_disponibles": _split_csv(indicator.mailles_disponibles),
        "maille_mini_disponible": indicator.maille_mini_disponible,
        "couverture_geographique": indicator.couverture_geographique,
        "completion": {
            level: indicator.get_completion_for_level(level)
            for level in GEOGRAPHIC_LEVELS
            if indicator.get_completion_for_level(level) is not None
        },
        "thematique": indicator.thematique_fnv,
    }
    result: dict[str, Any] = {
        "metadata": metadata,
        "sources": sources,
        "available_cubes": available_cubes,
    }
    if sources_warning:
        result["sources_warning"] = sources_warning
    if not available_cubes:
        result["hint"] = "No data cube found for this indicator: values cannot be queried."
    return _dump(result)


@log_tool_call
async def query_indicator_data(
    indicator_id: str,
    geographic_level: str,
    geographic_code: str = "",
    year: str = "",
    limit: int = QUERY_DEFAULT_LIMIT,
) -> str:
    """Get the values of an indicator for territories (one territory, or all territories of a level).

    Without geographic_code, returns every territory of the level (use `year` and `limit`
    to keep it small: there are 13 regions, ~100 départements, ~1 250 EPCI, ~35 000 communes).
    Without year, returns all available years. Use find_territory to get a code from a name.

    Args:
        indicator_id: Numeric indicator id (e.g. "611").
        geographic_level: "region", "departement", "epci" or "commune".
        geographic_code: Optional territory code: region 2 digits ("93" PACA, "11" Île-de-France), département 2-3 chars ("13", "2A", "974"), EPCI 9-digit SIREN, commune 5-digit INSEE ("75056" Paris).
        year: Optional year, e.g. "2022".
        limit: Max rows returned (default 100, max 500).

    Returns:
        JSON: indicator_id, indicator_name, unite, geographic_level, total_count, truncated, data [{geocode, libelle, annee, valeur}], annees_disponibles, cube.
    """
    err = await _ensure_cache_initialized()
    if err:
        return err

    try:
        ind_id = int(str(indicator_id).strip())
    except ValueError:
        return _error(f"Invalid indicator ID: {indicator_id}. Must be a number.")

    geo_level = (geographic_level or "").strip().lower()
    if geo_level not in GEOGRAPHIC_LEVELS:
        return _error(f"Invalid geographic level: {geographic_level}", valid_levels=GEOGRAPHIC_LEVELS)

    max_rows = _parse_int(limit, QUERY_DEFAULT_LIMIT, 1, QUERY_MAX_LIMIT)

    cache = get_cache()
    resolver = get_resolver()
    indicator = cache.get_indicator(ind_id)
    indicator_name = indicator.libelle if indicator else f"Indicator {ind_id}"
    indicator_unite = indicator.unite if indicator else None
    years_available = _split_csv(indicator.annees_disponibles) if indicator else []

    cube_name = resolver.find_cube_for_indicator(ind_id, geo_level)
    if cube_name is None:
        if not resolver.is_indicator_known(ind_id):
            return _error(
                f"Indicator {ind_id} not found in any data cube.",
                hint="Use search_indicators(query) to find a valid id.",
            )
        available = resolver.get_available_mailles(ind_id)
        return _error(
            f"Indicator {ind_id} is not available at {geo_level} level.",
            available_levels=available,
            hint=f"Try one of: {', '.join(available)}",
        )

    geo_patterns = GEO_DIMENSION_PATTERNS[geo_level]
    measure = resolver.get_measure_name(cube_name, ind_id)
    geocode_dim = resolver.get_dimension_name(cube_name, geo_patterns["geocode"])
    libelle_dim = resolver.get_dimension_name(cube_name, geo_patterns["libelle"])
    annee_dim = resolver.get_dimension_name(cube_name, "annee")

    # Ask for one extra row to detect truncation
    query: dict[str, Any] = {
        "measures": [measure],
        "dimensions": [geocode_dim, libelle_dim, annee_dim],
        "limit": max_rows + 1,
    }
    filters = []
    geo_code = (geographic_code or "").strip()
    if geo_code:
        filters.append({"member": geocode_dim, "operator": "equals", "values": [geo_code]})
    year_filter = (year or "").strip()
    if year_filter:
        filters.append({"member": annee_dim, "operator": "equals", "values": [year_filter]})
    if filters:
        query["filters"] = filters

    client = get_client()
    try:
        result = await client.load(query)
        data_rows = result.get("data", [])
    except (CubeJsClientError, httpx.HTTPError) as e:
        return _error(f"Query failed: {e}", cube=cube_name)

    data_points = [
        {
            "geocode": row.get(geocode_dim),
            "libelle": row.get(libelle_dim),
            "annee": row.get(annee_dim),
            "valeur": _to_number(row.get(measure)),
        }
        for row in data_rows
    ]
    data_points.sort(key=lambda x: (x.get("annee") or "", x.get("libelle") or ""))
    truncated = len(data_points) > max_rows
    data_points = data_points[:max_rows]

    payload: dict[str, Any] = {
        "indicator_id": ind_id,
        "indicator_name": indicator_name,
        "unite": indicator_unite,
        "geographic_level": geo_level,
        "geographic_code": geo_code or None,
        "year": year_filter or None,
        "total_count": len(data_points),
        "truncated": truncated,
        "data": data_points,
        "annees_disponibles": years_available,
        "cube": cube_name,
    }
    if truncated:
        payload["hint"] = (
            f"Only the first {max_rows} rows are returned: filter with geographic_code and/or year, "
            f"or raise limit (max {QUERY_MAX_LIMIT})."
        )
    if not data_points:
        hints = []
        if year_filter and years_available and year_filter not in years_available:
            hints.append(f"Year {year_filter} is not available (available: {', '.join(years_available)}).")
        if geo_code:
            hints.append(
                "Check the code format for this level (region: 2 digits, département: 2-3 chars, "
                "EPCI: 9-digit SIREN, commune: 5-digit INSEE) or use find_territory(name, level)."
            )
        payload["hint"] = " ".join(hints) or "No data for these filters."
    return _dump(payload)


@log_tool_call
async def find_territory(name: str, geographic_level: str = "") -> str:
    """Find the official code of a French territory from (part of) its name, e.g. "Lyon", "Bouches-du-Rhône", "Grand Paris".

    Use it before query_indicator_data when the code is unknown (EPCI SIREN and commune
    INSEE codes are hard to guess). Regions are resolved locally; other levels are looked
    up in the indicators API.

    Args:
        name: Territory name or fragment, accent-insensitive (e.g. "Métropole de Lyon", "Toulouse", "2A").
        geographic_level: Optional level to restrict the search: "region", "departement", "epci" or "commune". Default: all levels.

    Returns:
        JSON: name, total_count, territories [{geographic_level, geocode, libelle}] (max 20 per level).
    """
    search = (name or "").strip()
    if not search:
        return _error("name is required.")

    level = (geographic_level or "").strip().lower()
    if level and level not in GEOGRAPHIC_LEVELS:
        return _error(f"Invalid geographic level: {geographic_level}", valid_levels=GEOGRAPHIC_LEVELS)
    levels = [level] if level else GEOGRAPHIC_LEVELS

    found: list[dict[str, Any]] = []
    warnings: list[str] = []
    folded = fold_text(search)

    if "region" in levels:
        for code, label in REGION_CODES.items():
            if folded in fold_text(label) or folded == code:
                found.append({"geographic_level": "region", "geocode": code, "libelle": label})

    other_levels = [lv for lv in levels if lv != "region"]
    if other_levels:
        err = await _ensure_cache_initialized()
        if err:
            return err
        resolver = get_resolver()
        client = get_client()
        for lv in other_levels:
            cube_name = resolver.find_cube_with_maille(lv)
            if cube_name is None:
                warnings.append(f"No data cube exposes the {lv} level.")
                continue
            patterns = GEO_DIMENSION_PATTERNS[lv]
            try:
                rows = await client.find_territories(
                    cube_name, patterns["geocode"], patterns["libelle"], search, limit=TERRITORY_MAX * 3
                )
            except (CubeJsClientError, httpx.HTTPError) as e:
                warnings.append(f"{lv}: lookup failed ({e})")
                continue
            geocode_dim = f"{cube_name}.{patterns['geocode']}"
            libelle_dim = f"{cube_name}.{patterns['libelle']}"
            seen: set[str] = set()
            for row in rows:
                code = row.get(geocode_dim)
                if not code or code in seen:
                    continue
                seen.add(code)
                found.append({"geographic_level": lv, "geocode": code, "libelle": row.get(libelle_dim)})
                if len(seen) >= TERRITORY_MAX:
                    break

    payload: dict[str, Any] = {"name": search, "total_count": len(found), "territories": found}
    if warnings:
        payload["warnings"] = warnings
    if not found:
        payload["hint"] = "No territory matches; try a shorter fragment or another spelling."
    return _dump(payload)


# Export all tools
__all__ = [
    "list_indicators",
    "search_indicators",
    "get_indicator_details",
    "query_indicator_data",
    "find_territory",
]