import type { ProsthesisCatalogEntry } from '@/types/treatment-catalog';
import type { FdiToothId, LabCaseToothProsthesisDraft, ToothSelectionGroup } from '@/types/treatment';
import { PROSTHESIS_CATEGORY_COLORS } from '@/components/shared/catalog-type-colors';
import { prosthesisTypeColorFromCatalog } from '@/components/treatment/prosthesisTypeDisplay';
import {
  deriveTeethFromGroups,
  groupsFromFlatTeeth,
  newGroupId,
} from '@/components/treatment/toothSelectionGroups';

const PROSTHESIS_CATEGORY_ORDER = [
  'implant',
  'crown',
  'indirect',
  'post_core',
  'removable',
  'appliance',
  'digital',
] as const;

export const TOOTH_CATEGORIES = ['implant', 'crown', 'indirect', 'post_core', 'removable'] as const;
export const ARCH_CATEGORIES = ['removable', 'appliance', 'digital'] as const;

/** Not FDI — jaw-level prosthesis rows on LabCaseToothProsthesis.tooth */
export const ARCH_TOOTH_UPPER = 'UA';
export const ARCH_TOOTH_LOWER = 'LA';

export type ArchTarget = 'upper' | 'lower' | 'both';
export type PickerScope = 'tooth' | 'arch';

export function isArchSentinel(tooth: string): boolean {
  return tooth === ARCH_TOOTH_UPPER || tooth === ARCH_TOOTH_LOWER;
}

export function archSentinels(arch: ArchTarget): string[] {
  if (arch === 'upper') return [ARCH_TOOTH_UPPER];
  if (arch === 'lower') return [ARCH_TOOTH_LOWER];
  return [ARCH_TOOTH_UPPER, ARCH_TOOTH_LOWER];
}

export function catalogByCode(
  catalog: readonly ProsthesisCatalogEntry[],
): Map<string, ProsthesisCatalogEntry> {
  return new Map(catalog.map((entry) => [entry.code, entry]));
}

export const PARTIAL_DENTURE_CODE = 'partial_denture';

export function effectiveChartRegion(
  entry: Pick<ProsthesisCatalogEntry, 'code' | 'chartRegion'>,
): string {
  if (entry.code === PARTIAL_DENTURE_CODE) return 'crown';
  return entry.chartRegion;
}

export function matchesPickerScope(
  entry: Pick<ProsthesisCatalogEntry, 'chartRegion' | 'code'>,
  scope?: PickerScope,
): boolean {
  if (!scope) return true;
  const region = effectiveChartRegion(entry);
  if (scope === 'tooth') return region !== 'arch';
  return region === 'arch';
}

export function categoryCodesInCatalog(
  catalog: readonly ProsthesisCatalogEntry[],
  scope?: PickerScope,
): string[] {
  const present = new Set(
    catalog.filter((e) => matchesPickerScope(e, scope)).map((e) => e.category),
  );
  const order =
    scope === 'tooth'
      ? TOOTH_CATEGORIES
      : scope === 'arch'
        ? ARCH_CATEGORIES
        : PROSTHESIS_CATEGORY_ORDER;
  return order.filter((code) => present.has(code));
}

/** Matches category titles (e.g. Veneer/Inlay/Onlay/Overlay), not alphabetical codes. */
const SUBCATEGORY_ORDER: Record<string, readonly string[]> = {
  indirect: ['veneer', 'inlay', 'onlay', 'overlay'],
};

export function subcategoriesFor(
  catalog: readonly ProsthesisCatalogEntry[],
  category: string,
  scope?: PickerScope,
): string[] {
  const set = new Set(
    catalog
      .filter(
        (e) =>
          e.category === category &&
          e.subcategory &&
          matchesPickerScope(e, scope),
      )
      .map((e) => e.subcategory),
  );
  const preferred = SUBCATEGORY_ORDER[category];
  if (!preferred) return [...set].sort();
  const rest = [...set].filter((code) => !preferred.includes(code)).sort();
  return [...preferred.filter((code) => set.has(code)), ...rest];
}

