package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"

	_ "github.com/mattn/go-sqlite3"
	"github.com/modelcontextprotocol/go-sdk/mcp"
)

var db *sql.DB

// ==================== 数据库初始化 ====================

func initDB() error {
	// 获取可执行文件所在目录
	execPath, err := os.Executable()
	if err != nil {
		execPath = "."
	}
	dbPath := filepath.Join(filepath.Dir(execPath), "demo.db")

	// 如果数据库不存在，使用当前目录
	if _, err := os.Stat(dbPath); os.IsNotExist(err) {
		dbPath = "demo.db"
	}

	db, err = sql.Open("sqlite3", dbPath)
	if err != nil {
		return err
	}

	// 创建示例表
	_, err = db.Exec(`
		CREATE TABLE IF NOT EXISTS users (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			name TEXT NOT NULL,
			email TEXT UNIQUE,
			age INTEGER,
			created_at DATETIME DEFAULT CURRENT_TIMESTAMP
		);
		
		CREATE TABLE IF NOT EXISTS products (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			name TEXT NOT NULL,
			price REAL,
			stock INTEGER DEFAULT 0
		);
	`)
	if err != nil {
		return err
	}

	// 插入示例数据（如果表为空）
	var count int
	db.QueryRow("SELECT COUNT(*) FROM users").Scan(&count)
	if count == 0 {
		_, err = db.Exec(`
			INSERT INTO users (name, email, age) VALUES 
			('张三', 'zhangsan@example.com', 28),
			('李四', 'lisi@example.com', 32),
			('王五', 'wangwu@example.com', 25);
			
			INSERT INTO products (name, price, stock) VALUES 
			('iPhone 15', 6999.00, 100),
			('MacBook Pro', 14999.00, 50),
			('AirPods Pro', 1899.00, 200);
		`)
		if err != nil {
			return err
		}
	}

	return nil
}

// ==================== 工具输入结构 ====================

// ListTablesInput - 列出所有表
type ListTablesInput struct{}

// QueryInput - 执行 SQL 查询
type QueryInput struct {
	SQL string `json:"sql" mcp:"SQL query to execute (SELECT only for safety)"`
}

// GetUserInput - 获取用户
type GetUserInput struct {
	ID int `json:"id" mcp:"User ID to retrieve"`
}

// AddUserInput - 添加用户
type AddUserInput struct {
	Name  string `json:"name" mcp:"User name"`
	Email string `json:"email" mcp:"User email address"`
	Age   int    `json:"age" mcp:"User age"`
}

// DeleteUserInput - 删除用户
type DeleteUserInput struct {
	ID int `json:"id" mcp:"User ID to delete"`
}

// ==================== 辅助函数 ====================

func textResult(text string) (*mcp.CallToolResult, any, error) {
	return &mcp.CallToolResult{
		Content: []mcp.Content{
			&mcp.TextContent{Text: text},
		},
	}, nil, nil
}

func errorResult(text string) (*mcp.CallToolResult, any, error) {
	return &mcp.CallToolResult{
		Content: []mcp.Content{
			&mcp.TextContent{Text: text},
		},
		IsError: true,
	}, nil, nil
}

// ==================== 工具实现 ====================

// ListTables - 列出数据库中的所有表
func ListTables(ctx context.Context, req *mcp.CallToolRequest, input ListTablesInput) (*mcp.CallToolResult, any, error) {
	rows, err := db.Query("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'")
	if err != nil {
		return errorResult(fmt.Sprintf("Failed to list tables: %v", err))
	}
	defer rows.Close()

	var tables []string
	for rows.Next() {
		var name string
		rows.Scan(&name)
		tables = append(tables, name)
	}

	result, _ := json.MarshalIndent(tables, "", "  ")
	return textResult(fmt.Sprintf("Tables in database:\n%s", string(result)))
}

