//go:build llm_script

package main

import (
	"database/sql"
	"log"
	"os"
	"strings"
	"time"

	_ "github.com/lib/pq"
)

type ModelPricing struct {
	ModelID         string
	ModelName       string
	ProviderName    string
	ProviderCountry string
	OperatorName    string
	OperatorType    string
	Region          string
	Currency        string
	InputPrice      float64
	OutputPrice     float64
	ContextLength   int
	IsFree          bool
	SourceURL       string
	ModelSourceURL  string
	ReleaseDate     string
	DateConfidence  string
	DateSourceKind  string
	Modality        string
	SceneTags       []string
}

func releaseDateValue(raw string) any {
	if raw == "" {
		return nil
	}
	parsed, err := time.Parse("2006-01-02", raw)
	if err != nil {
		return nil
	}
	return parsed
}

type zhipuModelMetadata struct {
	Prefix         string
	ReleaseDate    string
	ModelSourceURL string
	DateConfidence string
	DateSourceKind string
}

var zhipuModelMetadataRules = []zhipuModelMetadata{
	{
		Prefix:         "glm-5-turbo",
		ReleaseDate:    "2026-03-15",
		ModelSourceURL: "https://www.zhipuai.cn/en/research/155",
		DateConfidence: "official_primary",
		DateSourceKind: "official_announcement",
	},
	{
		Prefix:         "glm-5.1",
		ReleaseDate:    "2026-04-07",
		ModelSourceURL: "https://www.zhipuai.cn/zh/research",
		DateConfidence: "official_primary",
		DateSourceKind: "official_announcement",
	},
	{
		Prefix:         "glm-5",
		ReleaseDate:    "2026-02-11",
		ModelSourceURL: "https://www.zhipuai.cn/zh/research/154",
		DateConfidence: "official_primary",
		DateSourceKind: "official_announcement",
	},
	{
		Prefix:         "glm-4.7-flash",
		ReleaseDate:    "2026-01-19",
		ModelSourceURL: "https://www.zhipuai.cn/zh/news/148",
		DateConfidence: "official_primary",
		DateSourceKind: "official_announcement",
	},
	{
		Prefix:         "glm-4.7",
		ReleaseDate:    "2025-12-21",
		ModelSourceURL: "https://www.zhipuai.cn/zh/research",
		DateConfidence: "official_primary",
		DateSourceKind: "official_announcement",
	},
	{
		Prefix:         "glm-4.6v",
		ReleaseDate:    "2025-12-07",
		ModelSourceURL: "https://www.zhipuai.cn/zh/research/144",
		DateConfidence: "official_primary",
		DateSourceKind: "official_announcement",
	},
	{
		Prefix:         "glm-tts",
		ReleaseDate:    "2025-12-10",
		ModelSourceURL: "https://www.zhipuai.cn/zh/research",
		DateConfidence: "official_primary",
		DateSourceKind: "official_announcement",
	},
	{
		Prefix:         "glm-4.5-air",
		DateConfidence: "unknown",
		DateSourceKind: "catalog_backfill",
	},
	{
		Prefix:         "glm-4.5v",
		DateConfidence: "unknown",
		DateSourceKind: "catalog_backfill",
	},
	{
		Prefix:         "glm-4-air",
		DateConfidence: "unknown",
		DateSourceKind: "catalog_backfill",
	},
	{
		Prefix:         "glm-4-long",
		DateConfidence: "unknown",
		DateSourceKind: "catalog_backfill",
	},
	{
		Prefix:         "glm-4v",
		DateConfidence: "unknown",
		DateSourceKind: "catalog_backfill",
	},
	{
		Prefix:         "chatglm3",
		DateConfidence: "unknown",
		DateSourceKind: "catalog_backfill",
	},
	{
		Prefix:         "glm-4-9b",
		DateConfidence: "unknown",
		DateSourceKind: "catalog_backfill",
	},
	{
		Prefix:         "glm-4-0520",
		ReleaseDate:    "2024-06-05",
		DateConfidence: "secondary_authoritative",
		DateSourceKind: "secondary_authoritative_report",
	},
	{
		Prefix:         "glm-realtime",
		DateConfidence: "unknown",
		DateSourceKind: "catalog_backfill",
	},
	{
		Prefix:         "glm-asr",
		DateConfidence: "unknown",
		DateSourceKind: "catalog_backfill",
	},
	{
		Prefix:         "glm-4-voice",
		ReleaseDate:    "2024-10-25",
		DateConfidence: "secondary_authoritative",
		DateSourceKind: "secondary_authoritative_report",
	},
}

