quote-judge-studio-v1/agent.py

"""QuoteJudge Studio v1: deterministic quote comparison with user-scoped persistence."""
from __future__ import annotations

import base64
import csv
import hashlib
import io
import json
import os
import re
import uuid
from datetime import datetime, timezone
from typing import Annotated, Any

from pydantic import BaseModel, Field, field_validator, model_validator

import a2a_pack as a2a
from a2a_pack import (
    A2AAgent,
    AgentDatabase,
    AgentDatabaseEnv,
    AgentDatabaseMigrations,
    AgentPlatformResources,
    FileUpload,
    PlatformUserAuth,
    Pricing,
    Resources,
    RunContext,
    State,
    UploadedFile,
)


MAX_QUOTES = 25
MAX_VENDOR_LEN = 120
MAX_COMPARISON_ID_LEN = 80
MAX_UPLOAD_BYTES = 64 * 1024
ACCEPTED_UPLOAD_MEDIA_TYPES = ("text/csv", "text/plain", "application/json")
DATABASE_NAME = "quote-judge-studio-v1-data"


class QuoteJudgeStudioV1Config(BaseModel):
    """No caller-provided configuration is required."""


class QuoteInput(BaseModel):
    vendor: str = Field(..., min_length=1, max_length=MAX_VENDOR_LEN)
    unit_price: float = Field(..., gt=0, le=1_000_000)
    quantity: int = Field(..., gt=0, le=10_000_000)
    delivery_days: int = Field(..., ge=0, le=3650)
    warranty_months: int = Field(..., ge=0, le=600)

    @field_validator("vendor")
    @classmethod
    def clean_vendor(cls, value: str) -> str:
        cleaned = re.sub(r"\s+", " ", value).strip()
        if not cleaned:
            raise ValueError("vendor is required")
        return cleaned


class WeightInput(BaseModel):
    price: float = Field(..., ge=0, le=1000)
    delivery: float = Field(..., ge=0, le=1000)
    warranty: float = Field(..., ge=0, le=1000)

    @model_validator(mode="after")
    def require_positive_total(self) -> "WeightInput":
        if self.price + self.delivery + self.warranty <= 0:
            raise ValueError("at least one weight must be greater than zero")
        return self