export function leavesFor(
  catalog: readonly ProsthesisCatalogEntry[],
  category: string,
  subcategory?: string,
  scope?: PickerScope,
): ProsthesisCatalogEntry[] {
  return catalog
    .filter((e) => {
      if (e.category !== category) return false;
      if (!matchesPickerScope(e, scope)) return false;
      if (subcategory) return e.subcategory === subcategory;
      const subs = subcategoriesFor(catalog, category, scope);
      if (subs.length === 0) return true;
      return !e.subcategory;
    })
    .sort((a, b) => a.sortOrder - b.sortOrder || a.code.localeCompare(b.code));
}

export const SCREW_RETAINED_CODE = 'screw_retained';

export type AddonSlot = 'crown';

export function jobsHaveStackGroup(
  jobCodes: readonly string[],
  byCode: Map<string, ProsthesisCatalogEntry>,
  group: string,
): boolean {
  return jobCodes.some((code) => byCode.get(code)?.stackGroup === group);
}

export function canStackLeaf(
  existingCodes: readonly string[],
  nextCode: string,
  byCode: Map<string, ProsthesisCatalogEntry>,
): boolean {
  const next = byCode.get(nextCode);
  if (!next) return false;
  const hasRestoration = jobsHaveStackGroup(existingCodes, byCode, 'restoration');
  const hasImplant = jobsHaveStackGroup(existingCodes, byCode, 'implant');
  const hasPost = jobsHaveStackGroup(existingCodes, byCode, 'post_core');
  if (next.stackGroup === 'post_core') return !hasRestoration && !hasImplant;
  if (next.stackGroup === 'implant') return !hasPost;
  // Screw-retained already paints the crown — do not stack a second restoration.
  if (next.stackGroup === 'restoration' && existingCodes.includes(SCREW_RETAINED_CODE)) {
    return false;
  }
  return true;
}

export function applyLeafToJobs(
  existingCodes: readonly string[],
  nextCode: string,
  byCode: Map<string, ProsthesisCatalogEntry>,
): string[] {
  const next = byCode.get(nextCode);
  if (!next) return [...existingCodes];
  if (!canStackLeaf(existingCodes, nextCode, byCode)) return [...existingCodes];

  const jobs = existingCodes.filter((code) => {
    const cur = byCode.get(code);
    if (!cur) return false;
    if (cur.stackGroup === next.stackGroup) return false;
    if (nextCode === SCREW_RETAINED_CODE && cur.stackGroup === 'restoration') return false;
    return true;
  });

  jobs.push(nextCode);
  return jobs;
}

export function crownRestorationCode(
  jobCodes: readonly string[],
  byCode: Map<string, ProsthesisCatalogEntry>,
): string | undefined {
  return jobCodes.find((code) => {
    const entry = byCode.get(code);
    return entry?.category === 'crown' && entry.stackGroup === 'restoration';
  });
}

/**
 * Crown suggestion slot after implant or post & core.
 * Stays visible when the restoration is a crown (filled chip). Hidden for
 * screw-retained and for non-crown restorations (veneer / inlay / onlay / overlay).
 */
export function jobsAllowCrownAddon(
  jobCodes: readonly string[],
  byCode: Map<string, ProsthesisCatalogEntry>,
): boolean {
  if (jobCodes.length === 0) return false;
  if (jobCodes.includes(SCREW_RETAINED_CODE)) return false;
  if (
    !jobsHaveStackGroup(jobCodes, byCode, 'implant') &&
    !jobsHaveStackGroup(jobCodes, byCode, 'post_core')
  ) {
    return false;
  }
  const restoration = jobCodes.find((code) => byCode.get(code)?.stackGroup === 'restoration');
  if (!restoration) return true;
  return byCode.get(restoration)?.category === 'crown';
}

export function categoryDisabledForJobs(
  category: string,
  jobCodes: readonly string[],
  byCode: Map<string, ProsthesisCatalogEntry>,
): boolean {
  const hasRestoration = jobsHaveStackGroup(jobCodes, byCode, 'restoration');
  const hasImplant = jobsHaveStackGroup(jobCodes, byCode, 'implant');
  const hasPost = jobsHaveStackGroup(jobCodes, byCode, 'post_core');
  if (category === 'post_core') return hasRestoration || hasImplant;
  if (category === 'implant') return hasPost;
  if (
    jobCodes.includes(SCREW_RETAINED_CODE) &&
    (category === 'crown' || category === 'indirect')
  ) {
    return true;
  }
  return false;
}

