package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"
	
	"github.com/dgrijalva/jwt-go"
	_ "github.com/mattn/go-sqlite3"
	"golang.org/x/crypto/bcrypt"
	"gopkg.in/yaml.v2"
)

type Config struct {
	Server struct {
		Port       string `yaml:"port"`
		SecretKey  string `yaml:"secret_key"`
	} `yaml:"server"`
	Database struct {
		Path string `yaml:"path"`
	} `yaml:"database"`
	Admin struct {
		Username string `yaml:"username"`
		Password string `yaml:"password"`
	} `yaml:"admin"`
}

type Site struct {
	ID       int    `json:"id"`
	Name     string `json:"name"`
	URL      string `json:"url"`
	Icon     string `json:"icon,omitempty"`
	Category string `json:"category"`
	Order    int    `json:"order"` // 新增排序字段
}

type Credentials struct {
	Username string `json:"username"`
	Password string `json:"password"`
}

type Claims struct {
	Username string `json:"username"`
	jwt.StandardClaims
}

var (
	db     *sql.DB
	config *Config
)

func main() {
	// 加载配置
	var err error
	config, err = loadConfig("config/config.yaml")
	if err != nil {
		log.Fatalf("Error loading config: %v", err)
	}

	// 初始化数据库
	err = initDB(config.Database.Path)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	// 设置路由
	http.HandleFunc("/api/login", loginHandler)
	http.HandleFunc("/api/sites", sitesHandler)
	http.HandleFunc("/api/sites/order", authMiddleware(updateSitesOrderHandler))
	http.HandleFunc("/api/sites/add", authMiddleware(addSiteHandler))
	http.HandleFunc("/api/sites/delete", authMiddleware(deleteSiteHandler))
	http.HandleFunc("/api/export", authMiddleware(exportHandler))
	http.HandleFunc("/api/import", authMiddleware(importHandler))
	
	// 静态文件服务
	fs := http.FileServer(http.Dir("../frontend"))
	http.Handle("/", fs)

	// 启动服务器
	log.Printf("Starting server on :%s", config.Server.Port)
	log.Fatal(http.ListenAndServe(":"+config.Server.Port, nil))
}

func loadConfig(path string) (*Config, error) {
	cfg := &Config{}
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	d := yaml.NewDecoder(file)
	if err := d.Decode(&cfg); err != nil {
		return nil, err
	}

	return cfg, nil
}

func initDB(dbPath string) error {
	var err error
	os.MkdirAll(filepath.Dir(dbPath), 0755)
	
	db, err = sql.Open("sqlite3", dbPath)
	if err != nil {
		return err
	}

	// 创建表
	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS sites (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		name TEXT NOT NULL,
		url TEXT NOT NULL,
		icon TEXT,
		category TEXT NOT NULL,
		sort_order INTEGER DEFAULT 0
	)`)
	if err != nil {
		return err
	}

	// 创建管理员表
	_, err = db.Exec(`CREATE TABLE IF NOT EXISTS admin_users (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		username TEXT UNIQUE NOT NULL,
		password TEXT NOT NULL
	)`)
	if err != nil {
		return err
	}

	// 初始化管理员账户
	hashedPassword, _ := bcrypt.GenerateFromPassword([]byte(config.Admin.Password), bcrypt.DefaultCost)
	db.Exec("INSERT OR IGNORE INTO admin_users (username, password) VALUES (?, ?)", 
		config.Admin.Username, string(hashedPassword))

	return nil
}

func loginHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var creds Credentials
	if err := json.NewDecoder(r.Body).Decode(&creds); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	// 验证用户
	var storedPassword string
	err := db.QueryRow("SELECT password FROM admin_users WHERE username = ?", creds.Username).Scan(&storedPassword)
	if err != nil {
		http.Error(w, "Invalid credentials", http.StatusUnauthorized)
		return
	}

	if err := bcrypt.CompareHashAndPassword([]byte(storedPassword), []byte(creds.Password)); err != nil {
		http.Error(w, "Invalid credentials", http.StatusUnauthorized)
		return
	}

	// 创建JWT令牌
	expirationTime := time.Now().Add(24 * time.Hour)
	claims := &Claims{
		Username: creds.Username,
		StandardClaims: jwt.StandardClaims{
			ExpiresAt: expirationTime.Unix(),
		},
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	tokenString, err := token.SignedString([]byte(config.Server.SecretKey))
	if err != nil {
		http.Error(w, "Error generating token", http.StatusInternalServerError)
		return
	}

	// 返回令牌
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{
		"token": tokenString,
	})
}

func authMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		authHeader := r.Header.Get("Authorization")
		if authHeader == "" {
			http.Error(w, "Authorization header required", http.StatusUnauthorized)
			return
		}

		tokenString := strings.TrimPrefix(authHeader, "Bearer ")
		claims := &Claims{}

		token, err := jwt.ParseWithClaims(tokenString, claims, func(token *jwt.Token) (interface{}, error) {
			return []byte(config.Server.SecretKey), nil
		})

		if err != nil || !token.Valid {
			http.Error(w, "Invalid token", http.StatusUnauthorized)
			return
		}

		next(w, r)
	}
}

func sitesHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	rows, err := db.Query("SELECT id, name, url, icon, category, sort_order FROM sites ORDER BY category, sort_order, name")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer rows.Close()

	var sites []Site
	for rows.Next() {
		var s Site
		if err := rows.Scan(&s.ID, &s.Name, &s.URL, &s.Icon, &s.Category, &s.Order); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		sites = append(sites, s)
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(sites)
}

func updateSitesOrderHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var sites []Site
	if err := json.NewDecoder(r.Body).Decode(&sites); err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	tx, err := db.Begin()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	for _, site := range sites {
		_, err = tx.Exec("UPDATE sites SET sort_order = ? WHERE id = ?", site.Order, site.ID)
		if err != nil {
			tx.Rollback()
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}

	if err := tx.Commit(); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
}

// 其他处理函数(addSiteHandler, deleteSiteHandler等)保持不变...