import type { Organization, Extension } from 'fhir/r4';

export const ORGANIZATION_ONBOARDING_STATUS_EXTENSION_URL =
  'https://akashichealthrecords.com/fhir/StructureDefinition/organization-onboarding-status';

export type OrgOnboardingTaskState = 'done' | 'skipped';
export type OrgOnboardingPersistedStatus = Record<string, OrgOnboardingTaskState>;

const TASK_SUB_URL = 'task';
const TASK_ID_SUB_URL = 'id';
const TASK_STATUS_SUB_URL = 'status';
const DISMISSED_SUB_URL = 'dismissed';

const findParentExtension = (
  organization?: Organization | null,
): Extension | undefined =>
  organization?.extension?.find(
    (e) => e.url === ORGANIZATION_ONBOARDING_STATUS_EXTENSION_URL,
  );

export const getOnboardingStatus = (
  organization?: Organization | null,
): OrgOnboardingPersistedStatus => {
  const parent = findParentExtension(organization);
  const status: OrgOnboardingPersistedStatus = {};
  for (const sub of parent?.extension ?? []) {
    if (sub.url !== TASK_SUB_URL) continue;
    const id = sub.extension?.find((e) => e.url === TASK_ID_SUB_URL)?.valueString;
    const raw = sub.extension?.find((e) => e.url === TASK_STATUS_SUB_URL)?.valueCode;
    if (id && (raw === 'done' || raw === 'skipped')) {
      status[id] = raw;
    }
  }
  return status;
};

export const getOnboardingDismissed = (
  organization?: Organization | null,
): boolean => {
  const parent = findParentExtension(organization);
  return (
    parent?.extension?.find((e) => e.url === DISMISSED_SUB_URL)?.valueBoolean === true
  );
};

// Rebuild the parent extension from a full status map + dismissed flag, preserving
// every OTHER organization extension. Omit the parent entirely when empty.
const writeOnboardingExtension = (
  organization: Organization,
  status: OrgOnboardingPersistedStatus,
  dismissed: boolean,
): Organization => {
  const others = (organization.extension || []).filter(
    (e) => e.url !== ORGANIZATION_ONBOARDING_STATUS_EXTENSION_URL,
  );

  const subExtensions: Extension[] = [];
  if (dismissed) {
    subExtensions.push({ url: DISMISSED_SUB_URL, valueBoolean: true });
  }
  for (const [id, state] of Object.entries(status)) {
    subExtensions.push({
      url: TASK_SUB_URL,
      extension: [
        { url: TASK_ID_SUB_URL, valueString: id },
        { url: TASK_STATUS_SUB_URL, valueCode: state },
      ],
    });
  }

  const extension =
    subExtensions.length > 0
      ? [...others, { url: ORGANIZATION_ONBOARDING_STATUS_EXTENSION_URL, extension: subExtensions }]
      : others;

  return { ...organization, extension };
};

export const setOnboardingTaskStatus = (
  organization: Organization,
  taskId: string,
  status: OrgOnboardingTaskState,
): Organization => {
  const next = { ...getOnboardingStatus(organization), [taskId]: status };
  return writeOnboardingExtension(organization, next, getOnboardingDismissed(organization));
};

export const setOnboardingDismissed = (
  organization: Organization,
  dismissed: boolean,
): Organization =>
  writeOnboardingExtension(
    organization,
    getOnboardingStatus(organization),
    dismissed,
  );