tools.py

import logging

import obd
from langchain_core.tools import tool
from pydantic import BaseModel, Field

from diagnosis_tools import DIAGNOSIS_TOOLS
from obd_connection import (
    get_poll_pids,
    obd_session_lock,
    probe_obd_vitals,
    read_dtcs,
    read_live_vitals,
)
from obd_vitals_catalog import format_obd_vitals_reference
from vitals_store import get_range, get_recent, get_vital_history, get_vital_history_in_range

logger = logging.getLogger("car-diagnostic-agent.tools")

# Labels persisted by the background poller (see POLL_PIDS in obd_connection.py).
STORED_VITAL_LABELS: tuple[str, ...] = tuple(
    label for label, _ in get_poll_pids()
)

_KEY_LABELS = STORED_VITAL_LABELS

RANGE_SNAPSHOT_CAP = 60


def _compact_snapshot(snapshot: dict) -> dict:
    readings = snapshot.get("readings") or {}
    if isinstance(readings, list):
        readings = {r["label"]: r["value"] for r in readings if "label" in r}
    compact_readings = {k: readings[k] for k in _KEY_LABELS if k in readings}
    return {
        "ts": snapshot["ts"],
        "dtcs": snapshot.get("dtcs") or [],
        "readings": compact_readings,
    }


def _agent_snapshot(snapshot: dict) -> dict:
    """Return compact snapshot fields visible to the diagnostic agent."""
    compact = _compact_snapshot(snapshot)
    compact["id"] = snapshot["id"]
    return compact


def _cap_snapshots(snapshots: list[dict]) -> tuple[list[dict], bool]:
    if len(snapshots) <= RANGE_SNAPSHOT_CAP:
        return snapshots, False
    return snapshots[:RANGE_SNAPSHOT_CAP], True


def _compact_live_vitals(vitals: list[dict]) -> list[dict]:
    return [{"name": v.get("name"), "value": v.get("value")} for v in vitals]


def _resolve_stored_vital_label(vital: str) -> str | None:
    """Map a tool argument to the label used in SQLite snapshots."""
    req = vital.strip()
    if not req:
        return None
    if req in STORED_VITAL_LABELS:
        return req
    by_lower = {label.lower(): label for label in STORED_VITAL_LABELS}
    if req.lower() in by_lower:
        return by_lower[req.lower()]
    key = req.upper().replace(" ", "_")
    for label, command in get_poll_pids():
        if command.name == key or label.lower() == req.lower():
            return label
    if hasattr(obd.commands, key):
        cmd = getattr(obd.commands, key)
        for label, command in get_poll_pids():
            if command.name == cmd.name:
                return label
    return None


class GetStoredVitalHistoryInput(BaseModel):
    vital: str = Field(
        description=(
            "Stored vital label or OBD command name "
            f"(e.g. {', '.join(STORED_VITAL_LABELS[:6])}, …)."
        ),
    )
    count: int = Field(
        default=30,
        ge=1,
        le=500,
        description="Number of most recent stored samples (default 30). Ignored if start_iso/end_iso set.",
    )
    start_iso: str | None = Field(
        default=None,
        description="Optional UTC ISO-8601 range start (use with end_iso instead of count).",
    )
    end_iso: str | None = Field(
        default=None,
        description="Optional UTC ISO-8601 range end (use with start_iso).",
    )


class ReadDtcCodesInput(BaseModel):
    protocol: str = Field(
        default="auto",
        description="OBD-II protocol to use when reading codes (default: auto).",
    )


class GetLiveVitalsInput(BaseModel):
    vital: str | None = Field(
        default=None,
        description=(
            "Optional OBD command name for one reading "
            "(e.g. RPM, COOLANT_TEMP, MAF, SHORT_FUEL_TRIM_1, GET_DTC). "
            "Omit for a compact summary of key vitals (not every PID)."
        ),
    )


class GetRecentVitalsInput(BaseModel):
    count: int = Field(
        default=10,
        ge=1,
        le=600,
        description=(
            "Number of most recent vitals snapshots (default 10; use 5–15 for trends). "
            "Returns compact key readings and DTCs per snapshot."
        ),
    )


class GetVitalsRangeInput(BaseModel):
    start_iso: str = Field(
        description="Start of time range (UTC ISO-8601), e.g. 2026-06-10T14:00:00+00:00.",
    )
    end_iso: str = Field(
        description="End of time range (UTC ISO-8601), e.g. 2026-06-10T14:30:00+00:00.",
    )


@tool
def obd_vitals_reference() -> str:
    """Compact OBD vitals reference: stored labels, command names, and how to query live/historical data.

    Returns a short guide (not the full 86-row catalog). Use get_live_vitals(vital=COMMAND)
    for any supported PID. Call BEFORE investigating symptoms.
    """
    logger.info("obd_vitals_reference invoked")
    return format_obd_vitals_reference()


@tool
def check_obd_vitals_link() -> dict:
    """Verify the OBD connection is working and vitals can be read from the ECU.

    Call after obd_vitals_reference and before get_live_vitals if readings fail
    or before starting an investigation. Returns connection status and a sample RPM read.
    """
    logger.info("check_obd_vitals_link invoked")
    result = probe_obd_vitals()
    logger.info(
        "check_obd_vitals_link: can_receive_vitals=%s status=%s",
        result.get("can_receive_vitals"),
        result.get("status"),
    )
    return result


