quote-judge-studio-v1/agent.py

"""QuoteJudge Studio v1: deterministic quote comparison with durable per-user state."""
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,
    State,
    WorkspaceAccess,
    WorkspaceMode,
)
from a2a_pack.workspace import FileUpload, UploadedFile

try:  # psycopg is installed in the deployed image via requirements.txt.
    import psycopg
    from psycopg.rows import dict_row
    from psycopg.types.json import Jsonb
except Exception:  # pragma: no cover - lets `a2a card` run before deps install.
    psycopg = None
    dict_row = None
    Jsonb = None


COMPARISON_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,79}$")
MAX_QUOTES = 10
MAX_BROWSER_DOCUMENTS = 4
MAX_UPLOAD_BYTES = 64 * 1024
ALLOWED_UPLOAD_MEDIA_TYPES = {
    "application/json",
    "text/json",
    "text/csv",
    "application/csv",
    "application/vnd.ms-excel",
    "text/plain",
}

# Local/dev fallback only. Hosted production uses DATABASE_URL and managed Postgres.
_LOCAL_COMPARISONS: dict[tuple[str, str], dict[str, Any]] = {}
_LOCAL_RECEIPTS: list[dict[str, Any]] = []


class QuoteJudgeStudioV1Config(BaseModel):
    """Runtime config is intentionally empty; the app uses platform auth + DB."""


class QuoteJudgeState(BaseModel):
    """Card-visible durable state marker; data lives in managed Postgres."""


class QuoteInput(BaseModel):
    vendor: str = Field(..., min_length=1, max_length=120)
    quantity: int = Field(..., ge=1, le=1_000_000)
    unit_price: float = Field(..., ge=0, le=1_000_000_000)
    delivery_days: int = Field(..., ge=0, le=3650)
    warranty_months: int = Field(..., ge=0, le=240)

    @field_validator("vendor")
    @classmethod
    def clean_vendor(cls, value: str) -> str:
        cleaned = " ".join(str(value or "").split())
        if not cleaned:
            raise ValueError("vendor is required")
        return cleaned


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

    @field_validator("warranty")
    @classmethod
    def at_least_one_weight(cls, value: float, info: Any) -> float:
        values = dict(info.data)
        total = float(values.get("price") or 0) + float(values.get("delivery") or 0) + float(value or 0)
        if total <= 0:
            raise ValueError("at least one weight must be greater than zero")
        return value