// QueryDB - 执行 SQL 查询
func QueryDB(ctx context.Context, req *mcp.CallToolRequest, input QueryInput) (*mcp.CallToolResult, any, error) {
	// 安全检查：只允许 SELECT 查询
	sqlUpper := strings.ToUpper(strings.TrimSpace(input.SQL))
	if !strings.HasPrefix(sqlUpper, "SELECT") {
		return errorResult("Only SELECT queries are allowed for safety")
	}

	rows, err := db.Query(input.SQL)
	if err != nil {
		return errorResult(fmt.Sprintf("Query failed: %v", err))
	}
	defer rows.Close()

	// 获取列名
	columns, _ := rows.Columns()

	// 读取所有行
	var results []map[string]interface{}
	for rows.Next() {
		values := make([]interface{}, len(columns))
		valuePtrs := make([]interface{}, len(columns))
		for i := range values {
			valuePtrs[i] = &values[i]
		}

		rows.Scan(valuePtrs...)

		row := make(map[string]interface{})
		for i, col := range columns {
			row[col] = values[i]
		}
		results = append(results, row)
	}

	jsonResult, _ := json.MarshalIndent(results, "", "  ")
	return textResult(fmt.Sprintf("Query results (%d rows):\n%s", len(results), string(jsonResult)))
}

// GetUser - 获取单个用户
func GetUser(ctx context.Context, req *mcp.CallToolRequest, input GetUserInput) (*mcp.CallToolResult, any, error) {
	var id int
	var name, email string
	var age int
	var createdAt string

	err := db.QueryRow("SELECT id, name, email, age, created_at FROM users WHERE id = ?", input.ID).
		Scan(&id, &name, &email, &age, &createdAt)
	if err == sql.ErrNoRows {
		return errorResult(fmt.Sprintf("User with ID %d not found", input.ID))
	}
	if err != nil {
		return errorResult(fmt.Sprintf("Query failed: %v", err))
	}

	user := map[string]interface{}{
		"id":         id,
		"name":       name,
		"email":      email,
		"age":        age,
		"created_at": createdAt,
	}

	jsonResult, _ := json.MarshalIndent(user, "", "  ")
	return textResult(string(jsonResult))
}

// AddUser - 添加用户
func AddUser(ctx context.Context, req *mcp.CallToolRequest, input AddUserInput) (*mcp.CallToolResult, any, error) {
	result, err := db.Exec("INSERT INTO users (name, email, age) VALUES (?, ?, ?)",
		input.Name, input.Email, input.Age)
	if err != nil {
		return errorResult(fmt.Sprintf("Failed to add user: %v", err))
	}

	id, _ := result.LastInsertId()
	return textResult(fmt.Sprintf("User added successfully with ID: %d", id))
}

// DeleteUser - 删除用户
func DeleteUser(ctx context.Context, req *mcp.CallToolRequest, input DeleteUserInput) (*mcp.CallToolResult, any, error) {
	result, err := db.Exec("DELETE FROM users WHERE id = ?", input.ID)
	if err != nil {
		return errorResult(fmt.Sprintf("Failed to delete user: %v", err))
	}

	affected, _ := result.RowsAffected()
	if affected == 0 {
		return errorResult(fmt.Sprintf("User with ID %d not found", input.ID))
	}

	return textResult(fmt.Sprintf("User with ID %d deleted successfully", input.ID))
}

// ==================== 主函数 ====================

func main() {
	// 初始化数据库
	if err := initDB(); err != nil {
		log.Fatalf("Failed to initialize database: %v", err)
	}
	defer db.Close()

	// 创建 MCP Server
	server := mcp.NewServer(
		&mcp.Implementation{
			Name:    "mcp-demo-go",
			Version: "1.0.0",
		},
		nil,
	)

	// 注册工具
	mcp.AddTool(server, &mcp.Tool{
		Name:        "list_tables",
		Description: "List all tables in the SQLite database",
	}, ListTables)

	mcp.AddTool(server, &mcp.Tool{
		Name:        "query_db",
		Description: "Execute a SELECT SQL query on the database",
	}, QueryDB)

	mcp.AddTool(server, &mcp.Tool{
		Name:        "get_user",
		Description: "Get a user by ID from the users table",
	}, GetUser)

	mcp.AddTool(server, &mcp.Tool{
		Name:        "add_user",
		Description: "Add a new user to the database",
	}, AddUser)

	mcp.AddTool(server, &mcp.Tool{
		Name:        "delete_user",
		Description: "Delete a user from the database by ID",
	}, DeleteUser)

	// 运行服务器
	log.Println("MCP Demo Go Server (SQLite) is running...")
	if err := server.Run(context.Background(), &mcp.StdioTransport{}); err != nil {
		log.Fatalf("Server error: %v", err)
	}
}