export function fdiTeethWithJobs(
  rows: readonly Pick<LabCaseToothProsthesisDraft, 'tooth' | 'prosthesisTypeCode' | 'detailClientId'>[],
  detailClientId: string,
): Set<string> {
  const teeth = new Set<string>();
  for (const row of rows) {
    if (
      row.detailClientId === detailClientId &&
      row.prosthesisTypeCode &&
      !isArchSentinel(row.tooth)
    ) {
      teeth.add(row.tooth);
    }
  }
  return teeth;
}

export function pruneGroupsToJobTeeth(
  groups: ToothSelectionGroup[],
  jobTeeth: Set<string>,
): ToothSelectionGroup[] {
  const extra = groups.flatMap((g) => g.teeth).filter((tooth) => !jobTeeth.has(tooth));
  if (extra.length === 0) return groups;
  return removeTeethFromGroups(groups, extra);
}

export function pruneDetailTeethToJobs<
  T extends {
    clientId: string;
    teeth: readonly string[];
    toothSelectionGroups: ToothSelectionGroup[];
  },
>(detail: T, jobs: readonly LabCaseToothProsthesisDraft[]): T {
  const jobTeeth = fdiTeethWithJobs(jobs, detail.clientId);
  const groups =
    detail.toothSelectionGroups.length > 0
      ? detail.toothSelectionGroups
      : groupsFromFlatTeeth(detail.teeth as FdiToothId[]);
  const nextGroups = pruneGroupsToJobTeeth(groups, jobTeeth);
  return {
    ...detail,
    toothSelectionGroups: nextGroups,
    teeth: deriveTeethFromGroups(nextGroups),
  };
}

export function codesOnTooth(
  rows: readonly Pick<LabCaseToothProsthesisDraft, 'tooth' | 'prosthesisTypeCode' | 'detailClientId'>[],
  detailClientId: string,
  tooth: string,
): string[] {
  return rows
    .filter((row) => row.detailClientId === detailClientId && row.tooth === tooth && row.prosthesisTypeCode)
    .map((row) => row.prosthesisTypeCode);
}

export function writeJobsForTeeth(
  rows: LabCaseToothProsthesisDraft[],
  detailClientId: string,
  teeth: readonly string[],
  codes: readonly string[],
  groupIdForTooth: (tooth: string) => string,
): LabCaseToothProsthesisDraft[] {
  const toothSet = new Set(teeth);
  const rest = rows.filter((row) => !(row.detailClientId === detailClientId && toothSet.has(row.tooth)));
  const next = teeth.flatMap((tooth) =>
    codes.map((prosthesisTypeCode) => ({
      detailClientId,
      tooth,
      prosthesisTypeCode,
      selectionGroupId: groupIdForTooth(tooth),
    })),
  );
  return [...rest, ...next];
}

export function clearJobsForTeeth(
  rows: LabCaseToothProsthesisDraft[],
  detailClientId: string,
  teeth: readonly string[],
): LabCaseToothProsthesisDraft[] {
  const toothSet = new Set(teeth);
  return rows.filter((row) => !(row.detailClientId === detailClientId && toothSet.has(row.tooth)));
}

export function ensureTeethInGroups(
  groups: ToothSelectionGroup[],
  teeth: readonly FdiToothId[],
): ToothSelectionGroup[] {
  const present = new Set(groups.flatMap((g) => g.teeth));
  const missing = teeth.filter((t) => !present.has(t));
  if (missing.length === 0) return groups;
  return [
    ...groups,
    ...missing.map((tooth) => ({
      groupId: newGroupId(),
      kind: 'single' as const,
      teeth: [tooth],
    })),
  ];
}

