X-Agents/server/cmd/api/main.go

package main

import (
	"bytes"
	"encoding/json"
	"io"
	"log"
	"os"
	"path/filepath"
	"time"
	"x-agents/server/internal/config"
	"x-agents/server/internal/handler"
	"x-agents/server/internal/middleware"
	"x-agents/server/internal/model"
	"x-agents/server/internal/repository"
	"x-agents/server/internal/service"

	"github.com/gin-gonic/gin"
	swaggerFiles "github.com/swaggo/files"
	ginSwagger "github.com/swaggo/gin-swagger"
)

// Logger 日志记录器
type Logger struct {
	successLog *log.Logger
	errorLog   *log.Logger
}

func NewLogger() *Logger {
	// 创建日志目录
	today := time.Now().Format("2006-01-02")
	logDir := filepath.Join("logs", today)
	os.MkdirAll(logDir, 0755)

	// 成功日志
	successFile, _ := os.OpenFile(filepath.Join(logDir, "success.log"), os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	successLogger := log.New(successFile, "", log.Ldate|log.Ltime|log.Lshortfile)

	// 错误日志
	errorFile, _ := os.OpenFile(filepath.Join(logDir, "failure.log"), os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	errorLogger := log.New(errorFile, "", log.Ldate|log.Ltime|log.Lshortfile)

	return &Logger{
		successLog: successLogger,
		errorLog:   errorLogger,
	}
}

// LogRequest 记录请求
func (l *Logger) LogRequest(method, path, body string, status int, duration time.Duration) {
	entry := "[%s] %s %s %d %v"

	if status >= 400 {
		l.errorLog.Printf(entry, method, path, body, status, duration)
	} else {
		l.successLog.Printf(entry, method, path, body, status, duration)
	}
}

var logger *Logger

// initDefaultAdmin 初始化默认管理员用户
func initDefaultAdmin(userRepo *repository.UserRepository) {
	// 确保 admin 用户的工作空间目录存在
	ensureAdminWorkspace()

	// 检查 admin 用户是否已存在
	_, err := userRepo.FindByUsername("admin")
	if err == nil {
		log.Println("Admin user already exists")
		return
	}

	// 使用 AuthService 创建用户（会自动加密密码）
	authService := service.NewAuthService("", userRepo)
	adminUser, err := authService.Register("admin", "admin", "admin@example.com")
	if err != nil {
		log.Printf("Failed to create admin user: %v", err)
		return
	}

	// 更新角色为管理员
	adminUser.RoleID = "admin"
	userRepo.Update(adminUser)

	// 创建管理员角色
	perms, _ := json.Marshal([]int{int(model.PermissionRead), int(model.PermissionWrite), int(model.PermissionExecute), int(model.PermissionAdmin)})
	adminRole := &model.Role{
		ID:          "admin",
		Name:        "admin",
		Permissions: string(perms),
	}
	userRepo.CreateRole(adminRole)

	log.Printf("Default admin user created: id=%s, username=admin", adminUser.ID)
}

// ensureAdminWorkspace 确保 admin 用户工作空间目录存在
func ensureAdminWorkspace() {
	execPath, _ := os.Getwd()
	projectRoot := execPath

	// 如果当前目录名为 server，向上找一级
	baseName := filepath.Base(execPath)
	if baseName == "server" {
		projectRoot = filepath.Dir(execPath)
	}

	// 尝试向上查找包含 .git 的目录
	if _, err := os.Stat(filepath.Join(projectRoot, ".git")); os.IsNotExist(err) {
		for i := 0; i < 3; i++ {
			parent := filepath.Dir(projectRoot)
			if parent == projectRoot {
				break
			}
			if _, err := os.Stat(filepath.Join(parent, ".git")); err == nil {
				projectRoot = parent
				break
			}
			projectRoot = parent
		}
	}

	// 创建 admin 工作空间
	workspacePath := filepath.Join(projectRoot, "account", "admin")
	if err := os.MkdirAll(workspacePath, 0755); err != nil {
		log.Printf("Warning: failed to create admin workspace: %v", err)
		return
	}

	// 创建子目录
	for _, dir := range []string{"projects", "files", "temp"} {
		os.MkdirAll(filepath.Join(workspacePath, dir), 0755)
	}

	log.Printf("Admin workspace created at: %s", workspacePath)
}

func main() {
	// 初始化日志
	logger = NewLogger()

	// 1. 加载配置
	cfg := config.Load()
	log.Printf("=== Server starting, port=%s ===", cfg.Port)

	// 2. 初始化数据库
	db, err := config.InitDB(cfg)
	if err != nil {
		log.Fatalf("Failed to connect database: %v", err)
	}

	// 3. 自动迁移表
	db.AutoMigrate(&model.DatabaseInfo{}, &model.SubTableInfo{}, &model.ModelInfo{}, &model.KnowledgeBase{}, &model.KnowledgeDocument{}, &model.User{}, &model.Role{}, &model.Tool{})

	// 3.1 确保 users 和 roles 表存在（使用 SQL 强制创建）
	db.Exec(`
		CREATE TABLE IF NOT EXISTS roles (
			id VARCHAR(191) PRIMARY KEY,
			name VARCHAR(191) UNIQUE,
			permissions TEXT,
			created_at DATETIME(3),
			updated_at DATETIME(3)
		)
	`)
	db.Exec(`
		CREATE TABLE IF NOT EXISTS users (
			id VARCHAR(191) PRIMARY KEY,
			username VARCHAR(50) UNIQUE NOT NULL,
			password VARCHAR(191) NOT NULL,
			email VARCHAR(191),
			role_id VARCHAR(50) NOT NULL,
			is_active BOOLEAN DEFAULT TRUE,
			created_at DATETIME(3),
			updated_at DATETIME(3),
			INDEX idx_users_username (username),
			INDEX idx_users_email (email)
		)
	`)

	// 3.2 确保 tools 表存在（使用 SQL 强制创建）
	db.Exec(`
		CREATE TABLE IF NOT EXISTS tools (
			id VARCHAR(191) PRIMARY KEY,
			name VARCHAR(100) UNIQUE NOT NULL,
			description TEXT,
			category VARCHAR(50) NOT NULL,
			provider VARCHAR(100),
			security_level VARCHAR(20) DEFAULT 'safe',
			require_approval BOOLEAN DEFAULT FALSE,
			parameters TEXT,
			status VARCHAR(20) DEFAULT 'active',
			created_at DATETIME(3),
			updated_at DATETIME(3),
			INDEX idx_tools_name (name),
			INDEX idx_tools_category (category),
			INDEX idx_tools_status (status)
		)
	`)

	// 使用GORM Migrator添加缺失的列
	migrator := db.Migrator()
	if !migrator.HasColumn(&model.Tool{}, "security_level") {
		migrator.AddColumn(&model.Tool{}, "security_level")
	}
	if !migrator.HasColumn(&model.Tool{}, "require_approval") {
		migrator.AddColumn(&model.Tool{}, "require_approval")
	}
	if !migrator.HasColumn(&model.Tool{}, "parameters") {
		migrator.AddColumn(&model.Tool{}, "parameters")
	}

	log.Println("Database tables verified/created")

	// 4. 初始化 Repository
	dbRepo := repository.NewDatabaseRepository(db)
	subTableRepo := repository.NewSubTableRepository(db)
	modelRepo := repository.NewModelRepository(db)
	knowledgeRepo := repository.NewKnowledgeRepository(db)
	userRepo := repository.NewUserRepository(db)
	toolRepo := repository.NewToolRepository(db)

	// 4.1 初始化默认管理员用户
	initDefaultAdmin(userRepo)

	// 5. 初始化 Service
	dbService := service.NewDatabaseService(dbRepo, subTableRepo)
	subTableService := service.NewSubTableService(subTableRepo, dbRepo)
	neo4jService := service.NewNeo4jService(dbRepo)
	modelService := service.NewModelService(modelRepo)
	uploadService, err := service.NewUploadService(cfg)
	if err != nil {
		log.Printf("Warning: Failed to initialize upload service: %v (files will not be available)", err)
	}
	knowledgeService := service.NewKnowledgeService(knowledgeRepo, modelRepo, uploadService, cfg.PythonServiceURL, cfg.AICoreServiceAddr, cfg.MarkdownLocalPath)
	authService := service.NewAuthService(cfg.JWTSecret, userRepo)
	toolService := service.NewToolService(toolRepo)

	// 4.2 初始化默认工具
	if err := toolService.InitDefaultTools(); err != nil {
		log.Printf("Warning: Failed to init default tools: %v", err)
	} else {
		log.Println("Default tools initialized")
	}

	// 6. 初始化 Handler
	dbHandler := handler.NewDatabaseHandler(dbService)
	subTableHandler := handler.NewSubTableHandler(subTableService)
	neo4jHandler := handler.NewNeo4jHandler(neo4jService)
	modelHandler := handler.NewModelHandler(modelService)
	systemHandler := handler.NewSystemHandler()
	knowledgeHandler := handler.NewKnowledgeHandler(knowledgeService)
	authHandler := handler.NewAuthHandler(authService)
	toolHandler := handler.NewToolHandler(toolService)
	var uploadHandler *handler.UploadHandler
	if uploadService != nil {
		uploadHandler = handler.NewUploadHandler(uploadService, knowledgeRepo)
	}

	// 7. 设置路由
	r := gin.New()

	// 添加日志和恢复中间件
	r.Use(gin.Logger())
	r.Use(gin.Recovery())

	// 请求日志中间件
	r.Use(func(c *gin.Context) {
		start := time.Now()
		path := c.Request.URL.Path

		// 记录请求体
		var requestBody []byte
		if c.Request.Method == "POST" || c.Request.Method == "PUT" {
			requestBody, _ = io.ReadAll(c.Request.Body)
			c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
		}

		c.Next()

		// 记录响应日志
		latency := time.Since(start)
		status := c.Writer.Status()

		// 使用日志系统记录
		logger.LogRequest(c.Request.Method, path, string(requestBody), status, latency)
	})

	// CORS 中间件
	r.Use(func(c *gin.Context) {
		c.Writer.Header().Set("Access-Control-Allow-Origin", "*")
		c.Writer.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS, PATCH")
		c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization, X-Requested-With")

		if c.Request.Method == "OPTIONS" {
			c.AbortWithStatus(204)
			return
		}

		c.Next()
	})

	// 认证模块（无需登录）
	authGroup := r.Group("/auth")
	{
		authGroup.POST("/register", authHandler.Register)
		authGroup.POST("/login", authHandler.Login)
	}

	// 需要登录的认证模块
	authProtectedGroup := r.Group("/auth")
	authProtectedGroup.Use(middleware.Auth(cfg.JWTSecret))
	{
		authProtectedGroup.GET("/me", authHandler.GetCurrentUser)
	}

	// 用户管理模块（需要登录）
	userGroup := r.Group("/user")
	userGroup.Use(middleware.Auth(cfg.JWTSecret))
	{
		userGroup.GET("/list", authHandler.ListUsers)
		userGroup.GET("/:id", authHandler.GetUserByID)
	}

	// 数据库管理模块
	databaseGroup := r.Group("/database")
	{
		databaseGroup.GET("/list", dbHandler.List)
		databaseGroup.GET("/:id", dbHandler.GetByID)
		databaseGroup.POST("/check", dbHandler.Check)
		databaseGroup.POST("/add", dbHandler.Create)
		databaseGroup.PUT("/:id", dbHandler.Update)
		databaseGroup.DELETE("/:id", dbHandler.Delete)
		databaseGroup.POST("/graph/save", dbHandler.SaveGraph)
	}

	// 子表映射管理模块
	subTableGroup := r.Group("/sub-table")
	{
		subTableGroup.POST("/add", subTableHandler.Create)
		subTableGroup.GET("/:id", subTableHandler.GetByID)
		subTableGroup.GET("/database/:database_id", subTableHandler.ListByDatabase)
		subTableGroup.GET("/mapping/:database_id", subTableHandler.GetMappingFromFile)
		subTableGroup.GET("/ddl/:database_id", subTableHandler.GetTablesDDL)
		subTableGroup.PUT("/:id", subTableHandler.Update)
		subTableGroup.DELETE("/:id", subTableHandler.Delete)
	}

	// Neo4j 连接管理模块
	neo4jGroup := r.Group("/neo4j")
	{
		neo4jGroup.POST("/check", neo4jHandler.Check)
		neo4jGroup.POST("/graphs", neo4jHandler.GetGraphs)
		neo4jGroup.POST("/nodes", neo4jHandler.GetNodes)
		neo4jGroup.POST("/relationships", neo4jHandler.GetRelationships)
	}

	// Model 管理模块
	modelGroup := r.Group("/model")
	{
		modelGroup.GET("/list", modelHandler.List)
		modelGroup.POST("/test", modelHandler.Test)
		modelGroup.POST("/add", modelHandler.Create)
		modelGroup.GET("/:id", modelHandler.GetByID)
		modelGroup.PUT("/:id", modelHandler.Update)
		modelGroup.DELETE("/:id", modelHandler.Delete)
	}

	// 知识库管理模块
	knowledgeGroup := r.Group("/api/knowledge")
	{
		knowledgeGroup.POST("/create", knowledgeHandler.Create)
		knowledgeGroup.GET("/list", knowledgeHandler.List)
		knowledgeGroup.GET("/:id", knowledgeHandler.GetByID)
		knowledgeGroup.PUT("/:id", knowledgeHandler.Update)
		knowledgeGroup.DELETE("/:id", knowledgeHandler.Delete)
		// 文档管理
		knowledgeGroup.GET("/:id/documents", knowledgeHandler.ListDocuments)
		knowledgeGroup.POST("/:id/documents", knowledgeHandler.UploadDocument)
		knowledgeGroup.DELETE("/:id/documents/:doc_id", knowledgeHandler.DeleteDocument)
		knowledgeGroup.POST("/:id/documents/:doc_id/reparse", knowledgeHandler.ReparseDocument)
		knowledgeGroup.GET("/:id/documents/:doc_id/preview", knowledgeHandler.GetDocumentPreview)
	}

	// 系统信息模块
	r.GET("/system/info", systemHandler.GetSystemInfo)

	// 工具管理模块
	toolGroup := r.Group("/tool")
	{
		toolGroup.GET("/list", toolHandler.List)
		toolGroup.GET("/sync", toolHandler.Sync) // 手动同步
		toolGroup.GET("/:id", toolHandler.GetByID)
		toolGroup.POST("/add", toolHandler.Create)
		toolGroup.PUT("/:id", toolHandler.Update)
		toolGroup.DELETE("/:id", toolHandler.Delete)
	}

	// Swagger 文档
	r.GET("/swagger/*any", ginSwagger.WrapHandler(swaggerFiles.Handler))

	// 文件上传模块
	if uploadHandler != nil {
		// 本地文件静态服务
		if cfg.UploadMode == "local" {
			r.Static("/files", cfg.UploadLocalPath)
		}
		// 上传路由
		r.POST("/api/file_upload", uploadHandler.Upload)
		r.DELETE("/api/file_upload/:filename", uploadHandler.Delete)
		// 文件代理路由（解决 MinIO 内网和 HTTPS 问题）
		r.GET("/api/file_proxy", uploadHandler.ProxyFile)
	}

	// 8. 启动服务
	log.Printf("Server starting on :%s", cfg.Port)
	if err := r.Run(":" + cfg.Port); err != nil {
		log.Fatalf("Failed to start server: %v", err)
	}
}