func enrichZhipuModelMetadata(model ModelPricing) ModelPricing {
	for _, metadata := range zhipuModelMetadataRules {
		if strings.HasPrefix(model.ModelID, metadata.Prefix) {
			if metadata.ReleaseDate != "" {
				model.ReleaseDate = metadata.ReleaseDate
			}
			if metadata.ModelSourceURL != "" {
				model.ModelSourceURL = metadata.ModelSourceURL
			}
			if metadata.DateConfidence != "" {
				model.DateConfidence = metadata.DateConfidence
			}
			if metadata.DateSourceKind != "" {
				model.DateSourceKind = metadata.DateSourceKind
			}
			return model
		}
	}
	if model.ModelSourceURL == "" {
		model.ModelSourceURL = model.SourceURL
	}
	if model.DateConfidence == "" {
		model.DateConfidence = "unknown"
	}
	if model.DateSourceKind == "" {
		model.DateSourceKind = "unknown"
	}
	return model
}

func hasExplicitModelMetadata(model ModelPricing) bool {
	return strings.TrimSpace(model.ReleaseDate) != "" ||
		firstNonEmpty(model.ModelSourceURL) != "" && model.ModelSourceURL != model.SourceURL ||
		strings.TrimSpace(model.DateConfidence) != "" && model.DateConfidence != "unknown" ||
		strings.TrimSpace(model.DateSourceKind) != "" && model.DateSourceKind != "unknown"
}

