src/app/page.tsx

'use client';

import { useState, useEffect } from 'react';
import * as duckdb from '@duckdb/duckdb-wasm';
import Table from './components/Table';

type ModelData = {
  ancestor: string;
  direct_children: string[] | null;
  all_children: string[];
  all_children_count: number;
  direct_children_count: number | null;
};

type OrgData = {
  org: string;
  family_model_count: number;
  family_direct_children_count: number;
  family_all_children_count: number;
};

export default function Home() {
  const [allModels, setAllModels] = useState<ModelData[]>([]);
  const [orgData, setOrgData] = useState<OrgData[]>([]);
  const [currentPage, setCurrentPage] = useState(1);
  const [pageSize, setPageSize] = useState(100);
  const [filterText, setFilterText] = useState('');
  const [isLoading, setIsLoading] = useState(true);
  const [orderBy, setOrderBy] = useState<'all_children' | 'direct_children'>('all_children');
  const [activeTab, setActiveTab] = useState<'models' | 'orgs'>('models');
  const [orgCurrentPage, setOrgCurrentPage] = useState(1);
  const [orgPageSize, setOrgPageSize] = useState(100);
  const [orgOrderBy, setOrgOrderBy] = useState<keyof OrgData>('family_all_children_count');

  useEffect(() => {
    const urlParams = new URLSearchParams(window.location.search);
    const tab = urlParams.get('tab');
    const page = urlParams.get('page');
    const order = urlParams.get('order');
    const filter = urlParams.get('filter');

    if (tab === 'orgs') {
      setActiveTab('orgs');
    }
    if (page) {
      setCurrentPage(parseInt(page, 10));
    }
    if (order === 'direct_children') {
      setOrderBy('direct_children');
    }
    if (filter) {
      setFilterText(filter);
    }
  }, []);

  useEffect(() => {
    const urlParams = new URLSearchParams();
    if (activeTab === 'orgs') {
      urlParams.set('tab', 'orgs');
    }
    if (currentPage > 1) {
      urlParams.set('page', currentPage.toString());
    }
    if (orderBy === 'direct_children') {
      urlParams.set('order', 'direct_children');
    }
    if (filterText) {
      urlParams.set('filter', filterText);
    }
    const newUrl = `${window.location.pathname}?${urlParams.toString()}`;
    window.history.replaceState(null, '', newUrl);
  }, [activeTab, currentPage, orderBy, filterText]);

  useEffect(() => {
    async function fetchData() {
      const JSDELIVR_BUNDLES = duckdb.getJsDelivrBundles();

      // Select a bundle based on browser checks
      const bundle = await duckdb.selectBundle(JSDELIVR_BUNDLES);

      const worker_url = URL.createObjectURL(
        new Blob([`importScripts("${bundle.mainWorker!}");`], { type: 'text/javascript' })
      );

      // Instantiate the asynchronous version of DuckDB-Wasm
      const worker = new Worker(worker_url);
      const logger = new duckdb.ConsoleLogger();
      const db = new duckdb.AsyncDuckDB(logger, worker);
      await db.instantiate(bundle.mainModule, bundle.pthreadWorker);

      // Register the Parquet file using the URL
      await db.registerFileURL(
        'ancestor_children.parquet',
        `${window.location.origin}/ancestor_children.parquet`,
        duckdb.DuckDBDataProtocol.HTTP,
        false
      );

      // Execute the SQL query using the registered Parquet file
      const query = `
        SELECT
          ancestor,
          direct_children,
          all_children,
          CAST(all_children_count AS INTEGER) AS all_children_count,
          CAST(direct_children_count AS INTEGER) AS direct_children_count
        FROM 'ancestor_children.parquet'
      `;
      const conn = await db.connect();
      const result = await conn.query(query);

      // Convert the result to a JavaScript array
      const data: ModelData[] = result.toArray();

      // Execute the SQL query to get the grouped org data with additional counts
      const orgQuery = `
        SELECT
          SPLIT_PART(ancestor, '/', 1) AS org,
          CAST(COUNT(DISTINCT ancestor) AS INTEGER) AS family_model_count,
          CAST(SUM(direct_children_count) AS INTEGER) AS family_direct_children_count,
          CAST(SUM(all_children_count) AS INTEGER) AS family_all_children_count
        FROM 'ancestor_children.parquet'
        GROUP BY org
        ORDER BY family_all_children_count DESC
      `;
      const orgResult = await conn.query(orgQuery);

      // Convert the org result to a JavaScript array
      const orgData: OrgData[] = orgResult.toArray();

      // Close the connection and terminate the worker
      await conn.close();
      await db.terminate();

      setAllModels(data);
      setOrgData(orgData);
      setIsLoading(false);
    }
    fetchData();
  }, []);

  const filteredModels = allModels.filter((model) =>
    model.ancestor.toLowerCase().includes(filterText.toLowerCase())
  );

  const sortedModels = filteredModels.sort((a, b) => {
    if (orderBy === 'all_children') {
      return b.all_children_count - a.all_children_count;
    } else {
      return (b.direct_children_count ?? 0) - (a.direct_children_count ?? 0);
    }
  });

  const handleTabChange = (tab: 'models' | 'orgs') => {
    setActiveTab(tab);
    setCurrentPage(1);
    setOrderBy('all_children');
    setFilterText('');
  };

  const handlePageChange = (page: number, tab: 'models' | 'orgs') => {
    if (tab === 'models') {
      setCurrentPage(page);
    } else {
      setOrgCurrentPage(page);
    }
  };

  const handleOrderByClick = (column: 'all_children' | 'direct_children') => {
    setOrderBy(column);
    setCurrentPage(1);
  };

  const sortedOrgData = orgData.sort((a, b) => {
    if (orgOrderBy === 'org') {
      return a.org.localeCompare(b.org);
    }
    return b[orgOrderBy] - a[orgOrderBy];
  });

  const paginatedOrgData = sortedOrgData.slice(
    (orgCurrentPage - 1) * orgPageSize,
    orgCurrentPage * orgPageSize
  );

  const orgTotalPages = Math.ceil(sortedOrgData.length / orgPageSize);

  return (
    <main className="container mx-auto py-8 bg-white dark:bg-gray-900 text-gray-900 dark:text-white">
      <h1 className="text-4xl font-bold mb-4">All Models</h1>
      <div className="mb-4 flex space-x-4">
        <a
          href={`?tab=models`}
          onClick={(e) => {
            e.preventDefault();
            handleTabChange('models');
          }}
          className={`px-4 py-2 rounded-md ${
            activeTab === 'models'
              ? 'bg-blue-500 dark:bg-blue-600 text-white'
              : 'bg-gray-200 dark:bg-gray-700 text-gray-700 dark:text-gray-200'
          }`}
        >
          Models
        </a>
        <a
          href={`?tab=orgs`}
          onClick={(e) => {
            e.preventDefault();
            handleTabChange('orgs');
          }}
          className={`px-4 py-2 rounded-md ${
            activeTab === 'orgs'
              ? 'bg-blue-500 dark:bg-blue-600 text-white'
              : 'bg-gray-200 dark:bg-gray-700 text-gray-700 dark:text-gray-200'
          }`}
        >
          Organizations
        </a>
      </div>
      {activeTab === 'models' ? (
        <>
          <div className="mb-4">
            <input
              type="text"
              placeholder="Filter by model name"
              value={filterText}
              onChange={(e) => setFilterText(e.target.value)}
              className="px-4 py-2 border border-gray-300 dark:border-gray-700 rounded-md bg-white dark:bg-gray-800 text-gray-900 dark:text-white"
            />
          </div>
          {isLoading ? (
            <p>Loading data...</p>
          ) : (
            <Table
              data={sortedModels}
              columns={[
                { key: 'ancestor', label: 'Model' },
                {
                  key: 'direct_children_count',
                  label: 'Direct Children',
                  render: (value) => <span className="text-right">{value ?? 0}</span>,
                },
                {
                  key: 'all_children_count',
                  label: 'All Children',
                  render: (value) => <span className="text-right">{value}</span>,
                },
              ]}
              orderBy={orderBy}
              onOrderByChange={(key) => {
                setOrderBy(key as 'all_children' | 'direct_children');
                setCurrentPage(1);
              }}
              pageSize={pageSize}
            />
          )}
        </>
      ) : (
        <>
          {isLoading ? (
            <p>Loading data...</p>
          ) : (
            <Table
              data={sortedOrgData}
              columns={[
                { key: 'org', label: 'Organization' },
                {
                  key: 'family_model_count',
                  label: 'Model Count',
                  render: (value) => <span className="text-right">{value}</span>,
                },
                {
                  key: 'family_direct_children_count',
                  label: 'Direct Children',
                  render: (value) => <span className="text-right">{value}</span>,
                },
                {
                  key: 'family_all_children_count',
                  label: 'All Children',
                  render: (value) => <span className="text-right">{value}</span>,
                },
              ]}
              orderBy={orgOrderBy}
              onOrderByChange={(key) => setOrgOrderBy(key)}
              pageSize={orgPageSize}
            />
          )}
        </>
      )}
    </main>
  );
}