import { sql } from 'drizzle-orm';
import type { Database } from '@cancerindex/database';
import { countryToIso3 } from './country-codes.js';

// The package index (owned by the integrator) already re-exports this module; surface the pure
// country mapping through it so apps import `countryToIso3` / `iso3ToName` from '@cancerindex/ranking'.
export * from './country-codes.js';

export const TRIAL_SITE_COUNTRY_FORMULA_VERSION = 'ci-trial-sites-v1';

/** Phase buckets exposed by the map; EARLY_PHASE1 is folded into PHASE1, NA and empty phases only count under "any phase". */
export const TRIAL_SITE_PHASES = ['PHASE1', 'PHASE2', 'PHASE3', 'PHASE4'] as const;
export type TrialSitePhase = (typeof TRIAL_SITE_PHASES)[number];

export interface TrialSiteCountsResult {
  /** Rows written to trial_site_country_counts. */
  rows: number;
  /** Distinct non-empty country names seen in trial_locations. */
  countries: number;
  /** Country names that could not be mapped to ISO 3166-1 alpha-3 (kept with iso3 = null). */
  unmapped: string[];
  /** Top-level cancers aggregated (in addition to the all-oncology scope). */
  cancers: number;
  ms: number;
}

/**
 * Recompute `trial_site_country_counts` (SPEC §11 — trial map), formula `ci-trial-sites-v1`.
 *
 * One row per (cancer_id, phase, recruiting_only, country) for the cartesian product
 *   cancer_id      ∈ {NULL = every oncology trial} ∪ {each active `top_level` cancer, trials attributed
 *                    through `trial_conditions` to the cancer or any NCIt descendant (depth ≤ 12,
 *                    same traversal as counters.ts)}
 *   phase          ∈ {NULL = any phase, PHASE1, PHASE2, PHASE3, PHASE4} — a PHASE2|PHASE3 study is
 *                    counted under both; EARLY_PHASE1 counts under PHASE1; NA / no phase → any only
 *   recruiting_only∈ {false = every location, true = location status RECRUITING, or, when the
 *                    location has no status, study overall_status RECRUITING}
 *
 *   sites  = number of `trial_locations` rows (a study with 40 US sites weighs 40)
 *   trials = distinct studies with ≥ 1 site in the country
 *
 * Interventional and observational studies are both included (no study_type filter). Locations
 * with an empty country are excluded. Deterministic, set-based, one transaction (delete + insert).
 */
