import { ACTIVITY_TABS } from "@/constants/constants";
import { z } from "zod";

export const entityTypeEnum = z.enum(["contact", "lead"]);

const { CALL, MEETING, SITE_VISIT, OTHER } = ACTIVITY_TABS;
export const activityTypeEnum = z.enum([CALL, MEETING, SITE_VISIT, OTHER]);

export const baseTaskSchema = z.object({
  // Note: entityType is not in baseTaskSchema anymore,
  // as it will be literalized in the final schemas.
  activityType: activityTypeEnum,
  date: z.date({
    required_error: "Date is required",
    invalid_type_error: "Invalid date",
  }),
  time: z.string().min(1, "Time is required"),
  assignedTo: z.string().min(1, "Assignee is required"),
});

// --- Activity-Specific Base Schemas ---

const callTaskSchema = baseTaskSchema.extend({
  activityType: z.literal(CALL),
});

const meetingTaskSchema = baseTaskSchema.extend({
  activityType: z.literal(MEETING),
  location: z.string().min(1, "Location is required"),
});

const siteVisitTaskSchema = baseTaskSchema.extend({
  activityType: z.literal(SITE_VISIT),
  project: z.string().min(1, "Project is required"),
});

const otherTaskSchema = baseTaskSchema.extend({
  activityType: z.literal(OTHER),
  activityName: z.string().min(1, "Activity Name is required"),
  project: z.string().optional(), // Related Project (Optional)
});

// --- Contact-Specific Schemas (Applying the Entity ID and Literal) ---

const contactTaskBase = z.object({
  entityType: z.literal("contact"),
  contactId: z.string().min(1, "Contact is required"),
  leadId: z.undefined().optional(),
});

const contactCallTaskSchema = callTaskSchema.merge(contactTaskBase);
const contactMeetingTaskSchema = meetingTaskSchema.merge(contactTaskBase);
const contactSiteVisitTaskSchema = siteVisitTaskSchema.merge(contactTaskBase);
const contactOtherTaskSchema = otherTaskSchema.merge(contactTaskBase);

// Define the 4 core activity-specific schemas (unchanged from prior solution)
// ... baseTaskSchema, callTaskSchema, meetingTaskSchema, etc.

const createEntityActivitySchemas = <
  E extends "contact" | "lead",
  ID extends string
>(
  entityType: E,
  idKey: ID
) => {
  // Define the base object containing the common entity fields
  const entityBase = z.object({
    entityType: z.literal(entityType),
    [idKey]: z.string().min(1, `${entityType} is required`),
    // Add the other ID as optional/undefined to prevent overlap
    ...(entityType === "contact"
      ? { leadId: z.undefined().optional() }
      : { contactId: z.undefined().optional() }),
  }) as z.ZodObject<
    {
      entityType: z.ZodLiteral<E>;
      [key in ID]: z.ZodString;
    } & {
      contactId?: z.ZodOptional<z.ZodUndefined>;
      leadId?: z.ZodOptional<z.ZodUndefined>;
    }
  >;

  // Use .merge() once for each activity type
  return {
    call: callTaskSchema.merge(entityBase),
    meeting: meetingTaskSchema.merge(entityBase),
    siteVisit: siteVisitTaskSchema.merge(entityBase),
    other: otherTaskSchema.merge(entityBase),
  };
};

// 1. Generate Contact Schemas
const contactSchemas = createEntityActivitySchemas("contact", "contactId");

// 2. Generate Lead Schemas
const leadSchemas = createEntityActivitySchemas("lead", "leadId");

// 3. Create the final, flat union of all 8 objects
export const taskFormSchema = z.union([
  // Contact activities
  contactSchemas.call,
  contactSchemas.meeting,
  contactSchemas.siteVisit,
  contactSchemas.other,

  // Lead activities
  leadSchemas.call,
  leadSchemas.meeting,
  leadSchemas.siteVisit,
  leadSchemas.other,
]);