func main() {
	dsn := os.Getenv("DATABASE_URL")
	if dsn == "" {
		dsn = "postgres://long@/llm_intelligence?host=/var/run/postgresql"
	}

	db, err := sql.Open("postgres", dsn)
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	// 智谱AI定价数据（从第一次无头浏览器抓取中提取）
	prices := []ModelPricing{
		// GLM-5.1系列
		{ModelID: "glm-5.1", ModelName: "GLM-5.1", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 6.0, OutputPrice: 24.0, ContextLength: 32000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话", "推理", "代码"}},
		{ModelID: "glm-5.1-32k", ModelName: "GLM-5.1 (32K+)", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 8.0, OutputPrice: 28.0, ContextLength: 200000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话", "推理", "代码"}},

		// GLM-5-Turbo
		{ModelID: "glm-5-turbo", ModelName: "GLM-5-Turbo", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 5.0, OutputPrice: 22.0, ContextLength: 32000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话", "推理"}},
		{ModelID: "glm-5-turbo-32k", ModelName: "GLM-5-Turbo (32K+)", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 7.0, OutputPrice: 26.0, ContextLength: 200000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话", "推理"}},

		// GLM-5
		{ModelID: "glm-5", ModelName: "GLM-5", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 4.0, OutputPrice: 18.0, ContextLength: 32000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话", "推理"}},
		{ModelID: "glm-5-32k", ModelName: "GLM-5 (32K+)", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 6.0, OutputPrice: 22.0, ContextLength: 200000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话", "推理"}},

		// GLM-4.7
		{ModelID: "glm-4.7", ModelName: "GLM-4.7", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 2.0, OutputPrice: 8.0, ContextLength: 32000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话", "推理", "代码"}},
		{ModelID: "glm-4.7-32k", ModelName: "GLM-4.7 (32K+)", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 4.0, OutputPrice: 16.0, ContextLength: 200000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话", "推理", "代码"}},

		// GLM-4.5-Air
		{ModelID: "glm-4.5-air", ModelName: "GLM-4.5-Air", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 0.8, OutputPrice: 2.0, ContextLength: 32000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话"}},
		{ModelID: "glm-4.5-air-32k", ModelName: "GLM-4.5-Air (32K+)", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 1.2, OutputPrice: 8.0, ContextLength: 128000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话"}},

		// GLM-4.7-FlashX
		{ModelID: "glm-4.7-flashx", ModelName: "GLM-4.7-FlashX", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 0.5, OutputPrice: 3.0, ContextLength: 200000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话"}},

		// GLM-4.7-Flash (Free)
		{ModelID: "glm-4.7-flash", ModelName: "GLM-4.7-Flash", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 0, OutputPrice: 0, ContextLength: 200000, IsFree: true, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话"}},

		// GLM-4.6V (Vision)
		{ModelID: "glm-4.6v", ModelName: "GLM-4.6V", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 2.0, OutputPrice: 6.0, ContextLength: 8000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "multimodal", SceneTags: []string{"视觉", "对话"}},

		// GLM-4.6V-FlashX
		{ModelID: "glm-4.6v-flashx", ModelName: "GLM-4.6V-FlashX", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 0.15, OutputPrice: 1.5, ContextLength: 8000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "multimodal", SceneTags: []string{"视觉", "对话"}},

		// GLM-4.5V
		{ModelID: "glm-4.5v", ModelName: "GLM-4.5V", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 2.0, OutputPrice: 6.0, ContextLength: 32000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "multimodal", SceneTags: []string{"视觉", "对话"}},

		// GLM-4系列 (Legacy)
		{ModelID: "glm-4-0520", ModelName: "GLM-4-0520", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 100.0, OutputPrice: 50.0, ContextLength: 128000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话", "推理"}},
		{ModelID: "glm-4-air", ModelName: "GLM-4-Air", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 0.5, OutputPrice: 0.25, ContextLength: 128000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话"}},
		{ModelID: "glm-4-airx", ModelName: "GLM-4-AirX", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 10.0, OutputPrice: 10.0, ContextLength: 8000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话", "极速"}},
		{ModelID: "glm-4-long", ModelName: "GLM-4-Long", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 1.0, OutputPrice: 0.5, ContextLength: 1000000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话", "长文本"}},

		// GLM-4V (Vision Legacy)
		{ModelID: "glm-4v-plus", ModelName: "GLM-4V-Plus", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 4.0, OutputPrice: 4.0, ContextLength: 8000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "multimodal", SceneTags: []string{"视觉", "对话"}},
		{ModelID: "glm-4v", ModelName: "GLM-4V", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 50.0, OutputPrice: 50.0, ContextLength: 2000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "multimodal", SceneTags: []string{"视觉", "对话"}},

		// ChatGLM3
		{ModelID: "chatglm3-6b", ModelName: "ChatGLM3-6B", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 0, OutputPrice: 0, ContextLength: 8000, IsFree: true, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话"}},

		// GLM-4-9B
		{ModelID: "glm-4-9b", ModelName: "GLM-4-9B", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 0, OutputPrice: 0, ContextLength: 8000, IsFree: true, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "text", SceneTags: []string{"对话"}},

		// GLM-Realtime
		{ModelID: "glm-realtime-flash", ModelName: "GLM-Realtime-Flash", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 0.18, OutputPrice: 0.18, ContextLength: 8000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "audio", SceneTags: []string{"实时", "音视频"}},
		{ModelID: "glm-realtime-air", ModelName: "GLM-Realtime-Air", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 0.3, OutputPrice: 0.3, ContextLength: 8000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "audio", SceneTags: []string{"实时", "音视频"}},

		// GLM-TTS
		{ModelID: "glm-tts", ModelName: "GLM-TTS", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 2.0, OutputPrice: 0, ContextLength: 8000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "audio", SceneTags: []string{"语音合成"}},
		{ModelID: "glm-tts-clone", ModelName: "GLM-TTS-Clone", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 6.0, OutputPrice: 0, ContextLength: 8000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "audio", SceneTags: []string{"音色克隆"}},

		// GLM-ASR
		{ModelID: "glm-asr-2512", ModelName: "GLM-ASR-2512", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 16.0, OutputPrice: 0, ContextLength: 8000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "audio", SceneTags: []string{"语音识别"}},

		// GLM-4-Voice
		{ModelID: "glm-4-voice", ModelName: "GLM-4-Voice", ProviderName: "Zhipu AI", ProviderCountry: "CN", OperatorName: "Zhipu", OperatorType: "official", Region: "CN", Currency: "CNY", InputPrice: 80.0, OutputPrice: 80.0, ContextLength: 8000, SourceURL: "https://open.bigmodel.cn/pricing", Modality: "audio", SceneTags: []string{"语音模型"}},
	}
	batchID := "manual-seed"

	log.Printf("Importing %d Zhipu AI models...", len(prices))

	// Save to database
	for _, rawPricing := range prices {
		p := enrichZhipuModelMetadata(rawPricing)
		// Find or create provider
		var providerID int64
		err := db.QueryRow("SELECT id FROM model_provider WHERE name = $1", p.ProviderName).Scan(&providerID)
		if err == sql.ErrNoRows {
			err = db.QueryRow(
				"INSERT INTO model_provider (name, country, website, status) VALUES ($1, $2, $3, 'active') RETURNING id",
				p.ProviderName, p.ProviderCountry, "",
			).Scan(&providerID)
		}
		if err != nil {
			log.Printf("Provider error: %v", err)
			continue
		}

		// Find or create operator
		var operatorID int64
		err = db.QueryRow("SELECT id FROM operator WHERE name = $1", p.OperatorName).Scan(&operatorID)
		if err == sql.ErrNoRows {
			err = db.QueryRow(
				"INSERT INTO operator (name, country, status) VALUES ($1, $2, 'active') RETURNING id",
				p.OperatorName, p.ProviderCountry,
			).Scan(&operatorID)
		}
		if err != nil {
			log.Printf("Operator error: %v", err)
			continue
		}

		// Find or create model
		var modelID int64
		err = db.QueryRow("SELECT id FROM models WHERE external_id = $1", p.ModelID).Scan(&modelID)
		if err == sql.ErrNoRows {
			err = db.QueryRow(
				`INSERT INTO models (external_id, name, provider_id, modality, context_length, status, source, batch_id, source_url, release_date, date_confidence, date_source_kind) 
				 VALUES ($1, $2, $3, $4, $5, 'active', $6, $7, $8, $9, $10, $11) RETURNING id`,
				p.ModelID, p.ModelName, providerID, p.Modality, p.ContextLength, p.OperatorName, batchID, firstNonEmpty(p.ModelSourceURL, p.SourceURL), releaseDateValue(p.ReleaseDate), p.DateConfidence, p.DateSourceKind,
			).Scan(&modelID)
		}
		if err != nil {
			log.Printf("Model error for %s: %v", p.ModelID, err)
			continue
		}
		if _, err := db.Exec(
			`UPDATE models
			 SET source_url = CASE
			                    WHEN $4 THEN $2
			                    ELSE COALESCE(NULLIF(source_url, ''), $2)
			                  END,
			     release_date = CASE
			                      WHEN $4 THEN $3::date
			                      ELSE COALESCE(release_date, $3::date)
			                    END,
			     date_confidence = CASE
			                         WHEN $4 THEN $5
			                         ELSE COALESCE(NULLIF(date_confidence, ''), $5, 'unknown')
			                       END,
			     date_source_kind = CASE
			                          WHEN $4 THEN $6
			                          ELSE COALESCE(NULLIF(date_source_kind, ''), $6, 'unknown')
			                        END,
			     updated_at = CURRENT_TIMESTAMP
			 WHERE id = $1`,
			modelID, firstNonEmpty(p.ModelSourceURL, p.SourceURL), releaseDateValue(p.ReleaseDate), hasExplicitModelMetadata(p), p.DateConfidence, p.DateSourceKind,
		); err != nil {
			log.Printf("Model metadata update error for %s: %v", p.ModelID, err)
		}

		// Insert pricing
		sourceType := p.OperatorType
		freeQuota := ""
		freeLimitations := "[]"
		rateLimit := "{}"
		if p.IsFree {
			sourceType = "free_tier"
			freeQuota = "Imported free-tier pricing entry"
			freeLimitations = `["See source_url for current quota and policy"]`
		}
		_, err = db.Exec(
			`INSERT INTO region_pricing 
			 (model_id, operator_id, region, currency, input_price_per_mtok, output_price_per_mtok, is_free, effective_date, source_url, source_type, free_quota, free_limitations, rate_limit)
			 VALUES ($1, $2, $3, $4, $5, $6, $7, CURRENT_DATE, $8, $9, $10, $11, $12)
			 ON CONFLICT (model_id, operator_id, region, currency, effective_date) 
			 DO UPDATE SET input_price_per_mtok = EXCLUDED.input_price_per_mtok,
			               output_price_per_mtok = EXCLUDED.output_price_per_mtok,
			               is_free = EXCLUDED.is_free,
			               source_type = EXCLUDED.source_type,
			               free_quota = EXCLUDED.free_quota,
			               free_limitations = EXCLUDED.free_limitations,
			               rate_limit = EXCLUDED.rate_limit,
			               updated_at = CURRENT_TIMESTAMP`,
			modelID, operatorID, p.Region, p.Currency, p.InputPrice, p.OutputPrice, p.IsFree, p.SourceURL,
			sourceType, freeQuota, freeLimitations, rateLimit,
		)
		if err != nil {
			log.Printf("Pricing error for %s: %v", p.ModelID, err)
			continue
		}
	}

	log.Printf("Successfully imported %d Zhipu AI models", len(prices))
}

func firstNonEmpty(values ...string) string {
	for _, value := range values {
		if value != "" {
			return value
		}
	}
	return ""
}