quote-judge-studio-v1/agent.py

"""QuoteJudge Studio full-stack A2A agent.

Deterministic quote comparison product with platform-user tenant isolation,
managed Postgres persistence, bounded browser uploads, typed external file
uploads, MCP-compatible tools, and execution receipt rows.
"""
from __future__ import annotations

import base64
import csv
import hashlib
import io
import json
import os
import re
from datetime import datetime, timezone
from decimal import Decimal, ROUND_HALF_UP
from typing import Annotated, Any

from pydantic import BaseModel, Field, field_validator

import a2a_pack as a2a
from a2a_pack import (
    A2AAgent,
    AgentDatabase,
    AgentDatabaseEnv,
    AgentDatabaseMigrations,
    AgentPlatformResources,
    PlatformUserAuth,
    Pricing,
    Resources,
    RunContext,
    WorkspaceAccess,
    WorkspaceMode,
)
from a2a_pack.workspace import FileUpload, UploadedFile

DATABASE_NAME = "quote-judge-studio-v1-data"
DB_CONNECT_TIMEOUT_SECONDS = 3
DB_STATEMENT_TIMEOUT_MS = 5000
DB_LOCK_TIMEOUT_MS = 3000
MAX_QUOTES = 20
MAX_UPLOAD_BYTES = 64_000
MAX_BROWSER_UPLOADS = 3
ALLOWED_UPLOAD_MEDIA_TYPES = ("text/csv", "text/plain", "application/json")
VERSION = "0.1.0"


class QuoteJudgeStudioV1Config(BaseModel):
    pass


