import { z } from 'zod';

export const BudgetScopeTypeSchema = z.enum(['product', 'agent', 'org', 'workspace']);
export const BudgetPolicyStatusSchema = z.enum(['active', 'paused', 'archived']);
export const BudgetPeriodSchema = z.enum(['daily', 'monthly']);
export const BudgetAlertSeveritySchema = z.enum(['warn', 'block']);
export const BudgetVerdictSchema = z.enum(['allow', 'warn', 'block']);

export const BudgetPolicySchema = z.object({
  id: z.string().min(1),
  productId: z.string().min(1),
  scopeType: BudgetScopeTypeSchema,
  scopeId: z.string().min(1),
  name: z.string().min(1),
  status: BudgetPolicyStatusSchema,
  period: BudgetPeriodSchema,
  budgetUsd: z.number().min(0),
  softThreshold: z.number().min(0).max(1),
  hardThreshold: z.number().min(0).max(1),
  modelAllowlist: z.array(z.string()).default([]),
  metadata: z.record(z.unknown()).optional(),
  createdBy: z.string().min(1),
  createdAt: z.string(),
  updatedAt: z.string(),
});

export type BudgetPolicyDoc = z.infer<typeof BudgetPolicySchema> & {
  _ts?: number;
  _etag?: string;
};

export const BudgetSpendEntrySchema = z.object({
  id: z.string().min(1),
  productId: z.string().min(1),
  policyId: z.string().optional(),
  scopeType: BudgetScopeTypeSchema,
  scopeId: z.string().min(1),
  agentId: z.string().optional(),
  agentVersionId: z.string().optional(),
  runId: z.string().optional(),
  evaluationRunId: z.string().optional(),
  model: z.string().optional(),
  tokensUsed: z.number().int().min(0).default(0),
  costUsd: z.number().min(0),
  source: z.string().optional(),
  verdict: BudgetVerdictSchema,
  recordedAt: z.string(),
});

export type BudgetSpendEntryDoc = z.infer<typeof BudgetSpendEntrySchema> & {
  _ts?: number;
  _etag?: string;
};

export const BudgetAlertSchema = z.object({
  id: z.string().min(1),
  productId: z.string().min(1),
  policyId: z.string().min(1),
  scopeType: BudgetScopeTypeSchema,
  scopeId: z.string().min(1),
  severity: BudgetAlertSeveritySchema,
  message: z.string().min(1),
  spendUsd: z.number().min(0),
  budgetUsd: z.number().min(0),
  percentUsed: z.number().min(0),
  createdAt: z.string(),
});

export type BudgetAlertDoc = z.infer<typeof BudgetAlertSchema> & {
  _ts?: number;
  _etag?: string;
};

export const CreateBudgetPolicySchema = z.object({
  scopeType: BudgetScopeTypeSchema,
  scopeId: z.string().min(1),
  name: z.string().min(1),
  period: BudgetPeriodSchema.default('monthly'),
  budgetUsd: z.number().min(0.01),
  softThreshold: z.number().min(0).max(1).default(0.8),
  hardThreshold: z.number().min(0).max(1).default(1),
  modelAllowlist: z.array(z.string()).default([]),
  metadata: z.record(z.unknown()).optional(),
});

export const UpdateBudgetPolicySchema = z.object({
  name: z.string().min(1).optional(),
  status: BudgetPolicyStatusSchema.optional(),
  period: BudgetPeriodSchema.optional(),
  budgetUsd: z.number().min(0.01).optional(),
  softThreshold: z.number().min(0).max(1).optional(),
  hardThreshold: z.number().min(0).max(1).optional(),
  modelAllowlist: z.array(z.string()).optional(),
  metadata: z.record(z.unknown()).optional(),
});

export const RecordBudgetSpendSchema = z.object({
  scopeType: BudgetScopeTypeSchema,
  scopeId: z.string().min(1),
  policyId: z.string().min(1).optional(),
  agentId: z.string().optional(),
  agentVersionId: z.string().optional(),
  runId: z.string().optional(),
  evaluationRunId: z.string().optional(),
  model: z.string().optional(),
  tokensUsed: z.number().int().min(0).default(0),
  costUsd: z.number().min(0),
  source: z.string().optional(),
});

export const ListBudgetPoliciesQuerySchema = z.object({
  scopeType: BudgetScopeTypeSchema.optional(),
  scopeId: z.string().min(1).optional(),
  status: BudgetPolicyStatusSchema.optional(),
  limit: z.coerce.number().min(1).max(100).default(20),
});

export const ListBudgetAlertsQuerySchema = z.object({
  severity: BudgetAlertSeveritySchema.optional(),
  limit: z.coerce.number().min(1).max(100).default(50),
});

export type ListBudgetPoliciesQuery = z.infer<typeof ListBudgetPoliciesQuerySchema>;
export type ListBudgetAlertsQuery = z.infer<typeof ListBudgetAlertsQuerySchema>;

// ── Cost Dashboard ──────────────────────────────────────────

export const CostDashboardQuerySchema = z.object({
  scopeType: BudgetScopeTypeSchema.optional(),
  scopeId: z.string().optional(),
  since: z.string().optional(),
  until: z.string().optional(),
  groupBy: z.enum(['model', 'agent', 'day', 'scope']).default('model'),
  limit: z.coerce.number().min(1).max(500).default(100),
});

export type CostDashboardQuery = z.infer<typeof CostDashboardQuerySchema>;

// ── Budget Rollover ─────────────────────────────────────────

export const BudgetRolloverSchema = z.object({
  id: z.string().min(1),
  productId: z.string().min(1),
  policyId: z.string().min(1),
  fromPeriod: z.string().min(1),
  toPeriod: z.string().min(1),
  spentUsd: z.number().min(0),
  budgetUsd: z.number().min(0),
  remainingUsd: z.number(),
  rolledOverUsd: z.number().min(0),
  createdAt: z.string(),
});

export type BudgetRolloverDoc = z.infer<typeof BudgetRolloverSchema> & {
  _ts?: number;
  _etag?: string;
};