class BrowserDocument(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=90_000)

    @field_validator("filename")
    @classmethod
    def safe_filename(cls, value: str) -> str:
        cleaned = str(value or "").strip().replace("\\", "/")
        if not cleaned or "/" in cleaned or cleaned in {".", ".."}:
            raise ValueError("filename must be a simple file name")
        return cleaned

    @field_validator("media_type")
    @classmethod
    def known_media_type(cls, value: str) -> str:
        cleaned = str(value or "").split(";", 1)[0].strip().lower()
        if cleaned not in ALLOWED_UPLOAD_MEDIA_TYPES:
            raise ValueError("unsupported media type")
        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 per-user comparisons, supports browser uploads, and serves "
        "a polished packed product UI."
    )
    version = "0.1.2"

    config_model = QuoteJudgeStudioV1Config
    auth_model = PlatformUserAuth

    state = State.DURABLE
    state_model = QuoteJudgeState
    pricing = Pricing(
        price_per_call_usd=0.0,
        caller_pays_llm=False,
        notes="Deterministic quote scoring; no LLM credential or provider key is required.",
    )
    resources = Resources(cpu="500m", memory="512Mi", max_runtime_seconds=600)
    tools_used = ("postgres", "mcp", "packed-react-frontend")
    workspace_access = WorkspaceAccess.dynamic(
        max_files=8,
        allowed_modes=(WorkspaceMode.READ_ONLY,),
        require_reason=False,
        max_total_size_bytes=MAX_UPLOAD_BYTES * MAX_BROWSER_DOCUMENTS,
    )
    platform_resources = AgentPlatformResources(
        databases=(
            AgentDatabase(
                name="quote-judge-studio-v1-data",
                scope="user",
                access_mode="read_write",
                env=AgentDatabaseEnv(url="DATABASE_URL"),
                migrations=AgentDatabaseMigrations(path="db/migrations"),
            ),
        )
    )

    @a2a.tool(
        description="Compare two to ten vendor quotes with weighted price, delivery, and warranty scoring; persist the result and a production receipt.",
        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]:
        tenant = _tenant_key(ctx)
        validation = _validate_comparison_request(comparison_id, quotes, weights)
        if validation is not None:
            return validation

        result = _score_quotes(comparison_id, quotes, weights)
        receipt = _build_receipt(
            tenant_key=tenant,
            skill_name="compare_quotes",
            comparison_id=comparison_id,
            inputs={
                "comparison_id": comparison_id,
                "quotes": [quote.model_dump(mode="json") for quote in quotes],
                "weights": weights.model_dump(mode="json"),
            },
            result=result,
        )
        persisted = _persist_comparison(tenant, comparison_id, result, receipt)
        result["receipt"] = {
            "receipt_id": receipt["receipt_id"],
            "persisted": persisted["ok"],
            "created_at": receipt["created_at"],
        }
        if persisted.get("warning"):
            result.setdefault("warnings", []).append(persisted["warning"])
        await ctx.emit_progress(f"saved comparison {comparison_id}")
        return result

    @a2a.tool(
        description="Reopen a saved quote comparison for the authenticated 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)
        if not _valid_comparison_id(comparison_id):
            return {"ok": False, "code": "invalid_comparison_id", "comparison_id": comparison_id}
        record = _load_comparison(tenant, comparison_id)
        if record is None:
            return {"ok": False, "code": "comparison_not_found", "comparison_id": comparison_id}
        result = dict(record["result"])
        result["comparison_id"] = comparison_id
        result["ok"] = True
        result["reloaded"] = True
        result["saved_at"] = record.get("updated_at")
        if record.get("latest_receipt_id"):
            result["latest_receipt_id"] = record["latest_receipt_id"]
        return result

    @a2a.tool(
        description="Browser-safe upload bridge: accept bounded base64 JSON/CSV/text quote documents, parse them, compare, persist, and receipt the result.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="deterministic",
    )
    async def compare_quotes_from_browser_upload(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        documents: list[BrowserDocument],
        weights: QuoteWeights,
    ) -> dict[str, Any]:
        parsed = _decode_browser_documents(documents)
        if not parsed["ok"]:
            return {"ok": False, "comparison_id": comparison_id, **parsed}
        quotes = parsed["quotes"]
        result = await self.compare_quotes(ctx, comparison_id, quotes, weights)
        result["source"] = "browser_upload"
        result["parsed_documents"] = parsed["documents"]
        return result

    @a2a.tool(
        description="External Agent API upload path: accept a typed FileUpload JSON/CSV/text document, parse quotes, compare, persist, and receipt the result.",
        timeout_seconds=30,
        idempotent=True,
        cost_class="deterministic",
    )
    async def compare_quotes_from_file(
        self,
        ctx: RunContext[PlatformUserAuth],
        comparison_id: str,
        document: Annotated[
            UploadedFile,
            FileUpload(
                accept=sorted(ALLOWED_UPLOAD_MEDIA_TYPES),
                max_bytes=MAX_UPLOAD_BYTES,
                description="JSON or CSV quote file with vendor, quantity, unit_price, delivery_days, warranty_months.",
            ),
        ],
        weights: QuoteWeights,
    ) -> dict[str, Any]:
        if document.size_bytes > MAX_UPLOAD_BYTES:
            return {"ok": False, "code": "file_too_large", "comparison_id": comparison_id}
        media_type = (document.media_type or "").split(";", 1)[0].strip().lower()
        if media_type not in ALLOWED_UPLOAD_MEDIA_TYPES:
            return {"ok": False, "code": "unsupported_media_type", "comparison_id": comparison_id}
        try:
            reader = getattr(ctx.workspace, "read_bytes", None)
            if reader is None:
                return {"ok": False, "code": "workspace_read_unavailable", "comparison_id": comparison_id}
            data = reader(document.path)
        except Exception:
            return {"ok": False, "code": "file_read_failed", "comparison_id": comparison_id}
        parsed = _parse_quote_document(data, filename=document.filename, media_type=media_type)
        if not parsed["ok"]:
            return {"ok": False, "comparison_id": comparison_id, **parsed}
        result = await self.compare_quotes(ctx, comparison_id, parsed["quotes"], weights)
        result["source"] = "file_upload"
        result["parsed_documents"] = [{"filename": document.filename, "quotes": len(parsed["quotes"])}]
        return result


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)
    return f"user:{stable}"