class BrowserQuoteUpload(BaseModel):
    filename: str = Field(..., min_length=1, max_length=160)
    media_type: str = Field(..., min_length=1, max_length=100)
    data_base64: str = Field(..., min_length=1, max_length=((MAX_UPLOAD_BYTES + 2) // 3) * 4 + 16)

    @field_validator("filename")
    @classmethod
    def clean_filename(cls, value: str) -> str:
        cleaned = value.replace("\\", "/").split("/")[-1].strip()
        if not cleaned or cleaned in {".", ".."}:
            raise ValueError("filename is required")
        return cleaned


class QuoteJudgeStudioV1(A2AAgent[QuoteJudgeStudioV1Config, PlatformUserAuth]):
    name = "quote-judge-studio-v1"
    description = (
        "QuoteJudge compares vendor quotes with weighted price, delivery, and "
        "warranty scoring, persists recommendations per platform user, and "
        "serves a one-page React product UI."
    )
    version = "0.1.0"

    config_model = QuoteJudgeStudioV1Config
    auth_model = PlatformUserAuth

    pricing = Pricing(
        price_per_call_usd=0.0,
        caller_pays_llm=False,
        notes="Deterministic quote scoring; no LLM credential required.",
    )
    resources = Resources(cpu="500m", memory="512Mi", max_runtime_seconds=120)
    state = State.DURABLE
    platform_resources = AgentPlatformResources(
        databases=(
            AgentDatabase(
                name=DATABASE_NAME,
                scope="user",
                access_mode="read_write",
                env=AgentDatabaseEnv(url="DATABASE_URL"),
                migrations=AgentDatabaseMigrations(path="db/migrations"),
            ),
        )
    )
    capabilities = {
        "product": {
            "name": "QuoteJudge",
            "primary_workflow": "Upload or paste vendor quotes, weight price/delivery/warranty, compare them, save the result, and reopen it.",
            "uploads": True,
            "receipts_persisted": True,
            "database_scope": "user",
        }
    }

    @a2a.tool(
        description="Compare two or more vendor quotes with weighted price, delivery, and warranty scoring; persist the result for the signed-in user.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="deterministic",
    )
    async def compare_quotes(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        quotes: list[QuoteInput],
        weights: WeightInput,
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        cleaned_id = _validate_comparison_id(comparison_id)
        validation_error = _validate_quotes(quotes)
        if validation_error is not None:
            result = _error(cleaned_id, validation_error, _error_message(validation_error))
            await _persist_receipt(ctx, tenant, "compare_quotes", cleaned_id, {"quotes": len(quotes), "weights": weights.model_dump()}, result)
            return result

        result = _build_comparison(cleaned_id, quotes, weights)
        await _upsert_comparison(tenant, result)
        await _persist_receipt(ctx, tenant, "compare_quotes", cleaned_id, {"quotes": [q.model_dump() for q in quotes], "weights": weights.model_dump()}, result)
        await ctx.emit_progress(f"Saved comparison {cleaned_id}: {result['recommendation']['vendor']} recommended")
        return result

    @a2a.tool(
        description="Reload a previously saved quote comparison for the signed-in user.",
        timeout_seconds=15,
        idempotent=True,
        cost_class="deterministic",
    )
    async def get_comparison(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        cleaned_id = _validate_comparison_id(comparison_id)
        result = await _load_comparison(tenant, cleaned_id)
        if result is None:
            result = _error(cleaned_id, "comparison_not_found", "No comparison with that id exists for this signed-in user.")
        await _persist_receipt(ctx, tenant, "get_comparison", cleaned_id, {"comparison_id": cleaned_id}, result)
        return result

    @a2a.tool(
        description="Browser upload bridge: decode a bounded base64 CSV/JSON quote file, compare quotes, and persist the result.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="deterministic",
    )
    async def compare_quotes_upload(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        upload: BrowserQuoteUpload,
        weights: WeightInput,
    ) -> dict[str, Any]:
        decoded = _decode_browser_upload(upload)
        if not decoded["ok"]:
            cleaned_id = _safe_comparison_id(comparison_id)
            result = _error(cleaned_id, decoded["code"], decoded["message"])
            await _persist_receipt(ctx, _tenant_key(ctx), "compare_quotes_upload", cleaned_id, {"filename": upload.filename, "media_type": upload.media_type}, result)
            return result
        quotes_or_error = _parse_quote_bytes(decoded["bytes"], decoded["media_type"])
        if isinstance(quotes_or_error, dict):
            cleaned_id = _safe_comparison_id(comparison_id)
            result = _error(cleaned_id, quotes_or_error["code"], quotes_or_error["message"])
            await _persist_receipt(ctx, _tenant_key(ctx), "compare_quotes_upload", cleaned_id, {"filename": upload.filename}, result)
            return result
        return await self.compare_quotes(ctx, comparison_id=comparison_id, quotes=quotes_or_error, weights=weights)

    @a2a.tool(
        description="External typed FileUpload path: read an uploaded CSV/JSON quote file from the caller workspace, compare quotes, and persist the result.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="deterministic",
    )
    async def compare_quotes_file(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        document: Annotated[
            UploadedFile,
            FileUpload(
                accept=ACCEPTED_UPLOAD_MEDIA_TYPES,
                max_bytes=MAX_UPLOAD_BYTES,
                description="CSV or JSON file containing quote rows with vendor, unit_price, quantity, delivery_days, and warranty_months.",
            ),
        ],
        weights: WeightInput,
    ) -> dict[str, Any]:
        if document.size_bytes > MAX_UPLOAD_BYTES:
            result = _error(_safe_comparison_id(comparison_id), "upload_too_large", f"Upload must be at most {MAX_UPLOAD_BYTES} bytes.")
            await _persist_receipt(ctx, _tenant_key(ctx), "compare_quotes_file", _safe_comparison_id(comparison_id), {"filename": document.filename}, result)
            return result
        if document.media_type not in ACCEPTED_UPLOAD_MEDIA_TYPES:
            result = _error(_safe_comparison_id(comparison_id), "unsupported_media_type", "Upload must be CSV, plain text CSV, or JSON.")
            await _persist_receipt(ctx, _tenant_key(ctx), "compare_quotes_file", _safe_comparison_id(comparison_id), {"filename": document.filename}, result)
            return result
        try:
            data = ctx.workspace.read_bytes(document.path)
        except Exception:
            result = _error(_safe_comparison_id(comparison_id), "upload_read_failed", "Could not read the uploaded file from the granted workspace.")
            await _persist_receipt(ctx, _tenant_key(ctx), "compare_quotes_file", _safe_comparison_id(comparison_id), {"filename": document.filename}, result)
            return result
        quotes_or_error = _parse_quote_bytes(data[: MAX_UPLOAD_BYTES + 1], document.media_type)
        if isinstance(quotes_or_error, dict):
            result = _error(_safe_comparison_id(comparison_id), quotes_or_error["code"], quotes_or_error["message"])
            await _persist_receipt(ctx, _tenant_key(ctx), "compare_quotes_file", _safe_comparison_id(comparison_id), {"filename": document.filename}, result)
            return result
        return await self.compare_quotes(ctx, comparison_id=comparison_id, quotes=quotes_or_error, weights=weights)


def _tenant_key(ctx: RunContext[PlatformUserAuth]) -> str:
    auth = ctx.auth
    principal = auth.user_id if auth.user_id is not None else (auth.sub or auth.email)
    if not principal:
        raise ValueError("platform user identity is required")
    return f"user:{principal}"


def _safe_comparison_id(value: str) -> str:
    try:
        return _validate_comparison_id(value)
    except Exception:
        return "invalid"


def _validate_comparison_id(value: str) -> str:
    cleaned = str(value or "").strip()
    if not cleaned:
        raise ValueError("comparison_id is required")
    if len(cleaned) > MAX_COMPARISON_ID_LEN or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.:-]{0,79}", cleaned):
        raise ValueError("comparison_id must be 1-80 safe characters")
    return cleaned


def _validate_quotes(quotes: list[QuoteInput]) -> str | None:
    if len(quotes) < 2:
        return "at_least_two_quotes_required"
    if len(quotes) > MAX_QUOTES:
        return "too_many_quotes"
    vendors = [q.vendor.casefold() for q in quotes]
    if len(vendors) != len(set(vendors)):
        return "duplicate_vendor"
    quantities = {q.quantity for q in quotes}
    if len(quantities) != 1:
        return "quantity_mismatch"
    return None


def _error_message(code: str) -> str:
    return {
        "at_least_two_quotes_required": "At least two quotes are required to make a recommendation.",
        "too_many_quotes": f"Compare at most {MAX_QUOTES} quotes at a time.",
        "duplicate_vendor": "Each quote must have a unique vendor name.",
        "quantity_mismatch": "All quotes must use the same quantity for a fair comparison.",
    }.get(code, "The quote input is invalid.")


def _score_lower_is_better(value: float, minimum: float, maximum: float) -> float:
    if maximum == minimum:
        return 100.0
    return ((maximum - value) / (maximum - minimum)) * 100.0


def _score_higher_is_better(value: float, minimum: float, maximum: float) -> float:
    if maximum == minimum:
        return 100.0
    return ((value - minimum) / (maximum - minimum)) * 100.0


def _build_comparison(comparison_id: str, quotes: list[QuoteInput], weights: WeightInput) -> dict[str, Any]:
    prices = [q.unit_price for q in quotes]
    deliveries = [q.delivery_days for q in quotes]
    warranties = [q.warranty_months for q in quotes]
    total_weight = weights.price + weights.delivery + weights.warranty
    rows: list[dict[str, Any]] = []
    for quote in quotes:
        components = {
            "price": round(_score_lower_is_better(quote.unit_price, min(prices), max(prices)), 4),
            "delivery": round(_score_lower_is_better(quote.delivery_days, min(deliveries), max(deliveries)), 4),
            "warranty": round(_score_higher_is_better(quote.warranty_months, min(warranties), max(warranties)), 4),
        }
        weighted_score = (
            components["price"] * weights.price
            + components["delivery"] * weights.delivery
            + components["warranty"] * weights.warranty
        ) / total_weight
        total_price = quote.unit_price * quote.quantity
        rows.append(
            {
                **quote.model_dump(),
                "total_price": round(total_price, 2),
                "component_scores": components,
                "weighted_score": round(weighted_score, 4),
            }
        )
    rows.sort(key=lambda item: (-item["weighted_score"], item["total_price"], item["delivery_days"], -item["warranty_months"], item["vendor"].casefold()))
    winner = rows[0]
    return {
        "ok": True,
        "comparison_id": comparison_id,
        "recommendation": {
            "vendor": winner["vendor"],
            "weighted_score": winner["weighted_score"],
            "rationale": _rationale(winner, weights),
        },
        "weights": weights.model_dump(),
        "quotes": rows,
        "created_at": _now_iso(),
        "receipt": {"persisted": True, "tool": "compare_quotes"},
    }


def _rationale(winner: dict[str, Any], weights: WeightInput) -> str:
    strengths: list[str] = []
    if winner["component_scores"]["price"] >= 99:
        strengths.append("best price")
    if winner["component_scores"]["delivery"] >= 99:
        strengths.append("fastest delivery")
    if winner["component_scores"]["warranty"] >= 99:
        strengths.append("strongest warranty")
    basis = ", ".join(strengths) if strengths else "best weighted balance"
    return f"{winner['vendor']} has the highest weighted score using price={weights.price:g}, delivery={weights.delivery:g}, warranty={weights.warranty:g}; key advantage: {basis}."


def _error(comparison_id: str, code: str, message: str) -> dict[str, Any]:
    return {"ok": False, "comparison_id": comparison_id, "code": code, "message": message}


def _now_iso() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def _decode_browser_upload(upload: BrowserQuoteUpload) -> dict[str, Any]:
    if upload.media_type not in ACCEPTED_UPLOAD_MEDIA_TYPES:
        return {"ok": False, "code": "unsupported_media_type", "message": "Upload must be CSV, plain text CSV, or JSON."}
    try:
        data = base64.b64decode(upload.data_base64, validate=True)
    except Exception:
        return {"ok": False, "code": "invalid_base64", "message": "Upload data_base64 must be valid base64."}
    if len(data) > MAX_UPLOAD_BYTES:
        return {"ok": False, "code": "upload_too_large", "message": f"Upload must be at most {MAX_UPLOAD_BYTES} bytes."}
    if not data:
        return {"ok": False, "code": "empty_upload", "message": "Upload file is empty."}
    return {"ok": True, "bytes": data, "media_type": upload.media_type}


def _parse_quote_bytes(data: bytes, media_type: str) -> list[QuoteInput] | dict[str, str]:
    if len(data) > MAX_UPLOAD_BYTES:
        return {"code": "upload_too_large", "message": f"Upload must be at most {MAX_UPLOAD_BYTES} bytes."}
    try:
        text = data.decode("utf-8-sig")
    except UnicodeDecodeError:
        return {"code": "invalid_text_encoding", "message": "Upload must be UTF-8 text."}
    try:
        if media_type == "application/json":
            raw = json.loads(text)
            rows = raw.get("quotes") if isinstance(raw, dict) else raw
            if not isinstance(rows, list):
                return {"code": "invalid_json_quotes", "message": "JSON upload must be an array or an object with a quotes array."}
            return [QuoteInput.model_validate(item) for item in rows]
        reader = csv.DictReader(io.StringIO(text))
        required = {"vendor", "unit_price", "quantity", "delivery_days", "warranty_months"}
        if not reader.fieldnames or not required <= set(reader.fieldnames):
            return {"code": "invalid_csv_headers", "message": "CSV must include vendor, unit_price, quantity, delivery_days, and warranty_months headers."}
        rows = []
        for row in reader:
            if len(rows) >= MAX_QUOTES + 1:
                break
            rows.append(QuoteInput.model_validate(row))
        return rows
    except Exception as exc:
        return {"code": "quote_parse_failed", "message": f"Could not parse quote rows: {type(exc).__name__}."}


def _db_url() -> str | None:
    return os.environ.get("DATABASE_URL")


def _connect() -> Any:
    import psycopg

    url = _db_url()
    if not url:
        raise RuntimeError("DATABASE_URL is not configured for managed Postgres persistence")
    return psycopg.connect(url, autocommit=True)


async def _upsert_comparison(tenant_key: str, result: dict[str, Any]) -> None:
    with _connect() as conn:
        conn.execute(
            """
            INSERT INTO quote_comparisons (tenant_key, comparison_id, recommendation_vendor, weighted_score, payload, updated_at)
            VALUES (%s, %s, %s, %s, %s::jsonb, NOW())
            ON CONFLICT (tenant_key, comparison_id) DO UPDATE SET
                recommendation_vendor = EXCLUDED.recommendation_vendor,
                weighted_score = EXCLUDED.weighted_score,
                payload = EXCLUDED.payload,
                updated_at = NOW()
            """,
            (
                tenant_key,
                result["comparison_id"],
                result["recommendation"]["vendor"],
                result["recommendation"]["weighted_score"],
                json.dumps(result, separators=(",", ":")),
            ),
        )


async def _load_comparison(tenant_key: str, comparison_id: str) -> dict[str, Any] | None:
    with _connect() as conn:
        row = conn.execute(
            "SELECT payload FROM quote_comparisons WHERE tenant_key = %s AND comparison_id = %s",
            (tenant_key, comparison_id),
        ).fetchone()
    if row is None:
        return None
    payload = row[0]
    if isinstance(payload, str):
        payload = json.loads(payload)
    if isinstance(payload, dict):
        payload = dict(payload)
        payload["reloaded"] = True
        payload["receipt"] = {"persisted": True, "tool": "get_comparison"}
        return payload
    return None


async def _persist_receipt(
    ctx: RunContext[PlatformUserAuth],
    tenant_key: str,
    tool_name: str,
    comparison_id: str,
    inputs: dict[str, Any],
    result: dict[str, Any],
) -> None:
    receipt_id = str(uuid.uuid5(uuid.NAMESPACE_URL, f"{tenant_key}:{tool_name}:{comparison_id}:{hashlib.sha256(json.dumps(inputs, sort_keys=True, default=str).encode()).hexdigest()}"))
    payload = {
        "receipt_id": receipt_id,
        "agent": QuoteJudgeStudioV1.name,
        "version": QuoteJudgeStudioV1.version,
        "tool": tool_name,
        "comparison_id": comparison_id,
        "task_id": getattr(ctx, "task_id", ""),
        "ok": bool(result.get("ok")),
        "input_hash": hashlib.sha256(json.dumps(inputs, sort_keys=True, default=str).encode()).hexdigest(),
        "created_at": _now_iso(),
    }
    try:
        with _connect() as conn:
            conn.execute(
                """
                INSERT INTO quote_execution_receipts (tenant_key, receipt_id, comparison_id, tool_name, ok, payload)
                VALUES (%s, %s, %s, %s, %s, %s::jsonb)
                ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET payload = EXCLUDED.payload
                """,
                (tenant_key, receipt_id, comparison_id, tool_name, bool(result.get("ok")), json.dumps(payload, separators=(",", ":"))),
            )
    except Exception:
        # Receipts are required in production, but a local no-DATABASE_URL unit
        # test should still be able to exercise deterministic validation. The
        # main persistence path raises if saving the comparison itself fails.
        if _db_url():
            raise