const { Pool } = require('pg');
require('dotenv').config({ path: '.env.local' });

const SYSTEM_USER_ID = '00000000-0000-0000-0000-000000000000';

(async () => {
  const pool = new Pool({ connectionString: process.env.LOCAL_DB_URL });
  const client = await pool.connect();
  try {
    await client.query('BEGIN');

    const byUniqueCustomModel = await client.query(`
      WITH model_owners AS (
        SELECT
          params->>'model' AS model,
          ARRAY_AGG(DISTINCT user_id) FILTER (WHERE user_id IS NOT NULL AND user_id <> $1) AS owner_ids
        FROM works
        WHERE status = 'completed'
          AND params->>'model' LIKE 'custom:%'
          AND EXISTS (SELECT 1 FROM profiles p WHERE p.id = works.user_id)
        GROUP BY params->>'model'
      ),
      unique_model_owners AS (
        SELECT model, owner_ids[1] AS owner_id
        FROM model_owners
        WHERE CARDINALITY(owner_ids) = 1
      ),
      updated AS (
        UPDATE works w
        SET user_id = umo.owner_id,
            updated_at = NOW()
        FROM unique_model_owners umo
        WHERE w.user_id = $1
          AND w.status = 'completed'
          AND w.is_public = true
          AND w.params->>'model' = umo.model
        RETURNING w.id, w.user_id, w.params->>'model' AS model
      )
      SELECT COUNT(*)::int AS fixed_count, json_agg(updated) AS rows
      FROM updated
    `, [SYSTEM_USER_ID]);

    const byExactPromptTime = await client.query(`
      WITH candidates AS (
        SELECT DISTINCT ON (public.id)
          public.id AS public_id,
          private.user_id AS owner_user_id,
          ABS(EXTRACT(EPOCH FROM (private.created_at - public.created_at))) AS time_distance
        FROM works public
        JOIN works private
          ON private.id <> public.id
         AND private.user_id IS NOT NULL
         AND private.user_id <> $1
         AND COALESCE(private.prompt, '') <> ''
         AND private.prompt = public.prompt
         AND private.created_at BETWEEN public.created_at - INTERVAL '30 minutes' AND public.created_at + INTERVAL '30 minutes'
        JOIN profiles p ON p.id = private.user_id
        WHERE public.user_id = $1
          AND public.is_public = true
          AND public.status = 'completed'
        ORDER BY public.id, time_distance
      ),
      unambiguous AS (
        SELECT public_id, MIN(owner_user_id::text)::uuid AS owner_user_id
        FROM candidates
        GROUP BY public_id
        HAVING COUNT(DISTINCT owner_user_id) = 1
      ),
      updated AS (
        UPDATE works w
        SET user_id = unambiguous.owner_user_id,
            updated_at = NOW()
        FROM unambiguous
        WHERE w.id = unambiguous.public_id
          AND w.user_id = $1
        RETURNING w.id, w.user_id, w.params->>'model' AS model
      )
      SELECT COUNT(*)::int AS fixed_count, json_agg(updated) AS rows
      FROM updated
    `, [SYSTEM_USER_ID]);

    await client.query('COMMIT');

    console.log(JSON.stringify({
      byUniqueCustomModel: byUniqueCustomModel.rows[0] || { fixed_count: 0, rows: [] },
      byExactPromptTime: byExactPromptTime.rows[0] || { fixed_count: 0, rows: [] },
    }, null, 2));
  } catch (error) {
    await client.query('ROLLBACK');
    throw error;
  } finally {
    client.release();
    await pool.end();
  }
})().catch(error => {
  console.error(error);
  process.exit(1);
});