def _valid_comparison_id(value: str) -> bool:
    return bool(COMPARISON_ID_RE.fullmatch(str(value or "").strip()))


def _validate_comparison_request(
    comparison_id: str,
    quotes: list[QuoteInput],
    weights: QuoteWeights,
) -> dict[str, Any] | None:
    if not _valid_comparison_id(comparison_id):
        return {"ok": False, "code": "invalid_comparison_id", "comparison_id": comparison_id}
    if len(quotes) < 2:
        return {"ok": False, "code": "at_least_two_quotes_required", "comparison_id": comparison_id}
    if len(quotes) > MAX_QUOTES:
        return {"ok": False, "code": "too_many_quotes", "comparison_id": comparison_id}
    vendor_names = [quote.vendor.casefold() for quote in quotes]
    if len(set(vendor_names)) != len(vendor_names):
        return {"ok": False, "code": "duplicate_vendor", "comparison_id": comparison_id}
    if weights.price + weights.delivery + weights.warranty <= 0:
        return {"ok": False, "code": "weights_sum_to_zero", "comparison_id": comparison_id}
    return None


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


def _score_quotes(comparison_id: str, quotes: list[QuoteInput], weights: QuoteWeights) -> dict[str, Any]:
    total_weight = Decimal(str(weights.price + weights.delivery + weights.warranty))
    prices = [_money(q.unit_price) for q in quotes]
    totals = [_money(q.unit_price * q.quantity) for q in quotes]
    deliveries = [Decimal(q.delivery_days) for q in quotes]
    warranties = [Decimal(q.warranty_months) for q in quotes]

    rows: list[dict[str, Any]] = []
    for index, quote in enumerate(quotes):
        price_score = _lower_is_better(totals[index], totals)
        delivery_score = _lower_is_better(deliveries[index], deliveries)
        warranty_score = _higher_is_better(warranties[index], warranties)
        weighted_score = (
            price_score * Decimal(str(weights.price))
            + delivery_score * Decimal(str(weights.delivery))
            + warranty_score * Decimal(str(weights.warranty))
        ) / total_weight
        rows.append(
            {
                "vendor": quote.vendor,
                "quantity": quote.quantity,
                "unit_price": _float2(prices[index]),
                "total_price": _float2(totals[index]),
                "delivery_days": quote.delivery_days,
                "warranty_months": quote.warranty_months,
                "score": _float2(weighted_score),
                "score_breakdown": {
                    "price": _float2(price_score),
                    "delivery": _float2(delivery_score),
                    "warranty": _float2(warranty_score),
                },
            }
        )
    rows.sort(key=_rank_sort_key)
    for rank, row in enumerate(rows, start=1):
        row["rank"] = rank
    recommendation = rows[0]
    return {
        "ok": True,
        "comparison_id": comparison_id,
        "recommendation": {
            "vendor": recommendation["vendor"],
            "score": recommendation["score"],
            "total_price": recommendation["total_price"],
            "delivery_days": recommendation["delivery_days"],
            "warranty_months": recommendation["warranty_months"],
            "rationale": (
                f"{recommendation['vendor']} has the best weighted score using "
                f"price={weights.price:g}, delivery={weights.delivery:g}, warranty={weights.warranty:g}."
            ),
        },
        "rankings": rows,
        "weights": weights.model_dump(mode="json"),
        "quote_count": len(quotes),
    }


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


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


