const express = require('express');
const { Pool } = require('pg');
const cors = require('cors');
const bodyParser = require('body-parser');
const { v4: uuidv4 } = require('uuid');

const app = express();
const PORT = process.env.PORT || 3000;

// 中间件
app.use(cors());
app.use(bodyParser.json());
app.use(express.static('public'));

// PostgreSQL 连接池
const pool = new Pool({
  host: process.env.DB_HOST || 'postgres',
  port: process.env.DB_PORT || 5432,
  database: process.env.DB_NAME || 'agent_tasks',
  user: process.env.DB_USER || 'postgres',
  password: process.env.DB_PASSWORD || 'postgres',
  max: 20,
  idleTimeoutMillis: 30000,
  connectionTimeoutMillis: 2000,
});

// 测试数据库连接
pool.connect((err, client, release) => {
  if (err) {
    console.error('❌ 数据库连接失败:', err.message);
    process.exit(1);
  } else {
    console.log('✅ 已连接到 PostgreSQL 数据库');
    release();
    initDatabase();
  }
});

// 创建表结构
async function initDatabase() {
  const client = await pool.connect();
  try {
    await client.query('BEGIN');
    
    // 智能体表
    await client.query(`
      CREATE TABLE IF NOT EXISTS agents (
        id VARCHAR(255) PRIMARY KEY,
        name VARCHAR(255) NOT NULL,
        type VARCHAR(100),
        capabilities JSONB,
        status VARCHAR(50) DEFAULT 'idle',
        registered_at BIGINT,
        last_heartbeat BIGINT
      )
    `);

    // 任务表
    await client.query(`
      CREATE TABLE IF NOT EXISTS tasks (
        id VARCHAR(255) PRIMARY KEY,
        title VARCHAR(500) NOT NULL,
        description TEXT,
        priority VARCHAR(50) DEFAULT 'normal',
        status VARCHAR(50) DEFAULT 'pending',
        created_by VARCHAR(255),
        assigned_to VARCHAR(255),
        created_at BIGINT,
        updated_at BIGINT,
        completed_at BIGINT,
        result TEXT
      )
    `);

    // 创建索引
    await client.query(`
      CREATE INDEX IF NOT EXISTS idx_agents_status ON agents(status)
    `);
    await client.query(`
      CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status)
    `);
    await client.query(`
      CREATE INDEX IF NOT EXISTS idx_tasks_assigned ON tasks(assigned_to)
    `);

    await client.query('COMMIT');
    console.log('✅ 数据库表已初始化');
  } catch (err) {
    await client.query('ROLLBACK');
    console.error('❌ 数据库初始化失败:', err);
  } finally {
    client.release();
  }
}

// ==================== API 路由 ====================

// 1. 智能体注册
app.post('/api/agents/register', async (req, res) => {
  const { name, type, capabilities } = req.body;
  
  if (!name) {
    return res.status(400).json({ error: '智能体名称不能为空' });
  }

  const id = uuidv4();
  const now = Date.now();

  try {
    await pool.query(
      `INSERT INTO agents (id, name, type, capabilities, status, registered_at, last_heartbeat) 
       VALUES ($1, $2, $3, $4, 'idle', $5, $6)`,
      [id, name, type || 'general', JSON.stringify(capabilities || []), now, now]
    );
    
    res.json({ 
      id, 
      name, 
      type: type || 'general',
      status: 'idle',
      message: '智能体注册成功' 
    });
  } catch (err) {
    res.status(500).json({ error: err.message });
  }
});

// 2. 获取所有智能体
app.get('/api/agents', async (req, res) => {
  try {
    const result = await pool.query('SELECT * FROM agents ORDER BY registered_at DESC');
    res.json(result.rows.map(row => ({
      ...row,
      capabilities: row.capabilities || []
    })));
  } catch (err) {
    res.status(500).json({ error: err.message });
  }
});

// 3. 更新智能体状态（心跳）
app.post('/api/agents/:id/heartbeat', async (req, res) => {
  const { id } = req.params;
  const { status } = req.body;
  
  try {
    const result = await pool.query(
      'UPDATE agents SET status = $1, last_heartbeat = $2 WHERE id = $3',
      [status || 'idle', Date.now(), id]
    );
    
    if (result.rowCount === 0) {
      return res.status(404).json({ error: '智能体不存在' });
    }
    
    res.json({ message: '心跳更新成功' });
  } catch (err) {
    res.status(500).json({ error: err.message });
  }
});

// 4. 创建任务
app.post('/api/tasks', async (req, res) => {
  const { title, description, priority, created_by } = req.body;
  
  if (!title) {
    return res.status(400).json({ error: '任务标题不能为空' });
  }

  const id = uuidv4();
  const now = Date.now();

  try {
    await pool.query(
      `INSERT INTO tasks (id, title, description, priority, status, created_by, created_at, updated_at) 
       VALUES ($1, $2, $3, $4, 'pending', $5, $6, $7)`,
      [id, title, description || '', priority || 'normal', created_by || 'human', now, now]
    );
    
    res.json({ 
      id, 
      title,
      status: 'pending',
      message: '任务创建成功' 
    });
  } catch (err) {
    res.status(500).json({ error: err.message });
  }
});

// 5. 获取所有任务
app.get('/api/tasks', async (req, res) => {
  const { status } = req.query;
  
  try {
    let query = 'SELECT * FROM tasks';
    let params = [];
    
    if (status) {
      query += ' WHERE status = $1';
      params.push(status);
    }
    
    query += ' ORDER BY created_at DESC';
    
    const result = await pool.query(query, params);
    res.json(result.rows);
  } catch (err) {
    res.status(500).json({ error: err.message });
  }
});