class QuoteInput(BaseModel):
    vendor: str = Field(min_length=1, max_length=120)
    unit_price: float = Field(gt=0, le=1_000_000)
    quantity: int = Field(gt=0, le=10_000_000)
    delivery_days: int = Field(gt=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 QuoteWeights(BaseModel):
    price: float = Field(ge=0, le=100)
    delivery: float = Field(ge=0, le=100)
    warranty: float = Field(ge=0, le=100)


class BrowserQuoteDocument(BaseModel):
    filename: str = Field(min_length=1, max_length=160)
    media_type: str = Field(min_length=1, max_length=120)
    data_base64: str = Field(min_length=1, max_length=MAX_UPLOAD_BYTES * 2)

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

    @field_validator("media_type")
    @classmethod
    def safe_media_type(cls, value: str) -> str:
        cleaned = value.split(";", 1)[0].strip().lower()
        if cleaned not in ALLOWED_UPLOAD_MEDIA_TYPES:
            raise ValueError(
                "unsupported media_type; use text/csv, text/plain, or application/json"
            )
        return cleaned


class QuoteJudgeStudioV1(A2AAgent[QuoteJudgeStudioV1Config, PlatformUserAuth]):
    name = "quote-judge-studio-v1"
    description = (
        "Compare vendor quotes by price, delivery, and warranty; persist the "
        "winner per signed-in A2A Cloud user; reopen saved comparisons."
    )
    version = VERSION

    config_model = QuoteJudgeStudioV1Config
    auth_model = PlatformUserAuth

    # Deterministic local logic only. No LLM credential is required.
    pricing = Pricing(
        price_per_call_usd=0.0,
        caller_pays_llm=False,
        notes="Deterministic QuoteJudge comparison; no LLM is called.",
    )
    resources = Resources(cpu="250m", memory="512Mi", max_runtime_seconds=120)
    workspace_access = WorkspaceAccess.dynamic(
        max_files=8,
        allowed_modes=(WorkspaceMode.READ_ONLY,),
        require_reason=False,
        max_total_size_bytes=MAX_UPLOAD_BYTES * MAX_BROWSER_UPLOADS,
    )
    tools_used = ("postgres", "managed-postgres", "mcp", "packed-frontend")
    platform_resources = AgentPlatformResources(
        databases=(
            AgentDatabase(
                name=DATABASE_NAME,
                provider="neon",
                engine="postgres",
                scope="user",
                branch="main",
                access_mode="read_write",
                env=AgentDatabaseEnv(url="DATABASE_URL"),
                migrations=AgentDatabaseMigrations(path="db/migrations"),
                scale_to_zero=True,
            ),
        )
    )

    @a2a.tool(
        description="Compare two or more vendor quotes, persist the result for this platform user, and return the recommendation.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="deterministic",
    )
    async def compare_quotes(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        quotes: list[QuoteInput],
        weights: QuoteWeights,
    ) -> dict[str, Any]:
        validation = _validate_request(comparison_id, quotes, weights)
        if validation is not None:
            await _persist_receipt_best_effort(ctx, "compare_quotes", comparison_id, "error", validation)
            return validation

        result = _score_quotes(comparison_id, quotes, weights)
        tenant = _tenant_key(ctx)
        try:
            await _save_comparison(tenant, result)
            await _persist_receipt_best_effort(ctx, "compare_quotes", comparison_id, "ok", result)
        except RuntimeError as exc:
            return _setup_required(comparison_id, str(exc))
        except Exception:  # noqa: BLE001
            return _temporary_db_error(comparison_id)
        await ctx.emit_progress(f"saved comparison {comparison_id}")
        return result

    @a2a.tool(
        description="Reload a saved quote comparison for the current signed-in platform user.",
        timeout_seconds=20,
        idempotent=True,
        cost_class="deterministic",
    )
    async def get_comparison(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
    ) -> dict[str, Any]:
        tenant = _tenant_key(ctx)
        clean_id = _clean_comparison_id(comparison_id)
        if clean_id is None:
            return {
                "ok": False,
                "code": "invalid_comparison_id",
                "comparison_id": comparison_id,
                "message": "comparison_id must be 1-120 URL-safe characters",
            }
        try:
            record = await _load_comparison(tenant, clean_id)
            await _persist_receipt_best_effort(ctx, "get_comparison", clean_id, "ok", record or {"ok": False})
        except RuntimeError as exc:
            return _setup_required(clean_id, str(exc))
        except Exception:  # noqa: BLE001
            return _temporary_db_error(clean_id)
        if record is None:
            return {
                "ok": False,
                "code": "comparison_not_found",
                "comparison_id": clean_id,
                "message": "No saved comparison exists for this user and id.",
            }
        return record

    @a2a.tool(
        description="Compare quotes from an external client FileUpload (CSV, JSON, or plain text) and persist the result.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="deterministic",
    )
    async def compare_quotes_file(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        weights: QuoteWeights,
        document: Annotated[
            UploadedFile,
            FileUpload(
                accept=ALLOWED_UPLOAD_MEDIA_TYPES,
                max_bytes=MAX_UPLOAD_BYTES,
                description="CSV/JSON/text quote file with vendor, unit_price, quantity, delivery_days, warranty_months.",
            ),
        ],
    ) -> dict[str, Any]:
        try:
            payload = ctx.workspace.read_bytes(document.path)  # type: ignore[attr-defined]
        except Exception:  # noqa: BLE001
            return {
                "ok": False,
                "code": "upload_read_failed",
                "comparison_id": comparison_id,
                "message": "Could not read the uploaded file from the invocation workspace.",
            }
        parsed = _parse_uploaded_quotes(
            filename=document.filename,
            media_type=document.media_type,
            data=payload,
        )
        if not parsed["ok"]:
            return {**parsed, "comparison_id": comparison_id}
        return await self.compare_quotes(ctx, comparison_id, parsed["quotes"], weights)

    @a2a.tool(
        description="Browser JSON bridge: accept bounded base64 quote files, parse them, compare quotes, and persist the result.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="deterministic",
    )
    async def compare_quotes_browser_upload(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        weights: QuoteWeights,
        documents: list[BrowserQuoteDocument],
    ) -> dict[str, Any]:
        if not documents or len(documents) > MAX_BROWSER_UPLOADS:
            return {
                "ok": False,
                "code": "invalid_upload_count",
                "comparison_id": comparison_id,
                "message": f"Upload 1 to {MAX_BROWSER_UPLOADS} files.",
            }
        all_quotes: list[QuoteInput] = []
        for document in documents:
            try:
                data = base64.b64decode(document.data_base64, validate=True)
            except Exception:  # noqa: BLE001
                return {
                    "ok": False,
                    "code": "invalid_base64_upload",
                    "comparison_id": comparison_id,
                    "message": f"{document.filename} is not valid base64.",
                }
            if len(data) > MAX_UPLOAD_BYTES:
                return {
                    "ok": False,
                    "code": "upload_too_large",
                    "comparison_id": comparison_id,
                    "message": f"{document.filename} exceeds {MAX_UPLOAD_BYTES} bytes.",
                }
            parsed = _parse_uploaded_quotes(
                filename=document.filename,
                media_type=document.media_type,
                data=data,
            )
            if not parsed["ok"]:
                return {**parsed, "comparison_id": comparison_id}
            all_quotes.extend(parsed["quotes"])
        return await self.compare_quotes(ctx, comparison_id, all_quotes, weights)


# ---------- deterministic business logic ----------


def _clean_comparison_id(value: str) -> str | None:
    cleaned = str(value or "").strip()
    if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.:-]{0,119}", cleaned):
        return None
    return cleaned