export async function computeTrialSiteCounts(db: Database): Promise<TrialSiteCountsResult> {
  const t0 = Date.now();
  return db.transaction(async (tx) => {
    // 1. Locations joined to their study, with the effective "recruiting" flag and phase array.
    await tx.execute(sql`
      CREATE TEMP TABLE _tl ON COMMIT DROP AS
      SELECT l.trial_id, l.country, (coalesce(l.status, t.overall_status) = 'RECRUITING') AS recruiting, t.phases
      FROM trial_locations l JOIN clinical_trials t ON t.id = l.trial_id
      WHERE l.country IS NOT NULL AND l.country <> ''`);

    // 2. Country → ISO3 through the pure mapping (JS side; ~200 names).
    const names = await tx.execute<{ country: string }>(sql`SELECT DISTINCT country FROM _tl ORDER BY country`);
    const mapping = Array.from(names).map((r) => ({ country: r.country, iso3: countryToIso3(r.country) }));
    const unmapped = mapping.filter((m) => m.iso3 === null).map((m) => m.country);
    await tx.execute(sql`CREATE TEMP TABLE _iso (country text PRIMARY KEY, iso3 text) ON COMMIT DROP`);
    const mapped = mapping.filter((m) => m.iso3 !== null);
    for (let i = 0; i < mapped.length; i += 100) {
      const chunk = mapped.slice(i, i + 100);
      await tx.execute(sql`INSERT INTO _iso (country, iso3) VALUES ${sql.join(chunk.map((m) => sql`(${m.country}, ${m.iso3})`), sql`, `)}`);
    }

    // 3. Location × phase bucket (NULL = any phase, plus each normalized phase present on the study).
    const phaseList = sql.raw(`(${['EARLY_PHASE1', ...TRIAL_SITE_PHASES].map((p) => `'${p}'`).join(',')})`);
    await tx.execute(sql`
      CREATE TEMP TABLE _tlp ON COMMIT DROP AS
      SELECT trial_id, country, recruiting, NULL::text AS phase FROM _tl
      UNION ALL
      SELECT l.trial_id, l.country, l.recruiting, p.phase
      FROM _tl l CROSS JOIN LATERAL (
        SELECT DISTINCT CASE WHEN x = 'EARLY_PHASE1' THEN 'PHASE1' ELSE x END AS phase
        FROM unnest(l.phases) x WHERE x IN ${phaseList}
      ) p`);
    await tx.execute(sql`CREATE INDEX ON _tlp (trial_id)`);

    // 4. Trial × cancer scope: NULL (all oncology trials) + every active top-level cancer whose
    //    descendant set (NCIt hierarchy, depth ≤ 12) contains a mapped condition of the trial.
    await tx.execute(sql`
      CREATE TEMP TABLE _desc ON COMMIT DROP AS
      WITH RECURSIVE d AS (
        SELECT id AS ancestor, id AS descendant, 0 AS depth FROM cancers WHERE status = 'active' AND top_level
        UNION
        SELECT d.ancestor, h.child_id, d.depth + 1 FROM d JOIN cancer_hierarchy h ON h.parent_id = d.descendant
        WHERE d.depth < 12
      )
      SELECT DISTINCT ancestor, descendant FROM d`);
    await tx.execute(sql`CREATE INDEX ON _desc (descendant)`);
    await tx.execute(sql`
      CREATE TEMP TABLE _tc ON COMMIT DROP AS
      SELECT NULL::varchar(32) AS cancer_id, t.id AS trial_id FROM clinical_trials t
      UNION ALL
      SELECT DISTINCT d.ancestor, tc.trial_id FROM trial_conditions tc JOIN _desc d ON d.descendant = tc.cancer_id WHERE tc.cancer_id IS NOT NULL`);
    await tx.execute(sql`CREATE INDEX ON _tc (trial_id)`);

    // 5. Rebuild.
    await tx.execute(sql`DELETE FROM trial_site_country_counts`);
    const inserted = await tx.execute<{ n: string }>(sql`
      WITH ins AS (
        -- Schema field computedAt is declared with the shared updatedAt() helper, hence column updated_at.
        INSERT INTO trial_site_country_counts (cancer_id, phase, recruiting_only, country, iso3, sites, trials, formula_version, updated_at)
        SELECT c.cancer_id, p.phase, r.recruiting_only, p.country, i.iso3, count(*)::int, count(DISTINCT p.trial_id)::int, ${TRIAL_SITE_COUNTRY_FORMULA_VERSION}, now()
        FROM _tlp p
        JOIN _tc c ON c.trial_id = p.trial_id
        CROSS JOIN (VALUES (false), (true)) AS r(recruiting_only)
        LEFT JOIN _iso i ON i.country = p.country
        WHERE NOT r.recruiting_only OR p.recruiting
        GROUP BY c.cancer_id, p.phase, r.recruiting_only, p.country, i.iso3
        RETURNING 1
      )
      SELECT count(*)::text AS n FROM ins`);
    const cancers = await tx.execute<{ n: string }>(sql`SELECT count(*)::text AS n FROM cancers WHERE status = 'active' AND top_level`);
    return { rows: Number(inserted[0]?.n ?? 0), countries: mapping.length, unmapped, cancers: Number(cancers[0]?.n ?? 0), ms: Date.now() - t0 };
  });
}