export function removeTeethFromGroups(
  groups: ToothSelectionGroup[],
  teeth: readonly FdiToothId[],
): ToothSelectionGroup[] {
  const drop = new Set(teeth);
  const next: ToothSelectionGroup[] = [];
  for (const group of groups) {
    const kept = group.teeth.filter((t) => !drop.has(t));
    if (kept.length === 0) continue;
    if (kept.length === 1) {
      next.push({ groupId: group.groupId, kind: 'single', teeth: kept });
    } else {
      next.push({ ...group, teeth: kept, kind: group.kind === 'connected' ? 'connected' : 'single' });
    }
  }
  return next.length > 0 ? next : groupsFromFlatTeeth([]);
}

export function toothRegionColors(
  rows: readonly Pick<LabCaseToothProsthesisDraft, 'tooth' | 'prosthesisTypeCode'>[],
  catalog: readonly ProsthesisCatalogEntry[],
): {
  crown: Partial<Record<FdiToothId, string>>;
  root: Partial<Record<FdiToothId, string>>;
} {
  const byCode = catalogByCode(catalog);
  const crown: Partial<Record<FdiToothId, string>> = {};
  const root: Partial<Record<FdiToothId, string>> = {};
  for (const row of rows) {
    if (isArchSentinel(row.tooth)) continue;
    const entry = byCode.get(row.prosthesisTypeCode);
    const color = prosthesisTypeColorFromCatalog(row.prosthesisTypeCode, catalog);
    const fdi = row.tooth as FdiToothId;
    const region = entry ? effectiveChartRegion(entry) : 'crown';
    if (region === 'root') root[fdi] = color;
    else if (region === 'arch') {
      crown[fdi] = crown[fdi] ?? color;
      root[fdi] = root[fdi] ?? color;
    } else {
      crown[fdi] = color;
    }
  }
  return { crown, root };
}

export function archTargetFromJobs(
  rows: readonly Pick<LabCaseToothProsthesisDraft, 'tooth' | 'prosthesisTypeCode' | 'detailClientId'>[],
  detailClientId: string,
): ArchTarget | null {
  const hasUpper = rows.some(
    (r) =>
      r.detailClientId === detailClientId &&
      r.tooth === ARCH_TOOTH_UPPER &&
      r.prosthesisTypeCode,
  );
  const hasLower = rows.some(
    (r) =>
      r.detailClientId === detailClientId &&
      r.tooth === ARCH_TOOTH_LOWER &&
      r.prosthesisTypeCode,
  );
  if (hasUpper && hasLower) return 'both';
  if (hasUpper) return 'upper';
  if (hasLower) return 'lower';
  return null;
}

export function hasArchJobs(
  rows: readonly Pick<LabCaseToothProsthesisDraft, 'tooth' | 'prosthesisTypeCode' | 'detailClientId'>[],
  detailClientId: string,
): boolean {
  return archTargetFromJobs(rows, detailClientId) !== null;
}

export function writeArchJobs(
  rows: LabCaseToothProsthesisDraft[],
  detailClientId: string,
  arch: ArchTarget,
  code: string,
): LabCaseToothProsthesisDraft[] {
  const rest = rows.filter(
    (row) => !(row.detailClientId === detailClientId && isArchSentinel(row.tooth)),
  );
  const next = archSentinels(arch).map((tooth) => ({
    detailClientId,
    tooth,
    prosthesisTypeCode: code,
    selectionGroupId: '',
  }));
  return [...rest, ...next];
}

/** Move / copy / drop UA·LA rows so they match the segmented arch control. */
export function retargetArchJobs(
  rows: LabCaseToothProsthesisDraft[],
  detailClientId: string,
  from: ArchTarget,
  to: ArchTarget,
): LabCaseToothProsthesisDraft[] {
  if (from === to) return rows;
  const keep = archSentinels(to).filter((tooth) => archSentinels(from).includes(tooth));
  const code =
    uniqueCodesOnGroup(rows, detailClientId, keep)[0] ??
    uniqueCodesOnGroup(rows, detailClientId, archSentinels(from))[0] ??
    uniqueCodesOnGroup(rows, detailClientId, [ARCH_TOOTH_UPPER, ARCH_TOOTH_LOWER])[0];
  if (!code) return rows;
  return writeArchJobs(rows, detailClientId, to, code);
}

