quote-judge-studio-v1/agent.py

"""QuoteJudge Studio full-stack A2A agent.

Deterministic quote comparison product with platform-user tenant isolation,
managed Postgres persistence, database timeouts, bounded browser uploads,
typed external file uploads, MCP-compatible tools, and receipt persistence.
"""
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.4"


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

    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:
            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:
            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:
            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(document.filename, document.media_type, 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 one bounded base64 quote file, parse it, 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: BrowserQuoteDocument, weights: QuoteWeights) -> dict[str, Any]:
        return await self.compare_quotes_browser_upload(ctx, comparison_id, weights, [upload])

    @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:
                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(document.filename, document.media_type, 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)


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
        delivery_score = min_delivery / quote.delivery_days
        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)
        rounded_score = _round(weighted_score)
        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": rounded_score,
            "weighted_score": rounded_score,
            "score_breakdown": {"price": _round(price_score * 100), "delivery": _round(delivery_score * 100), "warranty": _round(warranty_score * 100)},
        })
    ranked = sorted(scored, key=_quote_sort_key)
    winner = ranked[0]
    result = {
        "ok": True,
        "comparison_id": clean_id,
        "recommendation": {
            "vendor": winner["vendor"],
            "score": winner["score"],
            "weighted_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 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)
    result["receipt"] = result["execution_receipt"]
    return result


def _quote_sort_key(row: dict[str, Any]) -> tuple[Any, ...]:
    return (-row["score"], row["total_price"], row["delivery_days"], -row["warranty_months"], row["vendor"].lower())


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 not in {"execution_receipt", "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}


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:
        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({key.strip(): value for key, value in csv_row.items()}) for csv_row in reader]
    raise ValueError("CSV must include vendor, unit_price, quantity, delivery_days, warranty_months columns")


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_judge_comparisons (
                    tenant_key, comparison_id, recommendation_vendor, result_json, latest_receipt_id, updated_at
                ) VALUES (%s, %s, %s, %s::jsonb, %s, NOW())
                ON CONFLICT (tenant_key, comparison_id) DO UPDATE SET
                    recommendation_vendor = EXCLUDED.recommendation_vendor,
                    result_json = EXCLUDED.result_json,
                    latest_receipt_id = EXCLUDED.latest_receipt_id,
                    updated_at = NOW()
                """,
                (tenant_key, result["comparison_id"], result["recommendation"]["vendor"], payload, result.get("execution_receipt", {}).get("payload_sha256")),
            )
        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 result_json FROM quote_judge_comparisons WHERE tenant_key = %s AND comparison_id = %s LIMIT 1""", (tenant_key, comparison_id))
            db_row = cur.fetchone()
    if not db_row:
        return None
    payload = db_row["result_json"]
    return json.loads(payload) if isinstance(payload, str) else 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:
        return


def _persist_receipt_sync(tenant_key: str, comparison_id: str, skill_name: str, status: str, task_id: str, receipt: dict[str, Any]) -> None:
    receipt_json = json.dumps(receipt, default=str)
    with _connect() as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO quote_judge_receipts (
                    receipt_id, tenant_key, comparison_id, skill_name, tool_name, status,
                    input_hash, result_hash, receipt_json, payload
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s::jsonb, %s::jsonb)
                """,
                (receipt.get("payload_sha256"), tenant_key, comparison_id, skill_name, skill_name, status, receipt.get("payload_sha256"), receipt.get("payload_sha256"), receipt_json, receipt_json),
            )
        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."}