#!/usr/bin/env node

// # 显示所有表
// ./pg-query.js -t

// # 显示表结构
// ./pg-query.js -d influencers

// # 显示样本数据，限制5行
// ./pg-query.js -s posts -l 5

// # 查看表记录数
// ./pg-query.js -c posts

// # 显示索引
// ./pg-query.js -i posts

// # 显示外键
// ./pg-query.js -f posts

// # 显示引用
// ./pg-query.js -r influencers

// # 执行自定义查询
// ./pg-query.js -q "SELECT * FROM influencers WHERE platform = 'Instagram' LIMIT 5"

// # 执行SQL文件
// ./pg-query.js -e schema.sql

const { Client } = require('pg');
const path = require('path');
const fs = require('fs');
const yargs = require('yargs/yargs');
const { hideBin } = require('yargs/helpers');

// 加载.env文件 - 使用正确的相对路径
require('dotenv').config({ path: path.resolve(__dirname, '../../../.env') });

// 显示连接信息(不含密码)以便调试
function getConnectionString() {
  // 使用.env中的DATABASE_URL
  const databaseUrl = process.env.DATABASE_URL;
  
  if (!databaseUrl) {
    console.error('错误: 未找到DATABASE_URL环境变量');
    process.exit(1);
  }
  
  // 显示连接信息但隐藏密码
  const sanitizedUrl = databaseUrl.replace(/:[^:@]+@/, ':***@');
  console.log(`使用连接: ${sanitizedUrl}`);
  
  return databaseUrl;
}

// 创建一个新的客户端
async function runQuery(query, params = []) {
  const client = new Client({
    connectionString: getConnectionString()
  });
  
  try {
    await client.connect();
    console.log('数据库连接成功');
    const result = await client.query(query, params);
    return result.rows;
  } catch (err) {
    console.error('查询执行错误:', err.message);
    return null;
  } finally {
    await client.end();
  }
}

// 显示所有表
async function showTables() {
  const query = "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public' ORDER BY table_name;";
  const tables = await runQuery(query);
  
  if (tables && tables.length > 0) {
    console.log('数据库中的表:');
    console.table(tables);
  } else {
    console.log('没有找到表或连接失败');
  }
}

// 显示表结构
async function showTableStructure(tableName) {
  const query = `
    SELECT 
      column_name AS "列名", 
      data_type AS "数据类型", 
      CASE WHEN is_nullable = 'YES' THEN '允许为空' ELSE '不允许为空' END AS "是否可空", 
      column_default AS "默认值",
      character_maximum_length AS "最大长度"
    FROM 
      information_schema.columns 
    WHERE 
      table_schema = 'public' AND 
      table_name = $1
    ORDER BY 
      ordinal_position;
  `;
  
  const columns = await runQuery(query, [tableName]);
  
  if (columns && columns.length > 0) {
    console.log(`表 ${tableName} 的结构:`);
    console.table(columns);
  } else {
    console.log(`表 ${tableName} 不存在或连接失败`);
  }
}

// 显示样本数据
async function showSampleData(tableName, limit = 10) {
  const query = `SELECT * FROM "${tableName}" LIMIT ${limit};`;
  const data = await runQuery(query);
  
  if (data && data.length > 0) {
    console.log(`表 ${tableName} 的样本数据 (${limit} 行):`);
    console.table(data);
  } else {
    console.log(`表 ${tableName} 为空或不存在`);
  }
}

// 显示记录计数
async function showRecordCount(tableName) {
  const query = `SELECT COUNT(*) AS "记录数" FROM "${tableName}";`;
  const count = await runQuery(query);
  
  if (count) {
    console.log(`表 ${tableName} 的记录数:`);
    console.table(count);
  } else {
    console.log(`表 ${tableName} 不存在或连接失败`);
  }
}

// 显示索引信息
async function showIndexes(tableName) {
  const query = `
    SELECT
      indexname AS "索引名称",
      indexdef AS "索引定义"
    FROM
      pg_indexes
    WHERE
      tablename = $1
    ORDER BY
      indexname;
  `;
  
  const indexes = await runQuery(query, [tableName]);
  
  if (indexes && indexes.length > 0) {
    console.log(`表 ${tableName} 的索引:`);
    console.table(indexes);
  } else {
    console.log(`表 ${tableName} 没有索引或不存在`);
  }
}

