FrenoCorp/src/lib/analytics/cohort-analysis.ts

import { eq, and, gte, lte } from "drizzle-orm";
import { cohorts, cohortMembers } from "../../db/schema";
import type { DrizzleDB } from "../../db/config/migrations";
import type { NewCohort, Cohort, NewCohortMember } from "../../db/schema";

export interface CohortDefinition {
  name: string;
  description: string;
  periodStart: Date;
  periodEnd?: Date;
  filterCriteria: Record<string, unknown>;
}

export interface CohortAnalysisResult {
  cohort: Cohort;
  retention: Record<number, number>;
  metrics: CohortMetrics;
}

export interface CohortMetrics {
  totalUsers: number;
  activeUsers: number;
  retentionRate: number;
  avgEngagement: number;
  conversionRate: number;
}

export async function createCohort(
  db: DrizzleDB,
  definition: CohortDefinition
): Promise<Cohort> {
  const cohort: NewCohort = {
    name: definition.name,
    definition: JSON.stringify(definition.filterCriteria),
    periodStart: definition.periodStart,
    periodEnd: definition.periodEnd ?? null,
    size: 0,
    retentionData: null,
    metadata: definition.description ? JSON.stringify({ description: definition.description }) : null,
  };

  const result = await db.insert(cohorts).values(cohort).returning();
  return result[0];
}

export async function addCohortMember(
  db: DrizzleDB,
  cohortId: number,
  userId: number
): Promise<void> {
  const member: NewCohortMember = {
    cohortId,
    userId,
    joinedAt: new Date(),
  };
  await db.insert(cohortMembers).values(member);

  await db
    .update(cohorts)
    .set({
      size: await getCohortSize(db, cohortId),
    })
    .where(eq(cohorts.id, cohortId));
}

export async function getCohortSize(db: DrizzleDB, cohortId: number): Promise<number> {
  const rows = await db
    .select({ count: cohortMembers.id })
    .from(cohortMembers)
    .where(eq(cohortMembers.cohortId, cohortId));
  return rows.length;
}

export async function getCohortAnalysis(
  db: DrizzleDB,
  cohortId: number
): Promise<CohortAnalysisResult | null> {
  const cohortRows = await db.select().from(cohorts).where(eq(cohorts.id, cohortId)).limit(1);
  if (cohortRows.length === 0) return null;

  const cohort = cohortRows[0];
  const members = await db
    .select()
    .from(cohortMembers)
    .where(eq(cohortMembers.cohortId, cohortId));

  const totalUsers = members.length;
  const activeUsers = members.filter((m) => {
    const daysSinceJoin = (Date.now() - m.joinedAt.getTime()) / (1000 * 60 * 60 * 24);
    return daysSinceJoin <= 30;
  }).length;

  const retentionRate = totalUsers > 0 ? activeUsers / totalUsers : 0;

  const retention = computeRetentionCurve(members);

  const metrics: CohortMetrics = {
    totalUsers,
    activeUsers,
    retentionRate,
    avgEngagement: 0,
    conversionRate: 0,
  };

  return {
    cohort,
    retention,
    metrics,
  };
}

function computeRetentionCurve(members: typeof cohortMembers.$inferSelect[]): Record<number, number> {
  const retention: Record<number, number> = {};
  const now = new Date();

  for (let week = 0; week <= 12; week++) {
    const weekStart = new Date(now.getTime() - week * 7 * 24 * 60 * 60 * 1000);
    const weekEnd = new Date(now.getTime() - (week - 1) * 7 * 24 * 60 * 60 * 1000);

    const activeInWeek = members.filter((m) => {
      return m.joinedAt >= weekStart && m.joinedAt < weekEnd;
    }).length;

    retention[week] = activeInWeek;
  }

  return retention;
}

export async function listCohorts(
  db: DrizzleDB,
  periodStart?: Date,
  periodEnd?: Date
): Promise<Cohort[]> {
  const conditions: import("drizzle-orm").SQL[] = [];

  if (periodStart) {
    conditions.push(gte(cohorts.periodStart, periodStart));
  }
  if (periodEnd) {
    conditions.push(lte(cohorts.periodEnd ?? new Date(), periodEnd));
  }

  if (conditions.length === 0) {
    return await db.select().from(cohorts);
  }

  return await db.select().from(cohorts).where(and(...conditions));
}

export function createMonthlyCohortTemplate(): CohortDefinition {
  const now = new Date();
  const monthStart = new Date(now.getFullYear(), now.getMonth(), 1);

  return {
    name: `Monthly Cohort - ${now.toLocaleDateString("en-US", { month: "long", year: "numeric" })}`,
    description: `Users who joined in ${now.toLocaleDateString("en-US", { month: "long", year: "numeric" })}`,
    periodStart: monthStart,
    periodEnd: new Date(now.getFullYear(), now.getMonth() + 1, 0),
    filterCriteria: {
      type: "signup_date",
      granularity: "month",
    },
  };
}

export function createWeeklyCohortTemplate(): CohortDefinition {
  const now = new Date();
  const weekStart = new Date(now);
  weekStart.setDate(weekStart.getDate() - weekStart.getDay());

  return {
    name: `Weekly Cohort - Week of ${weekStart.toLocaleDateString()}`,
    description: `Users who joined in the week of ${weekStart.toLocaleDateString()}`,
    periodStart: weekStart,
    periodEnd: new Date(weekStart.getTime() + 6 * 24 * 60 * 60 * 1000),
    filterCriteria: {
      type: "signup_date",
      granularity: "week",
    },
  };
}

export function createFeatureCohortTemplate(featureName: string): CohortDefinition {
  return {
    name: `Feature Cohort - ${featureName}`,
    description: `Users who have used the ${featureName} feature`,
    periodStart: new Date(),
    filterCriteria: {
      type: "feature_usage",
      feature: featureName,
    },
  };
}