def _lower_is_better(value: Decimal, values: list[Decimal]) -> Decimal:
    low, high = min(values), max(values)
    if low == high:
        return Decimal("100")
    return ((high - value) / (high - low) * Decimal("100")).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP)


def _higher_is_better(value: Decimal, values: list[Decimal]) -> Decimal:
    low, high = min(values), max(values)
    if low == high:
        return Decimal("100")
    return ((value - low) / (high - low) * Decimal("100")).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP)


def _decode_browser_documents(documents: list[BrowserDocument]) -> dict[str, Any]:
    if not documents:
        return {"ok": False, "code": "no_documents_provided"}
    if len(documents) > MAX_BROWSER_DOCUMENTS:
        return {"ok": False, "code": "too_many_documents"}
    all_quotes: list[QuoteInput] = []
    summaries: list[dict[str, Any]] = []
    for doc in documents:
        try:
            data = base64.b64decode(doc.data_base64.encode("ascii"), validate=True)
        except Exception:
            return {"ok": False, "code": "invalid_base64", "filename": doc.filename}
        if len(data) > MAX_UPLOAD_BYTES:
            return {"ok": False, "code": "file_too_large", "filename": doc.filename}
        parsed = _parse_quote_document(data, filename=doc.filename, media_type=doc.media_type)
        if not parsed["ok"]:
            return {"ok": False, "filename": doc.filename, **parsed}
        all_quotes.extend(parsed["quotes"])
        summaries.append({"filename": doc.filename, "quotes": len(parsed["quotes"])})
    return {"ok": True, "quotes": all_quotes, "documents": summaries}


def _parse_quote_document(data: bytes, *, filename: str, media_type: str) -> dict[str, Any]:
    if len(data) > MAX_UPLOAD_BYTES:
        return {"ok": False, "code": "file_too_large"}
    text = data.decode("utf-8-sig", errors="replace").strip()
    if not text:
        return {"ok": False, "code": "empty_document"}
    try:
        if media_type in {"application/json", "text/json"} or filename.lower().endswith(".json"):
            raw = json.loads(text)
            items = raw.get("quotes", raw) if isinstance(raw, dict) else raw
            if not isinstance(items, list):
                return {"ok": False, "code": "quotes_array_required"}
            return {"ok": True, "quotes": [QuoteInput.model_validate(item) for item in items]}
        rows = list(csv.DictReader(io.StringIO(text)))
        if rows and rows[0]:
            return {"ok": True, "quotes": [_quote_from_row(row) for row in rows]}
        return {"ok": False, "code": "unsupported_document_format"}
    except Exception:
        return {"ok": False, "code": "document_parse_failed"}


def _quote_from_row(row: dict[str, Any]) -> QuoteInput:
    normalized = {str(key).strip().lower(): value for key, value in row.items()}
    return QuoteInput(
        vendor=str(normalized.get("vendor") or normalized.get("supplier") or ""),
        quantity=int(float(normalized.get("quantity") or normalized.get("qty") or 0)),
        unit_price=float(normalized.get("unit_price") or normalized.get("price") or 0),
        delivery_days=int(float(normalized.get("delivery_days") or normalized.get("delivery") or 0)),
        warranty_months=int(float(normalized.get("warranty_months") or normalized.get("warranty") or 0)),
    )


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


def _connect() -> Any:
    if psycopg is None:
        raise RuntimeError("psycopg is not installed")
    url = _database_url()
    if not url:
        raise RuntimeError("DATABASE_URL is not configured")
    return psycopg.connect(url, row_factory=dict_row)


