import { z } from "zod";

/** The 7 mutually exclusive content categories from SEC Release 33-11216. */
export const ContentCategory = z.enum([
  "Board Governance",
  "Management Role",
  "Risk Management Process",
  "Third-Party Risk",
  "Incident Disclosure",
  "Strategy Integration",
  "None/Other",
]);
export type ContentCategory = z.infer<typeof ContentCategory>;

/** Ordinal specificity as semantic labels — we convert to 1-4 in post-processing. */
export const SpecificityLabel = z.enum([
  "Generic Boilerplate",
  "Sector-Adapted",
  "Firm-Specific",
  "Quantified-Verifiable",
]);
export type SpecificityLabel = z.infer<typeof SpecificityLabel>;

/** Map specificity labels to ordinal values. */
export const SPECIFICITY_ORDINAL: Record<SpecificityLabel, number> = {
  "Generic Boilerplate": 1,
  "Sector-Adapted": 2,
  "Firm-Specific": 3,
  "Quantified-Verifiable": 4,
};

// Keep this for downstream compatibility (Annotation schema, consensus, etc.)
export const SpecificityLevel = z.union([
  z.literal(1),
  z.literal(2),
  z.literal(3),
  z.literal(4),
]);
export type SpecificityLevel = z.infer<typeof SpecificityLevel>;

/**
 * Per-dimension confidence for annotator self-assessment.
 * Accepts strings ("high", "High") or numbers (0-1 mapped to low/medium/high).
 */
export const Confidence = z
  .union([z.string(), z.number()])
  .transform((v) => {
    if (typeof v === "number") return v >= 0.8 ? "high" : v >= 0.5 ? "medium" : "low";
    return v.toLowerCase();
  })
  .pipe(z.enum(["high", "medium", "low"]));
export type Confidence = z.infer<typeof Confidence>;

/** Type of specific fact found in the text. */
export const FactType = z.enum([
  "framework",       // Named standard: NIST, ISO 27001, SOC 2, etc.
  "named_role",      // Specific title: CISO, VP of Security, CTO, etc.
  "named_committee", // Named committee: Audit Committee, Technology Committee, etc.
  "named_team",      // Team/dept composition: "Legal, Compliance, and Finance"
  "specific_date",   // Date: "On March 10, 2024", "In December 2023"
  "named_system",    // Product/system/tool: Splunk, CrowdStrike, Dropbox Sign
  "quantified",      // Numbers: dollar amounts, headcounts, percentages, frequencies
  "named_firm",      // Third-party firm: Mandiant, Deloitte, AIG
  "other_specific",  // Other verifiable fact not covered above
]);
export type FactType = z.infer<typeof FactType>;

/** A specific fact extracted from the paragraph text. */
const SpecificFactObject = z.object({
  fact: z.string().describe("The specific fact, quoted or paraphrased from the text"),
  type: FactType.describe("Classification of this fact"),
});

/**
 * Accepts either a {fact, type} object or a plain string (coerced to
 * {fact: string, type: "other_specific"}). Some models return flat strings
 * instead of the structured format.
 */
export const SpecificFact = z
  .union([SpecificFactObject, z.string().transform((s) => ({ fact: s, type: "other_specific" as const }))])
  .pipe(SpecificFactObject);
export type SpecificFact = z.infer<typeof SpecificFact>;

/**
 * Raw schema passed to `generateText` — what the model produces.
 * Uses text labels for specificity to give the model semantic signal.
 * The `specific_facts` array forces chain-of-thought: the model must
 * enumerate evidence BEFORE assigning specificity.
 */
export const LabelOutputRaw = z.object({
  content_category: ContentCategory.describe(
    "The single most applicable content category for this paragraph",
  ),
  specific_facts: z.array(SpecificFact).describe(
    "List ONLY facts from the ✓ IS list (cybersecurity-specific roles, named non-generic committees, team compositions, dates, named programs, products, firms, specific numbers). Do NOT include items from the ✗ NOT list (CEO, CFO, Audit Committee, generic cadences, boilerplate phrases). Empty array if no qualifying facts exist.",
  ),
  specificity: SpecificityLabel.describe(
    "Derived from specific_facts: no facts → Generic Boilerplate, only frameworks → Sector-Adapted, any firm-unique fact → Firm-Specific, 2+ verifiable facts → Quantified-Verifiable",
  ),
  category_confidence: Confidence.describe(
    "high=clear-cut, medium=some ambiguity, low=genuinely torn between categories",
  ),
  specificity_confidence: Confidence.describe(
    "high=clear-cut, medium=borderline adjacent levels, low=could argue for 2+ levels",
  ),
  reasoning: z
    .string()
    .describe(
      "Brief 1-2 sentence justification for the content_category assignment",
    ),
});
export type LabelOutputRaw = z.infer<typeof LabelOutputRaw>;

/** Convert raw model output to the stored format with ordinal specificity. */
export function toLabelOutput(raw: LabelOutputRaw): LabelOutput {
  return {
    content_category: raw.content_category,
    specificity_level: SPECIFICITY_ORDINAL[raw.specificity] as SpecificityLevel,
    category_confidence: raw.category_confidence,
    specificity_confidence: raw.specificity_confidence,
    reasoning: raw.reasoning,
  };
}

/**
 * Stored label format — what goes in annotations and downstream.
 * Keeps ordinal specificity_level for all existing consumers.
 */
export const LabelOutput = z.object({
  content_category: ContentCategory.describe(
    "The single most applicable content category for this paragraph",
  ),
  specificity_level: SpecificityLevel.describe(
    "1=generic boilerplate, 2=sector-adapted, 3=firm-specific, 4=quantified-verifiable",
  ),
  category_confidence: Confidence.describe(
    "high=clear-cut, medium=some ambiguity, low=genuinely torn between categories",
  ),
  specificity_confidence: Confidence.describe(
    "high=clear-cut, medium=borderline adjacent levels, low=could argue for 2+ levels",
  ),
  reasoning: z
    .string()
    .describe(
      "Brief 1-2 sentence justification for the content_category assignment",
    ),
});
export type LabelOutput = z.infer<typeof LabelOutput>;