export function uniqueCodesOnGroup(
  rows: readonly LabCaseToothProsthesisDraft[],
  detailClientId: string,
  teeth: readonly string[],
): string[] {
  const codes = new Set<string>();
  const toothSet = new Set(teeth);
  for (const row of rows) {
    if (row.detailClientId === detailClientId && toothSet.has(row.tooth) && row.prosthesisTypeCode) {
      codes.add(row.prosthesisTypeCode);
    }
  }
  return [...codes];
}

function sortProsthesisTeeth(teeth: readonly string[]): string[] {
  return [...teeth].sort((a, b) => {
    const rank = (tooth: string) =>
      tooth === ARCH_TOOTH_UPPER ? 0 : tooth === ARCH_TOOTH_LOWER ? 1 : 2;
    const diff = rank(a) - rank(b);
    if (diff !== 0) return diff;
    return a.localeCompare(b, undefined, { numeric: true });
  });
}

export type ProsthesisJobDisplayRow = {
  teeth: string[];
  codes: string[];
  connected: boolean;
};

/**
 * Dispatch / summary rows.
 * Connected teeth that share a type collapse to one row (bridge). An extra type on
 * one unit of that bridge is its own row. Unconnected teeth stay one row per tooth
 * (stacked codes stay together).
 */
export function toothJobRowsForDetail(
  rows: readonly Pick<
    LabCaseToothProsthesisDraft,
    'tooth' | 'prosthesisTypeCode' | 'detailClientId' | 'selectionGroupId'
  >[],
  detailClientId: string,
): ProsthesisJobDisplayRow[] {
  const byTooth = new Map<string, { codes: string[]; selectionGroupId: string }>();
  for (const row of rows) {
    if (row.detailClientId !== detailClientId || !row.prosthesisTypeCode) continue;
    const entry = byTooth.get(row.tooth) ?? { codes: [], selectionGroupId: '' };
    if (!entry.codes.includes(row.prosthesisTypeCode)) {
      entry.codes.push(row.prosthesisTypeCode);
    }
    const groupId = row.selectionGroupId?.trim() || '';
    if (!entry.selectionGroupId && groupId) {
      entry.selectionGroupId = groupId;
    }
    byTooth.set(row.tooth, entry);
  }

  const teethByGroup = new Map<string, string[]>();
  for (const [tooth, entry] of byTooth) {
    if (!entry.selectionGroupId) continue;
    const list = teethByGroup.get(entry.selectionGroupId) ?? [];
    list.push(tooth);
    teethByGroup.set(entry.selectionGroupId, list);
  }
  const bridgeGroupIds = new Set(
    [...teethByGroup.entries()]
      .filter(([, teeth]) => teeth.length > 1)
      .map(([groupId]) => groupId),
  );

  const connectedByType = new Map<string, { code: string; teeth: string[] }>();
  const singles: ProsthesisJobDisplayRow[] = [];

  for (const [tooth, entry] of byTooth) {
    if (entry.selectionGroupId && bridgeGroupIds.has(entry.selectionGroupId)) {
      for (const code of entry.codes) {
        const key = `${entry.selectionGroupId}::${code}`;
        const group = connectedByType.get(key) ?? { code, teeth: [] };
        if (!group.teeth.includes(tooth)) group.teeth.push(tooth);
        connectedByType.set(key, group);
      }
      continue;
    }
    singles.push({ teeth: [tooth], codes: entry.codes, connected: false });
  }

  const connectedRows: ProsthesisJobDisplayRow[] = [...connectedByType.values()].map((group) => {
    const teeth = sortProsthesisTeeth(group.teeth);
    return { teeth, codes: [group.code], connected: teeth.length > 1 };
  });

  return [...connectedRows, ...singles].sort((a, b) => {
    const toothA = a.teeth[0] ?? '';
    const toothB = b.teeth[0] ?? '';
    if (toothA !== toothB) {
      return sortProsthesisTeeth([toothA, toothB])[0] === toothA ? -1 : 1;
    }
    return (a.codes[0] ?? '').localeCompare(b.codes[0] ?? '');
  });
}

export function categoryTileColor(category: string): string {
  return PROSTHESIS_CATEGORY_COLORS[category] ?? '#e2e8f0';
}
