import { z } from 'zod';

const evidenceArray = z.array(z.string()).default([]);
const confidenceValue = z.number().min(0).max(1).default(0);
const completionScore = z.number().min(0).max(1).default(0);

const defaultWeightedString = {
  description: null as string | null,
  confidence: 0,
  evidence: [] as string[],
};

const weightedStringField = z
  .object({
    description: z.union([z.string(), z.null()]).default(null),
    confidence: confidenceValue.default(0),
    evidence: evidenceArray.default([]),
  })
  .default(defaultWeightedString);

const weightedListItem = z.object({
  id: z.string(),
  description: z.string(),
  confidence: confidenceValue,
  evidence: evidenceArray,
});

const stageEnum = z.enum([
  'idea',
  'prototype',
  'mvp_in_progress',
  'live_beta',
  'live_paid',
  'unknown',
]);

const severityEnum = z.enum(['low', 'medium', 'high', 'unknown']);
const frequencyEnum = z.enum(['rare', 'occasional', 'frequent', 'constant', 'unknown']);
const competitorTypeEnum = z.enum(['direct', 'indirect', 'alternative', 'unknown']);
const relatedAreaEnum = z.enum(['product', 'tech', 'market', 'business_model', 'other']);
const priorityEnum = z.enum(['high', 'medium', 'low']);

export const ChatExtractionSchema = z.object({
  project_summary: z.object({
    working_title: z.union([z.string(), z.null()]).default(null),
    one_liner: z.union([z.string(), z.null()]).default(null),
    stage: stageEnum.default('unknown'),
    overall_confidence: confidenceValue,
    evidence: evidenceArray,
  }),
  product_vision: z.object({
    problem_statement: weightedStringField,
    target_outcome: weightedStringField,
    founder_intent: weightedStringField,
    completion_score: completionScore,
  }),
  target_users: z.object({
    primary_segment: weightedStringField,
    segments: z
      .array(
        z.object({
          id: z.string(),
          description: z.string(),
          jobs_to_be_done: z.array(z.string()).default([]),
          environment: z.union([z.string(), z.null()]),
          confidence: confidenceValue,
          evidence: evidenceArray,
        }),
      )
      .default([]),
    completion_score: completionScore,
  }),
  problems_and_pains: z.object({
    problems: z
      .array(
        z.object({
          id: z.string(),
          description: z.string(),
          severity: severityEnum,
          frequency: frequencyEnum,
          confidence: confidenceValue,
          evidence: evidenceArray,
        }),
      )
      .default([]),
    completion_score: completionScore,
  }),
  solution_and_features: z.object({
    core_solution: weightedStringField,
    core_features: z
      .array(
        z.object({
          id: z.string(),
          name: z.string(),
          description: z.string(),
          is_must_have_for_v1: z.boolean(),
          confidence: confidenceValue,
          evidence: evidenceArray,
        }),
      )
      .default([]),
    nice_to_have_features: z
      .array(
        z.object({
          id: z.string(),
          name: z.string(),
          description: z.string(),
          confidence: confidenceValue,
          evidence: evidenceArray,
        }),
      )
      .default([]),
    completion_score: completionScore,
  }),
  market_and_competition: z.object({
    market_category: weightedStringField,
    competitors: z
      .array(
        z.object({
          id: z.string(),
          name: z.string(),
          description: z.string(),
          type: competitorTypeEnum,
          confidence: confidenceValue,
          evidence: evidenceArray,
        }),
      )
      .default([]),
    differentiation_points: weightedListItem.array().default([]),
    completion_score: completionScore,
  }),
  tech_and_constraints: z.object({
    stack_mentions: weightedListItem.array().default([]),
    constraints: weightedListItem.array().default([]),
    completion_score: completionScore,
  }),
  execution_status: z.object({
    current_stage: weightedStringField,
    work_done: weightedListItem.array().default([]),
    work_in_progress: weightedListItem.array().default([]),
    blocked_items: weightedListItem.array().default([]),
    completion_score: completionScore,
  }),
  goals_and_success: z.object({
    short_term_goals: weightedListItem.array().default([]),
    long_term_goals: weightedListItem.array().default([]),
    success_criteria: weightedListItem.array().default([]),
    completion_score: completionScore,
  }),
  unknowns_and_questions: z.object({
    unknowns: z
      .array(
        z.object({
          id: z.string(),
          description: z.string(),
          related_area: relatedAreaEnum,
          evidence: evidenceArray,
          confidence: confidenceValue,
        }),
      )
      .default([]),
    questions_to_ask_user: z
      .array(
        z.object({
          id: z.string(),
          question: z.string(),
          priority: priorityEnum,
        }),
      )
      .default([]),
  }),
  summary_scores: z.object({
    overall_completion: completionScore,
    overall_confidence: confidenceValue,
  }),
});

export type ChatExtractionData = z.infer<typeof ChatExtractionSchema>;