@tool
def get_car_info() -> dict:
    """Get vehicle make, model, year, and engine information."""
    logger.info("get_car_info invoked")
    result = {
        "make": "Toyota",
        "model": "Auris Hybrid",
        "year": 2012,
        "engine": "1.8L Hybrid",
    }
    logger.info("get_car_info result: %s", result)
    return result


@tool(args_schema=ReadDtcCodesInput)
def read_dtc_codes(protocol: str = "auto") -> dict:
    """Read stored OBD-II diagnostic trouble codes (DTCs) from the vehicle ECU."""
    logger.info("read_dtc_codes invoked (protocol=%s)", protocol)
    with obd_session_lock:
        codes = read_dtcs()
    result = {
        "protocol": protocol,
        "codes": codes,
        "count": len(codes),
    }
    logger.info("read_dtc_codes result: %s", result)
    return result


@tool(args_schema=GetLiveVitalsInput)
def get_live_vitals(vital: str | None = None) -> dict:
    """Query the ECU now for live OBD-II sensor readings. Pass vital for one PID, or omit for key vitals."""
    logger.info("get_live_vitals invoked (vital=%s)", vital)
    try:
        with obd_session_lock:
            data = read_live_vitals(command=vital)
    except ValueError as exc:
        logger.warning("get_live_vitals failed: %s", exc)
        return {
            "command": vital,
            "count": 0,
            "vitals": [],
            "dtcs": [],
            "error": str(exc),
        }
    vitals = data["vitals"]
    if vital is None:
        vitals = _compact_live_vitals(vitals)
    result = {
        "command": data.get("command"),
        "count": len(vitals),
        "vitals": vitals,
        "dtcs": data["dtcs"],
    }
    logger.info(
        "get_live_vitals returned %d reading(s), %d DTC(s)",
        result["count"],
        len(result["dtcs"]),
    )
    return result


@tool(args_schema=GetRecentVitalsInput)
def get_recent_vitals(count: int = 10) -> dict:
    """Return recent stored OBD snapshots with timestamps, DTCs, and key sensor readings."""
    logger.info("get_recent_vitals invoked (count=%s)", count)
    snapshots = [_agent_snapshot(s) for s in get_recent(count)]
    result = {"count": len(snapshots), "snapshots": snapshots}
    logger.info("get_recent_vitals returned %d snapshot(s)", len(snapshots))
    return result


@tool(args_schema=GetVitalsRangeInput)
def get_vitals_in_range(start_iso: str, end_iso: str) -> dict:
    """Return OBD vitals snapshots between two UTC timestamps (compact key readings + DTCs)."""
    logger.info(
        "get_vitals_in_range invoked (start=%s, end=%s)",
        start_iso,
        end_iso,
    )
    raw = get_range(start_iso, end_iso)
    truncated, was_truncated = _cap_snapshots(raw)
    snapshots = [_agent_snapshot(s) for s in truncated]
    result = {
        "start_iso": start_iso,
        "end_iso": end_iso,
        "count": len(snapshots),
        "snapshots": snapshots,
        "truncated": was_truncated,
        "total_in_range": len(raw),
    }
    logger.info("get_vitals_in_range returned %d snapshot(s)", len(snapshots))
    return result


@tool(args_schema=GetStoredVitalHistoryInput)
def get_stored_vital_history(
    vital: str,
    count: int = 30,
    start_iso: str | None = None,
    end_iso: str | None = None,
) -> dict:
    """Return chronological history of one vital from stored OBD snapshots (oldest → newest)."""
    logger.info(
        "get_stored_vital_history invoked (vital=%s, count=%s, start=%s, end=%s)",
        vital,
        count,
        start_iso,
        end_iso,
    )
    label = _resolve_stored_vital_label(vital)
    if label is None:
        return {
            "vital": vital,
            "count": 0,
            "history": [],
            "error": (
                f"Unknown stored vital '{vital}'. "
                f"Use a label like: {', '.join(STORED_VITAL_LABELS)}"
            ),
        }

    if start_iso and end_iso:
        history = get_vital_history_in_range(label, start_iso, end_iso)
        result = {
            "vital": label,
            "start_iso": start_iso,
            "end_iso": end_iso,
            "count": len(history),
            "history": history,
        }
    elif start_iso or end_iso:
        return {
            "vital": label,
            "count": 0,
            "history": [],
            "error": "Provide both start_iso and end_iso for a time range, or omit both to use count.",
        }
    else:
        history = get_vital_history(label, count)
        result = {
            "vital": label,
            "count": len(history),
            "history": history,
        }

    logger.info(
        "get_stored_vital_history returned %d point(s) for %s",
        result["count"],
        label,
    )
    return result


DIAGNOSTIC_TOOLS = [
    obd_vitals_reference,
    check_obd_vitals_link,
    get_car_info,
    read_dtc_codes,
    get_live_vitals,
    get_recent_vitals,
    get_vitals_in_range,
    get_stored_vital_history,
    *DIAGNOSIS_TOOLS,
]