// 6. 领取任务
app.post('/api/tasks/:id/claim', async (req, res) => {
  const { id } = req.params;
  const { agent_id } = req.body;
  
  if (!agent_id) {
    return res.status(400).json({ error: '需要提供智能体ID' });
  }

  const client = await pool.connect();
  try {
    await client.query('BEGIN');
    
    const taskResult = await client.query('SELECT status FROM tasks WHERE id = $1', [id]);
    
    if (taskResult.rows.length === 0) {
      await client.query('ROLLBACK');
      return res.status(404).json({ error: '任务不存在' });
    }
    
    if (taskResult.rows[0].status !== 'pending') {
      await client.query('ROLLBACK');
      return res.status(400).json({ error: '任务已被领取或已完成' });
    }

    await client.query(
      'UPDATE tasks SET status = $1, assigned_to = $2, updated_at = $3 WHERE id = $4',
      ['in_progress', agent_id, Date.now(), id]
    );
    
    await client.query('UPDATE agents SET status = $1 WHERE id = $2', ['busy', agent_id]);
    
    await client.query('COMMIT');
    res.json({ message: '任务领取成功' });
  } catch (err) {
    await client.query('ROLLBACK');
    res.status(500).json({ error: err.message });
  } finally {
    client.release();
  }
});

// 7. 更新任务状态
app.put('/api/tasks/:id', async (req, res) => {
  const { id } = req.params;
  const { status, result } = req.body;
  
  const client = await pool.connect();
  try {
    await client.query('BEGIN');
    
    const now = Date.now();
    const updates = [];
    const params = [];
    let paramIndex = 1;
    
    if (status) {
      updates.push(`status = $${paramIndex++}`);
      params.push(status);
    }
    if (result !== undefined) {
      updates.push(`result = $${paramIndex++}`);
      params.push(result);
    }
    if (status === 'completed') {
      updates.push(`completed_at = $${paramIndex++}`);
      params.push(now);
    }
    
    updates.push(`updated_at = $${paramIndex++}`);
    params.push(now);
    params.push(id);
    
    const updateResult = await client.query(
      `UPDATE tasks SET ${updates.join(', ')} WHERE id = $${paramIndex}`,
      params
    );
    
    if (updateResult.rowCount === 0) {
      await client.query('ROLLBACK');
      return res.status(404).json({ error: '任务不存在' });
    }
    
    // 如果任务完成，将智能体状态改为 idle
    if (status === 'completed') {
      const taskResult = await client.query('SELECT assigned_to FROM tasks WHERE id = $1', [id]);
      if (taskResult.rows.length > 0 && taskResult.rows[0].assigned_to) {
        await client.query('UPDATE agents SET status = $1 WHERE id = $2', ['idle', taskResult.rows[0].assigned_to]);
      }
    }
    
    await client.query('COMMIT');
    res.json({ message: '任务状态更新成功' });
  } catch (err) {
    await client.query('ROLLBACK');
    res.status(500).json({ error: err.message });
  } finally {
    client.release();
  }
});

// 8. 获取单个任务详情
app.get('/api/tasks/:id', async (req, res) => {
  try {
    const result = await pool.query('SELECT * FROM tasks WHERE id = $1', [req.params.id]);
    if (result.rows.length === 0) {
      return res.status(404).json({ error: '任务不存在' });
    }
    res.json(result.rows[0]);
  } catch (err) {
    res.status(500).json({ error: err.message });
  }
});

// 9. 统计信息
app.get('/api/stats', async (req, res) => {
  try {
    const stats = {
      total_agents: 0,
      active_agents: 0,
      total_tasks: 0,
      pending_tasks: 0,
      in_progress_tasks: 0,
      completed_tasks: 0
    };
    
    const results = await Promise.all([
      pool.query('SELECT COUNT(*) as count FROM agents'),
      pool.query('SELECT COUNT(*) as count FROM agents WHERE status = $1', ['busy']),
      pool.query('SELECT COUNT(*) as count FROM tasks'),
      pool.query('SELECT COUNT(*) as count FROM tasks WHERE status = $1', ['pending']),
      pool.query('SELECT COUNT(*) as count FROM tasks WHERE status = $1', ['in_progress']),
      pool.query('SELECT COUNT(*) as count FROM tasks WHERE status = $1', ['completed'])
    ]);
    
    stats.total_agents = parseInt(results[0].rows[0].count);
    stats.active_agents = parseInt(results[1].rows[0].count);
    stats.total_tasks = parseInt(results[2].rows[0].count);
    stats.pending_tasks = parseInt(results[3].rows[0].count);
    stats.in_progress_tasks = parseInt(results[4].rows[0].count);
    stats.completed_tasks = parseInt(results[5].rows[0].count);
    
    res.json(stats);
  } catch (err) {
    res.status(500).json({ error: err.message });
  }
});

// 健康检查
app.get('/health', async (req, res) => {
  try {
    await pool.query('SELECT 1');
    res.json({ status: 'healthy', database: 'connected' });
  } catch (err) {
    res.status(503).json({ status: 'unhealthy', error: err.message });
  }
});

// 启动服务器
app.listen(PORT, () => {
  console.log(`🚀 多智能体任务服务已启动: http://localhost:${PORT}`);
});

// 优雅关闭
process.on('SIGINT', async () => {
  console.log('\n正在关闭服务...');
  await pool.end();
  console.log('数据库连接已关闭');
  process.exit(0);
});