def _persist_comparison(
    tenant_key: str,
    comparison_id: str,
    result: dict[str, Any],
    receipt: dict[str, Any],
) -> dict[str, Any]:
    if not _database_url() or psycopg is None:
        _LOCAL_COMPARISONS[(tenant_key, comparison_id)] = {
            "result": json.loads(json.dumps(result)),
            "updated_at": receipt["created_at"],
            "latest_receipt_id": receipt["receipt_id"],
        }
        _LOCAL_RECEIPTS.append(receipt)
        return {"ok": True, "warning": "local_ephemeral_persistence"}
    try:
        with _connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO quote_judge_comparisons
                        (tenant_key, comparison_id, input_json, result_json, recommendation_vendor, latest_receipt_id, updated_at)
                    VALUES (%s, %s, %s, %s, %s, %s, NOW())
                    ON CONFLICT (tenant_key, comparison_id) DO UPDATE SET
                        input_json = EXCLUDED.input_json,
                        result_json = EXCLUDED.result_json,
                        recommendation_vendor = EXCLUDED.recommendation_vendor,
                        latest_receipt_id = EXCLUDED.latest_receipt_id,
                        updated_at = NOW()
                    """,
                    (
                        tenant_key,
                        comparison_id,
                        Jsonb({"weights": result["weights"], "quote_count": result["quote_count"]}),
                        Jsonb(result),
                        result["recommendation"]["vendor"],
                        receipt["receipt_id"],
                    ),
                )
                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, %s)
                    ON CONFLICT (receipt_id) DO UPDATE SET
                        tenant_key = EXCLUDED.tenant_key,
                        comparison_id = EXCLUDED.comparison_id,
                        skill_name = EXCLUDED.skill_name,
                        tool_name = EXCLUDED.tool_name,
                        status = EXCLUDED.status,
                        input_hash = EXCLUDED.input_hash,
                        result_hash = EXCLUDED.result_hash,
                        receipt_json = EXCLUDED.receipt_json,
                        payload = EXCLUDED.payload
                    """,
                    (
                        receipt["receipt_id"],
                        tenant_key,
                        comparison_id,
                        receipt["skill_name"],
                        receipt["skill_name"],
                        receipt["status"],
                        receipt["input_hash"],
                        receipt["result_hash"],
                        Jsonb(receipt),
                        Jsonb(receipt),
                    ),
                )
            conn.commit()
        return {"ok": True}
    except Exception:
        return {"ok": False, "warning": "persistence_unavailable"}


def _load_comparison(tenant_key: str, comparison_id: str) -> dict[str, Any] | None:
    if not _database_url() or psycopg is None:
        return _LOCAL_COMPARISONS.get((tenant_key, comparison_id))
    try:
        with _connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    SELECT result_json, updated_at, latest_receipt_id
                    FROM quote_judge_comparisons
                    WHERE tenant_key = %s AND comparison_id = %s
                    """,
                    (tenant_key, comparison_id),
                )
                row = cur.fetchone()
        if not row:
            return None
        updated_at = row["updated_at"]
        return {
            "result": row["result_json"],
            "updated_at": updated_at.isoformat() if hasattr(updated_at, "isoformat") else str(updated_at),
            "latest_receipt_id": row.get("latest_receipt_id"),
        }
    except Exception:
        return None


def _build_receipt(
    *,
    tenant_key: str,
    skill_name: str,
    comparison_id: str,
    inputs: dict[str, Any],
    result: dict[str, Any],
) -> dict[str, Any]:
    now = datetime.now(timezone.utc).isoformat()
    tenant_hash = _hash_json({"tenant_key": tenant_key})
    input_hash = _hash_json(inputs)
    result_hash = _hash_json(result)
    receipt_id = "qjr_" + _hash_json({
        "tenant_hash": tenant_hash,
        "skill_name": skill_name,
        "comparison_id": comparison_id,
        "input_hash": input_hash,
        "result_hash": result_hash,
    })[:24]
    return {
        "schema": "quotejudge.production_receipt.v1",
        "receipt_id": receipt_id,
        "agent_name": "quote-judge-studio-v1",
        "agent_version": "0.1.0",
        "skill_name": skill_name,
        "comparison_id": comparison_id,
        "status": "ok" if result.get("ok") else "error",
        "tenant_hash": tenant_hash,
        "input_hash": input_hash,
        "result_hash": result_hash,
        "created_at": now,
    }


def _hash_json(value: Any) -> str:
    encoded = json.dumps(value, sort_keys=True, separators=(",", ":"), default=str).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()