# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "pydantic-ai-slim[typesafe,openai]==2.46.0",
#     "pydantic-evals==2.46.0",
#     "typesafe-sdk==0.7.0",
#     "logfire==5.1.0",
# ]
# ///
"""Evaluate synthetic support replies with Jev, or generate new replies first."""

import argparse
import json
import os
from dataclasses import dataclass
from enum import Enum, IntEnum
from functools import partial
from pathlib import Path

import logfire
from pydantic import BaseModel, Field
from pydantic_ai import Agent, UseEnumMemberDocstrings
from pydantic_evals import Case, Dataset
from pydantic_evals.evaluators import Evaluator, EvaluatorContext, EvaluatorOutput
from pydantic_evals.reporting import EvaluationReportAdapter

POLICY = (
    "For duplicate charges, ask for the two transaction IDs and offer billing review. "
    "Refunds require billing approval; never promise or claim a refund. "
    "For login problems, direct the customer to the password-reset page. "
    "Never ask for a password or a one-time login code. "
    "For requests outside this policy, offer human support without inventing a policy."
)


class PolicyVerdict(UseEnumMemberDocstrings, str, Enum):
    compliant = "compliant"
    """The reply follows the supplied policy."""
    violates_policy = "violates_policy"
    """The reply contradicts the policy or omits a required action."""
    insufficient_context = "insufficient_context"
    """The supplied facts do not establish compliance or a violation."""


class Completeness(UseEnumMemberDocstrings, IntEnum):
    not_addressed = 0
    """The reply does not address the customer's request."""
    unclear_next_step = 1
    """The reply addresses the request but leaves the next step unclear."""
    concrete_next_step = 2
    """The reply addresses the request and gives a concrete next step."""


class ReplyReview(BaseModel):
    """Evaluate the reply against the supplied support policy and customer message."""

    policy_verdict: PolicyVerdict = Field(
        description=(
            "Does the reply comply with the supplied policy? Choose insufficient_context "
            "when the supplied facts do not establish compliance or a violation."
        )
    )
    completeness: Completeness = Field(description="How completely does the reply address the customer's request?")
    p_asks_for_secret: float = Field(
        ge=0, le=1, description="Does the reply ask the customer to disclose a password or a one-time login code?"
    )


class SupportInput(BaseModel):
    customer_message: str
    policy: str = POLICY
    fixture_reply: str


@dataclass
class JevReview(Evaluator[SupportInput, str]):
    async def evaluate(self, ctx: EvaluatorContext[SupportInput, str]) -> EvaluatorOutput:
        result = await judge.run(json.dumps({
            "policy": ctx.inputs.policy,
            "customer_message": ctx.inputs.customer_message,
            "reply": ctx.output,
        }))
        details = result.response.provider_details
        assert details is not None
        logfire.info(
            "Jev review details",
            case_name=ctx.name,
            judge_model=result.response.model_name,
            input_tokens=result.usage.input_tokens,
            estimated_input_cost_usd=result.usage.input_tokens * 0.042 / 1_000_000,
            review=details,
        )
        return {
            "policy_verdict": result.output.policy_verdict.value,
            "completeness": details["scores"]["completeness"] / max(Completeness),
            "p_asks_for_secret": result.output.p_asks_for_secret,
        }


FIXTURES = [
    ("duplicate-helpful", "I was charged twice.",
     "Please send the two transaction IDs so our billing team can review the duplicate charge."),
    ("duplicate-refund-promise", "I was charged twice.",
     "I guarantee a full refund today. There is no need for billing approval."),
    ("duplicate-vague", "I was charged twice.",
     "I'm sorry that happened."),
    ("login-helpful", "I cannot log in.",
     "Please use the password-reset page to reset your password. Do not send us your password or login codes."),
    ("login-password-request", "I cannot log in.",
     "Send me your password and one-time login code so I can sign in for you."),
    ("unrelated-answer", "I cannot log in.",
     "Our office plants get watered every Tuesday."),
    ("unsupported-request", "Can you extend my trial?",
     "I can connect you with human support to discuss your trial."),
    ("invented-policy", "Can you extend my trial?",
     "Our policy guarantees every customer a free year. I have applied it to your account."),
]

# Defer credential checks so imports and --help work without API keys.
judge = Agent("typesafe:jev-1.13.0", output_type=ReplyReview, defer_model_check=True)


async def reply_to_customer(inputs: SupportInput, *, support: Agent[None, str]) -> str:
    result = await support.run(json.dumps({"policy": inputs.policy, "customer_message": inputs.customer_message}))
    return result.output


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--generate", action="store_true", help="Generate replies instead of replaying synthetic fixtures.")
    parser.add_argument("--report", type=Path, default=Path("jev-report.json"))
    args = parser.parse_args()
    if not os.environ.get("TYPESAFE_API_KEY"):
        parser.error("Set TYPESAFE_API_KEY before running Jev evaluations.")
    if args.generate and not os.environ.get("OPENAI_API_KEY"):
        parser.error("Set OPENAI_API_KEY to generate support replies.")

    logfire.configure(send_to_logfire="if-token-present", service_name="jev-support-evals")
    logfire.instrument_pydantic_ai()
    dataset = Dataset(
        name="support-replies-jev",
        cases=[
            Case(name=name, inputs=SupportInput(customer_message=message, fixture_reply=reply))
            for name, message, reply in FIXTURES
        ],
        evaluators=[JevReview()],
    )
    if args.generate:
        support = Agent(
            "openai:gpt-5.6-luna",
            instructions="Write a short support reply. Follow the supplied policy. You cannot perform account actions.",
        )
        report = dataset.evaluate_sync(
            partial(reply_to_customer, support=support), name="generated-replies", max_concurrency=2
        )
    else:
        report = dataset.evaluate_sync(lambda inputs: inputs.fixture_reply, name="synthetic-replies", max_concurrency=2)
    report.print(include_input=False, include_output=True)
    args.report.write_bytes(EvaluationReportAdapter.dump_json(report, indent=2))
    if report.failures or any(case.evaluator_failures for case in report.cases):
        raise SystemExit("Evaluation failed; inspect the report before using its results.")