llm-intelligence/scripts/import_zhipu_coding_plan.go

//go:build llm_script && !scripts_pkg

package main

import (
	"database/sql"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"
)

type zhipuCodingPlanImportConfig struct {
	OverviewURL      string
	PromotionURL     string
	OverviewFixture  string
	PromotionFixture string
	DryRun           bool
	Timeout          time.Duration
}

func main() {
	loadSubscriptionImportEnv()

	var overviewURL string
	var promotionURL string
	var overviewFixture string
	var promotionFixture string
	var dryRun bool
	var timeoutSeconds int

	flag.StringVar(&overviewURL, "overview-url", defaultZhipuCodingPlanOverviewURL, "智谱 Coding Plan 概览 URL")
	flag.StringVar(&promotionURL, "promotion-url", defaultZhipuCodingPlanPromotionURL, "智谱 Coding Plan 活动页 URL")
	flag.StringVar(&overviewFixture, "overview-fixture", "", "智谱 Coding Plan 概览样例文件")
	flag.StringVar(&promotionFixture, "promotion-fixture", "", "智谱 Coding Plan 活动页样例文件")
	flag.BoolVar(&dryRun, "dry-run", false, "仅解析并打印摘要，不写入数据库")
	flag.IntVar(&timeoutSeconds, "timeout", 20, "请求超时（秒）")
	flag.Parse()

	cfg := zhipuCodingPlanImportConfig{
		OverviewURL:      overviewURL,
		PromotionURL:     promotionURL,
		OverviewFixture:  overviewFixture,
		PromotionFixture: promotionFixture,
		DryRun:           dryRun,
		Timeout:          time.Duration(timeoutSeconds) * time.Second,
	}

	var db *sql.DB
	var err error
	if !cfg.DryRun {
		db, err = subscriptionImportDB()
		if err != nil {
			fmt.Fprintf(os.Stderr, "open db: %v\n", err)
			os.Exit(1)
		}
		defer db.Close()
	}

	if err := runZhipuCodingPlanImport(cfg, db, os.Stdout); err != nil {
		fmt.Fprintf(os.Stderr, "import_zhipu_coding_plan: %v\n", err)
		os.Exit(1)
	}
}

func runZhipuCodingPlanImport(cfg zhipuCodingPlanImportConfig, db *sql.DB, out io.Writer) error {
	client := &http.Client{Timeout: cfg.Timeout}
	overviewRaw, err := fetchSubscriptionPage(cfg.OverviewURL, cfg.OverviewFixture, client)
	if err != nil {
		return err
	}
	promotionRaw, err := fetchSubscriptionPage(cfg.PromotionURL, cfg.PromotionFixture, client)
	if err != nil {
		return err
	}

	records, err := parseZhipuCodingPlanCatalog(overviewRaw, promotionRaw)
	if err != nil {
		return err
	}
	if cfg.DryRun {
		_, err = fmt.Fprintf(out, "source=zhipu-coding-plan-import plans=%d provider=%s operator=%s dry_run=true\n", len(records), records[0].ProviderName, records[0].OperatorName)
		return err
	}
	if db == nil {
		return fmt.Errorf("db is required when dry-run=false")
	}
	if err := upsertSubscriptionImportRecords(db, records); err != nil {
		return err
	}

	var tableRows int
	if err := db.QueryRow(`SELECT COUNT(*) FROM subscription_plan`).Scan(&tableRows); err != nil {
		return fmt.Errorf("count subscription_plan: %w", err)
	}
	_, err = fmt.Fprintf(out, "source=zhipu-coding-plan-import plans=%d provider=%s operator=%s table_rows=%d dry_run=false\n", len(records), records[0].ProviderName, records[0].OperatorName, tableRows)
	return err
}