def _validate_request(
    comparison_id: str,
    quotes: list[QuoteInput],
    weights: QuoteWeights,
) -> dict[str, Any] | None:
    clean_id = _clean_comparison_id(comparison_id)
    if clean_id is None:
        return {
            "ok": False,
            "code": "invalid_comparison_id",
            "comparison_id": comparison_id,
            "message": "comparison_id must be 1-120 URL-safe characters",
        }
    if len(quotes) < 2:
        return {
            "ok": False,
            "code": "at_least_two_quotes_required",
            "comparison_id": clean_id,
            "message": "Compare at least two vendor quotes.",
        }
    if len(quotes) > MAX_QUOTES:
        return {
            "ok": False,
            "code": "too_many_quotes",
            "comparison_id": clean_id,
            "message": f"Compare at most {MAX_QUOTES} quotes at a time.",
        }
    if (weights.price + weights.delivery + weights.warranty) <= 0:
        return {
            "ok": False,
            "code": "weights_total_zero",
            "comparison_id": clean_id,
            "message": "At least one weight must be greater than zero.",
        }
    return None


def _score_quotes(
    comparison_id: str,
    quotes: list[QuoteInput],
    weights: QuoteWeights,
) -> dict[str, Any]:
    clean_id = _clean_comparison_id(comparison_id) or comparison_id
    min_unit_price = min(q.unit_price for q in quotes)
    min_delivery = min(q.delivery_days for q in quotes)
    max_warranty = max(q.warranty_months for q in quotes) or 1
    weight_total = weights.price + weights.delivery + weights.warranty

    scored: list[dict[str, Any]] = []
    for quote in quotes:
        price_score = min_unit_price / quote.unit_price if quote.unit_price else 0
        delivery_score = min_delivery / quote.delivery_days if quote.delivery_days else 0
        warranty_score = quote.warranty_months / max_warranty if max_warranty else 0
        weighted_score = (
            price_score * weights.price
            + delivery_score * weights.delivery
            + warranty_score * weights.warranty
        ) / weight_total * 100
        total_price = Decimal(str(quote.unit_price)) * Decimal(quote.quantity)
        scored.append(
            {
                "vendor": quote.vendor,
                "unit_price": _money(quote.unit_price),
                "quantity": quote.quantity,
                "total_price": _money(total_price),
                "delivery_days": quote.delivery_days,
                "warranty_months": quote.warranty_months,
                "score": _round(weighted_score),
                "score_breakdown": {
                    "price": _round(price_score * 100),
                    "delivery": _round(delivery_score * 100),
                    "warranty": _round(warranty_score * 100),
                },
            }
        )
    ranked = sorted(
        scored,
        key=lambda row: (
            -row["score"],
            row["total_price"],
            row["delivery_days"],
            -row["warranty_months"],
            row["vendor"].lower(),
        ),
    )
    winner = ranked[0]
    result = {
        "ok": True,
        "comparison_id": clean_id,
        "recommendation": {
            "vendor": winner["vendor"],
            "score": winner["score"],
            "total_price": winner["total_price"],
            "delivery_days": winner["delivery_days"],
            "warranty_months": winner["warranty_months"],
            "rationale": (
                f"{winner['vendor']} has the highest weighted score using "
                f"price={weights.price:g}, delivery={weights.delivery:g}, warranty={weights.warranty:g}."
            ),
        },
        "weights": {
            "price": weights.price,
            "delivery": weights.delivery,
            "warranty": weights.warranty,
        },
        "quotes": scored,
        "ranked_quotes": ranked,
        "saved": True,
        "version": VERSION,
        "generated_at": datetime.now(timezone.utc).isoformat(),
    }
    result["execution_receipt"] = _receipt_payload("compare_quotes", clean_id, result)
    return result


def _money(value: Any) -> float:
    return float(Decimal(str(value)).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP))


def _round(value: Any) -> float:
    return float(Decimal(str(value)).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP))


def _receipt_payload(skill_name: str, comparison_id: str, payload: dict[str, Any]) -> dict[str, Any]:
    public_payload = {k: v for k, v in payload.items() if k != "execution_receipt"}
    digest = hashlib.sha256(
        json.dumps(public_payload, sort_keys=True, default=str).encode("utf-8")
    ).hexdigest()
    return {
        "agent": "quote-judge-studio-v1",
        "agent_version": VERSION,
        "skill": skill_name,
        "comparison_id": comparison_id,
        "payload_sha256": digest,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "persisted": True,
    }


# ---------- uploads ----------


