weather-station/cmd/weatherstation/main.go

package main

import (
	"context"
	"flag"
	"log"
	"os"
	"sync"
	"time"
	"weatherstation/internal/config"
	"weatherstation/internal/database"
	"weatherstation/internal/forecast"
	"weatherstation/internal/selftest"
	"weatherstation/internal/server"
	"weatherstation/internal/tools"
)

func main() {
	// 命令行参数
	var webOnly = flag.Bool("web", false, "只启动Web服务器（Gin）")
	var udpOnly = flag.Bool("udp", false, "只启动UDP服务器")
	// 调试：回填10分钟表
	var doBackfill = flag.Bool("backfill", false, "将16s原始数据聚合写入10分钟表（调试用途）")
	var bfStation = flag.String("station", "", "指定站点ID（为空则全站回填）")
	var bfFrom = flag.String("from", "", "回填起始时间，格式：YYYY-MM-DD HH:MM:SS")
	var bfTo = flag.String("to", "", "回填结束时间，格式：YYYY-MM-DD HH:MM:SS")
	var bfWrap = flag.Float64("wrap", 0, "回绕一圈对应毫米值（mm），<=0 则降级为仅计当前值")
	// 自检控制
	var noSelftest = flag.Bool("no-selftest", false, "跳过启动自检")
	var selftestOnly = flag.Bool("selftest_only", false, "仅执行自检后退出")
	// 预报抓取
	var forecastOnly = flag.Bool("forecast_only", false, "仅执行一次open-meteo拉取并退出")
	var caiyunOnly = flag.Bool("caiyun_only", false, "仅执行一次彩云拉取并退出")
	var forecastDay = flag.String("forecast_day", "", "按日期抓取当天0点到当前时间+3h（格式YYYY-MM-DD）")
	// 历史数据补完
	var historicalOnly = flag.Bool("historical_only", false, "仅执行历史数据补完并退出")
	var historicalStart = flag.String("historical_start", "", "历史数据开始日期（格式YYYY-MM-DD）")
	var historicalEnd = flag.String("historical_end", "", "历史数据结束日期（格式YYYY-MM-DD）")
	// 覆盖风：使用彩云实况替换导出中的风速/风向
	var useWindOverride = flag.Bool("wind", false, "使用彩云实况覆盖导出CSV中的风速/风向")
	flag.Parse()

	// 设置日志
	server.SetupLogger()

	// 初始化数据库连接
	_ = database.GetDB() // 确保PostgreSQL连接已初始化
	defer database.Close()
	// 初始化MySQL连接（如果配置存在）
	_ = database.GetMySQL()
	defer database.CloseMySQL()

	// 启动前自检
	if !*noSelftest {
		if err := selftest.Run(context.Background()); err != nil {
			log.Fatalf("启动自检失败: %v", err)
		}
		if *selftestOnly {
			log.Println("自检完成，按 --selftest_only 要求退出")
			return
		}
	}

	// 单次 open-meteo 拉取
	if *forecastOnly {
		if err := forecast.RunOpenMeteoFetch(context.Background()); err != nil {
			log.Fatalf("open-meteo 拉取失败: %v", err)
		}
		log.Println("open-meteo 拉取完成")
		return
	}

	// 单次 彩云 拉取（token 从环境变量 CAIYUN_TOKEN 或命令行 -caiyun_token 读取）
	if *caiyunOnly {
		token := os.Getenv("CAIYUN_TOKEN")
		if token == "" {
			// 退回配置
			token = config.GetConfig().Forecast.CaiyunToken
			if token == "" {
				log.Fatalf("未提供彩云 token，请设置环境变量 CAIYUN_TOKEN 或配置文件 forecast.caiyun_token")
			}
		}
		if err := forecast.RunCaiyunFetch(context.Background(), token); err != nil {
			log.Fatalf("caiyun 拉取失败: %v", err)
		}
		log.Println("caiyun 拉取完成")
		return
	}

	// 工具：按日期抓取当天0点到当前时间+3小时（两家）
	if *forecastDay != "" {
		if err := tools.RunForecastFetchForDay(context.Background(), *forecastDay); err != nil {
			log.Fatalf("forecast_day 运行失败: %v", err)
		}
		log.Println("forecast_day 完成")
		return
	}

	// 历史数据补完
	if *historicalOnly {
		if *historicalStart == "" || *historicalEnd == "" {
			log.Fatalln("历史数据补完需要提供 --historical_start 与 --historical_end 日期（格式YYYY-MM-DD）")
		}
		if err := forecast.RunOpenMeteoHistoricalFetch(context.Background(), *historicalStart, *historicalEnd); err != nil {
			log.Fatalf("历史数据补完失败: %v", err)
		}
		log.Println("历史数据补完成")
		return
	}

	// Backfill 调试路径
	if *doBackfill {
		if *bfFrom == "" || *bfTo == "" {
			log.Fatalln("backfill 需要提供 --from 与 --to 时间")
		}
		loc, _ := time.LoadLocation("Asia/Shanghai")
		if loc == nil {
			loc = time.FixedZone("CST", 8*3600)
		}
		fromT, err := time.ParseInLocation("2006-01-02 15:04:05", *bfFrom, loc)
		if err != nil {
			log.Fatalf("解析from失败: %v", err)
		}
		toT, err := time.ParseInLocation("2006-01-02 15:04:05", *bfTo, loc)
		if err != nil {
			log.Fatalf("解析to失败: %v", err)
		}
		ctx := context.Background()
		if err := tools.RunBackfill10Min(ctx, tools.BackfillOptions{
			StationID:     *bfStation,
			FromTime:      fromT,
			ToTime:        toT,
			WrapCycleMM:   *bfWrap,
			BucketMinutes: 10,
		}); err != nil {
			log.Fatalf("回填失败: %v", err)
		}
		log.Println("回填完成")
		return
	}

	// 根据命令行参数启动服务
	startExporterBackground := func(wg *sync.WaitGroup) {
		if wg != nil {
			wg.Add(1)
		}
		go func() {
			defer func() {
				if wg != nil {
					wg.Done()
				}
			}()
			log.Println("启动数据导出器（10分钟）...")
			ctx := context.Background()
			// 处理 --wind 覆盖
			var opts tools.ExporterOptions
			if *useWindOverride {
				token := os.Getenv("CAIYUN_TOKEN")
				if token == "" {
					token = config.GetConfig().Forecast.CaiyunToken
				}
				if token == "" {
					log.Println("警告: 指定了 --wind 但未提供彩云 token，忽略风覆盖")
				} else {
					opts.OverrideWindWithCaiyun = true
					opts.CaiyunToken = token
				}
			}
			exporter := tools.NewExporterWithOptions(opts)
			if err := exporter.Start(ctx); err != nil {
				log.Printf("导出器退出: %v", err)
			}
		}()
	}

	if *webOnly {
		// 只启动Web服务器 + 导出器
		startExporterBackground(nil)
		log.Println("启动Web服务器模式...")
		if err := server.StartGinServer(); err != nil {
			log.Fatalf("启动Web服务器失败: %v", err)
		}
	} else if *udpOnly {
		// 只启动UDP服务器 + 导出器
		startExporterBackground(nil)
		log.Println("启动UDP服务器模式...")
		if err := server.StartUDPServer(); err != nil {
			log.Fatalf("启动UDP服务器失败: %v", err)
		}
	} else {
		// 同时启动UDP和Web服务器 + 导出器
		log.Println("启动完整模式：UDP + Web服务器 + 导出器...")

		var wg sync.WaitGroup
		wg.Add(2)

		// 启动UDP服务器
		go func() {
			defer wg.Done()
			log.Println("正在启动UDP服务器...")
			if err := server.StartUDPServer(); err != nil {
				log.Printf("UDP服务器异常退出: %v", err)
			}
		}()

		// 启动Web服务器
		go func() {
			defer wg.Done()
			log.Println("正在启动Web服务器...")
			if err := server.StartGinServer(); err != nil {
				log.Printf("Web服务器异常退出: %v", err)
			}
		}()

		startExporterBackground(&wg)
		wg.Wait()
	}
}