import * as z from "zod"
import { tool } from "langchain"
import os from 'os'
import { promisify } from 'util'
import { exec } from 'child_process'

const execPromise = promisify(exec)

/**
 * 获取系统活动连接和Socket信息工具
 * @returns {string} - 返回Socket信息的 JSON 字符串，包括活动连接和监听端口等
 */
export const getSocketInfoTool = tool(
  async (input) => {
    try {
      const socketInfo = {
        platform: os.platform(),
        activeConnections: [],
        listeningPorts: []
      }

      // 获取活动连接（Windows）
      if (os.platform() === 'win32') {
        try {
          const { stdout } = await execPromise('netstat -ano')
          const lines = stdout.trim().split('\n').slice(4)
          socketInfo.activeConnections = lines.map(line => {
            const parts = line.trim().split(/\s+/)
            if (parts.length >= 5) {
              return {
                protocol: parts[0],
                localAddress: parts[1],
                foreignAddress: parts[2],
                state: parts[3],
                processId: parts[4]
              }
            }
            return null
          }).filter(item => item)
        } catch (error) {
          console.error('Error getting active connections:', error)
        }

        // 获取监听端口（Windows）
        try {
          const { stdout } = await execPromise('netstat -an | findstr LISTENING')
          const lines = stdout.trim().split('\n')
          socketInfo.listeningPorts = lines.map(line => {
            const parts = line.trim().split(/\s+/)
            if (parts.length >= 4) {
              return {
                protocol: parts[0],
                localAddress: parts[1],
                state: parts[3]
              }
            }
            return null
          }).filter(item => item)
        } catch (error) {
          console.error('Error getting listening ports:', error)
        }
      }
      // 获取活动连接（Linux）
      else if (os.platform() === 'linux') {
        try {
          const { stdout } = await execPromise('ss -tuln')
          const lines = stdout.trim().split('\n').slice(1)
          socketInfo.activeConnections = lines.map(line => {
            const parts = line.trim().split(/\s+/)
            if (parts.length >= 5) {
              return {
                protocol: parts[0],
                state: parts[1],
                localAddress: parts[4],
                peerAddress: parts[5]
              }
            }
            return null
          }).filter(item => item)
        } catch (error) {
          console.error('Error getting active connections:', error)
        }

        // 获取监听端口（Linux）
        try {
          const { stdout } = await execPromise('ss -tulpn')
          const lines = stdout.trim().split('\n').slice(1)
          socketInfo.listeningPorts = lines.map(line => {
            const parts = line.trim().split(/\s+/)
            if (parts.length >= 6) {
              return {
                protocol: parts[0],
                state: parts[1],
                localAddress: parts[4],
                processInfo: parts[5]
              }
            }
            return null
          }).filter(item => item)
        } catch (error) {
          console.error('Error getting listening ports:', error)
        }
      }
      // 获取活动连接（macOS）
      else if (os.platform() === 'darwin') {
        try {
          const { stdout } = await execPromise('netstat -an')
          const lines = stdout.trim().split('\n').slice(2)
          socketInfo.activeConnections = lines.map(line => {
            const parts = line.trim().split(/\s+/)
            if (parts.length >= 6) {
              return {
                protocol: parts[0],
                localAddress: parts[3],
                foreignAddress: parts[4],
                state: parts[5]
              }
            }
            return null
          }).filter(item => item)
        } catch (error) {
          console.error('Error getting active connections:', error)
        }

        // 获取监听端口（macOS）
        try {
          const { stdout } = await execPromise('netstat -an | grep LISTEN')
          const lines = stdout.trim().split('\n')
          socketInfo.listeningPorts = lines.map(line => {
            const parts = line.trim().split(/\s+/)
            if (parts.length >= 6) {
              return {
                protocol: parts[0],
                localAddress: parts[3],
                foreignAddress: parts[4],
                state: parts[5]
              }
            }
            return null
          }).filter(item => item)
        } catch (error) {
          console.error('Error getting listening ports:', error)
        }
      }

      return JSON.stringify(socketInfo, null, 2)
    } catch (error) {
      console.error('Error getting socket info:', error)
      return JSON.stringify({ error: error.message }, null, 2)
    }
  },
  {
    name: "get_socket_info",
    description: `获取系统活动连接和Socket信息，包括TCP/UDP连接、监听端口、进程ID等`,
    schema: z.object({})
  }
)