// 显示外键
async function showForeignKeys(tableName) {
  const query = `
    SELECT
      conname AS "外键名称",
      pg_get_constraintdef(oid) AS "外键定义"
    FROM
      pg_constraint
    WHERE
      conrelid = $1::regclass AND contype = 'f';
  `;
  
  const foreignKeys = await runQuery(query, [tableName]);
  
  if (foreignKeys && foreignKeys.length > 0) {
    console.log(`表 ${tableName} 的外键:`);
    console.table(foreignKeys);
  } else {
    console.log(`表 ${tableName} 没有外键或不存在`);
  }
}

// 显示引用当前表的外键
async function showReferencingKeys(tableName) {
  const query = `
    SELECT
      c.conname AS "外键名称",
      t.relname AS "引用表",
      pg_get_constraintdef(c.oid) AS "外键定义"
    FROM
      pg_constraint c
    JOIN
      pg_class t ON c.conrelid = t.oid
    WHERE
      c.confrelid = $1::regclass AND c.contype = 'f';
  `;
  
  const referencingKeys = await runQuery(query, [tableName]);
  
  if (referencingKeys && referencingKeys.length > 0) {
    console.log(`引用表 ${tableName} 的外键关系:`);
    console.table(referencingKeys);
  } else {
    console.log(`没有找到引用表 ${tableName} 的外键关系`);
  }
}

// 执行自定义查询
async function executeQuery(query) {
  const result = await runQuery(query);
  
  if (result) {
    console.log('查询结果:');
    console.table(result);
  } else {
    console.log('查询执行失败或无结果');
  }
}

// 执行SQL文件
async function executeSqlFile(filename) {
  try {
    const sql = fs.readFileSync(filename, 'utf8');
    console.log(`执行SQL文件: ${filename}`);
    await executeQuery(sql);
  } catch (err) {
    console.error(`执行SQL文件失败: ${err.message}`);
  }
}

// 主函数
async function main() {
  try {
    const argv = yargs(hideBin(process.argv))
      .usage('PostgreSQL 查询工具\n\n用法: $0 [选项]')
      .option('t', {
        alias: 'tables',
        describe: '显示所有表',
        type: 'boolean'
      })
      .option('d', {
        alias: 'describe',
        describe: '显示表结构',
        type: 'string'
      })
      .option('s', {
        alias: 'sample',
        describe: '显示表样本数据',
        type: 'string'
      })
      .option('l', {
        alias: 'limit',
        describe: '样本数据行数限制',
        type: 'number',
        default: 10
      })
      .option('c', {
        alias: 'count',
        describe: '计算表中的记录数',
        type: 'string'
      })
      .option('i', {
        alias: 'indexes',
        describe: '显示表索引',
        type: 'string'
      })
      .option('f', {
        alias: 'foreign-keys',
        describe: '显示表外键关系',
        type: 'string'
      })
      .option('r', {
        alias: 'references',
        describe: '显示引用此表的外键',
        type: 'string'
      })
      .option('q', {
        alias: 'query',
        describe: '执行自定义SQL查询',
        type: 'string'
      })
      .option('e', {
        alias: 'execute-file',
        describe: '执行SQL文件',
        type: 'string'
      })
      .example('$0 -t', '显示所有表')
      .example('$0 -d influencers', '显示influencers表结构')
      .example('$0 -s posts -l 5', '显示posts表前5行数据')
      .epilog('更多信息请访问项目文档')
      .help()
      .alias('h', 'help')
      .argv;

    if (argv.tables) {
      await showTables();
    } else if (argv.describe) {
      await showTableStructure(argv.describe);
    } else if (argv.sample) {
      await showSampleData(argv.sample, argv.limit);
    } else if (argv.count) {
      await showRecordCount(argv.count);
    } else if (argv.indexes) {
      await showIndexes(argv.indexes);
    } else if (argv.foreignKeys) {
      await showForeignKeys(argv.foreignKeys);
    } else if (argv.references) {
      await showReferencingKeys(argv.references);
    } else if (argv.query) {
      await executeQuery(argv.query);
    } else if (argv.executeFile) {
      await executeSqlFile(argv.executeFile);
    } else {
      yargs(hideBin(process.argv)).showHelp();
    }
  } catch (err) {
    console.error('程序执行错误:', err.message);
    process.exit(1);
  }
}

// 执行主函数
main().catch(err => {
  console.error('程序执行错误:', err);
  process.exit(1);
});