def _parse_uploaded_quotes(filename: str, media_type: str, data: bytes) -> dict[str, Any]:
    safe_media_type = media_type.split(";", 1)[0].strip().lower()
    if safe_media_type not in ALLOWED_UPLOAD_MEDIA_TYPES:
        return {"ok": False, "code": "unsupported_upload_media_type"}
    if len(data) > MAX_UPLOAD_BYTES:
        return {"ok": False, "code": "upload_too_large"}
    text = data.decode("utf-8-sig", errors="replace")
    try:
        if safe_media_type == "application/json" or filename.lower().endswith(".json"):
            raw = json.loads(text)
            rows = raw.get("quotes", raw) if isinstance(raw, dict) else raw
            if not isinstance(rows, list):
                raise ValueError("JSON upload must be a list or {'quotes': [...]} object")
            quotes = [QuoteInput.model_validate(row) for row in rows]
        else:
            quotes = _parse_csv_or_text(text)
    except Exception as exc:  # noqa: BLE001
        return {
            "ok": False,
            "code": "quote_parse_failed",
            "message": f"Could not parse quote upload: {str(exc)[:160]}",
        }
    return {"ok": True, "quotes": quotes}


def _parse_csv_or_text(text: str) -> list[QuoteInput]:
    sample = text.strip()
    if not sample:
        raise ValueError("file is empty")
    reader = csv.DictReader(io.StringIO(sample))
    required = {"vendor", "unit_price", "quantity", "delivery_days", "warranty_months"}
    if reader.fieldnames and required.issubset({name.strip() for name in reader.fieldnames}):
        return [QuoteInput.model_validate({k.strip(): v for k, v in row.items()}) for row in reader]
    raise ValueError(
        "CSV must include vendor, unit_price, quantity, delivery_days, warranty_months columns"
    )


# ---------- tenant and database ----------


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


def _database_url() -> str:
    url = os.environ.get("DATABASE_URL", "").strip()
    if not url:
        raise RuntimeError("Managed Postgres is not configured; DATABASE_URL is missing.")
    return url


def _connect():
    import psycopg
    from psycopg.rows import dict_row

    return psycopg.connect(
        _database_url(),
        connect_timeout=DB_CONNECT_TIMEOUT_SECONDS,
        row_factory=dict_row,
        options=f"-c statement_timeout={DB_STATEMENT_TIMEOUT_MS} -c lock_timeout={DB_LOCK_TIMEOUT_MS}",
    )


async def _save_comparison(tenant_key: str, result: dict[str, Any]) -> None:
    import asyncio

    await asyncio.to_thread(_save_comparison_sync, tenant_key, result)


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


async def _load_comparison(tenant_key: str, comparison_id: str) -> dict[str, Any] | None:
    import asyncio

    return await asyncio.to_thread(_load_comparison_sync, tenant_key, comparison_id)


def _load_comparison_sync(tenant_key: str, comparison_id: str) -> dict[str, Any] | None:
    with _connect() as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT payload
                  FROM quote_comparisons
                 WHERE tenant_key = %s AND comparison_id = %s
                 LIMIT 1
                """,
                (tenant_key, comparison_id),
            )
            row = cur.fetchone()
    if not row:
        return None
    payload = row["payload"]
    if isinstance(payload, str):
        return json.loads(payload)
    return dict(payload)


async def _persist_receipt_best_effort(
    ctx: RunContext[PlatformUserAuth],
    skill_name: str,
    comparison_id: str,
    status: str,
    payload: dict[str, Any],
) -> None:
    import asyncio

    try:
        tenant = _tenant_key(ctx)
        receipt = _receipt_payload(skill_name, comparison_id, payload)
        await asyncio.to_thread(
            _persist_receipt_sync,
            tenant,
            comparison_id,
            skill_name,
            status,
            getattr(ctx, "task_id", "") or "",
            receipt,
        )
    except Exception:  # noqa: BLE001
        # Persistence of the comparison itself fails closed; duplicate receipt
        # persistence must not mask the primary skill result.
        return


def _persist_receipt_sync(
    tenant_key: str,
    comparison_id: str,
    skill_name: str,
    status: str,
    task_id: str,
    receipt: dict[str, Any],
) -> None:
    with _connect() as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO quote_execution_receipts (
                    tenant_key, comparison_id, skill_name, status, task_id, receipt
                ) VALUES (%s, %s, %s, %s, %s, %s::jsonb)
                """,
                (tenant_key, comparison_id, skill_name, status, task_id, json.dumps(receipt, default=str)),
            )
        conn.commit()


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


def _temporary_db_error(comparison_id: str) -> dict[str, Any]:
    return {
        "ok": False,
        "code": "database_unavailable",
        "comparison_id": comparison_id,
        "message": "QuoteJudge could not persist or read the comparison before the database timeout.",
    }