big-qmt/go-client/apps/zt/logic/state.go

package logic

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
)

type SymbolState struct {
	Code     string  `json:"code"`
	BaseQty  int     `json:"base_qty"`
	BaseCost float64 `json:"base_cost"`
	AddQty   int     `json:"add_qty"`
	AddCost  float64 `json:"add_cost"`
	Pending  string  `json:"pending"`
}

type filePayload struct {
	Version int            `json:"version"`
	Data    map[string]any `json:"data"`
}

type ZTState struct {
	path      string
	Items     map[string]*SymbolState
	LoadError string
	mu        sync.Mutex
}

var (
	statesMu sync.Mutex
	states   = map[string]*ZTState{}
)

func getState(cfg Config) *ZTState {
	statesMu.Lock()
	defer statesMu.Unlock()
	if s, ok := states[cfg.AccountID]; ok {
		return s
	}
	s := loadZTState(cfg.DataDir, cfg.AccountID)
	states[cfg.AccountID] = s
	return s
}

func loadZTState(dataDir, accountID string) *ZTState {
	st := &ZTState{
		path:  filepath.Join(dataDir, fmt.Sprintf("zt_%s_state.json", accountID)),
		Items: map[string]*SymbolState{},
	}
	raw, err := os.ReadFile(st.path)
	if err != nil {
		if os.IsNotExist(err) {
			return st
		}
		st.rebuild(err)
		return st
	}
	var payload filePayload
	if err := json.Unmarshal(raw, &payload); err != nil || payload.Version != 1 {
		st.rebuild(fmt.Errorf("状态文件版本无效"))
		return st
	}
	data := payload.Data
	if data == nil {
		st.rebuild(fmt.Errorf("状态文件内容无效"))
		return st
	}
	symbolsAny, _ := data["symbols"]
	symbols, _ := symbolsAny.(map[string]any)
	if symbols == nil {
		if _, ok := data["code"]; ok {
			symbols = map[string]any{}
		} else {
			symbols = data
		}
	}
	for code, value := range symbols {
		m, ok := value.(map[string]any)
		if !ok {
			continue
		}
		item := &SymbolState{Code: code}
		b, _ := json.Marshal(m)
		_ = json.Unmarshal(b, item)
		item.Code = code
		st.Items[code] = item
	}
	return st
}

func (s *ZTState) rebuild(err error) {
	if err := os.Remove(s.path); err != nil && !os.IsNotExist(err) {
		s.LoadError = err.Error()
		logf("ERROR", "[ZT][状态] 状态文件重建失败: %s", s.LoadError)
		return
	}
	s.Items = map[string]*SymbolState{}
	if saveErr := s.saveUnlocked(); saveErr != nil {
		s.LoadError = fmt.Sprintf("%v；重建失败: %v", err, saveErr)
		logf("ERROR", "[ZT][状态] 状态文件重建失败: %s", s.LoadError)
		return
	}
	logf("WARNING", "[ZT][状态] 状态文件损坏，已删除并重建: %v", err)
}

func (s *ZTState) Get(code string) *SymbolState {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.Items[code]
}

func (s *ZTState) Ensure(code string) *SymbolState {
	s.mu.Lock()
	defer s.mu.Unlock()
	if item, ok := s.Items[code]; ok {
		return item
	}
	item := &SymbolState{Code: code}
	s.Items[code] = item
	return item
}

func (s *ZTState) Remove(code string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	delete(s.Items, code)
}

func (s *ZTState) Codes() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]string, 0, len(s.Items))
	for code := range s.Items {
		out = append(out, code)
	}
	return out
}

func (s *ZTState) Save() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.LoadError != "" {
		return
	}
	if err := s.saveUnlocked(); err != nil {
		logf("ERROR", "[ZT][状态] 保存失败: %v", err)
	}
}

func (s *ZTState) saveUnlocked() error {
	symbols := map[string]any{}
	for code, item := range s.Items {
		symbols[code] = item
	}
	payload := filePayload{Version: 1, Data: map[string]any{"symbols": symbols}}
	raw, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(s.path), 0o755); err != nil {
		return err
	}
	tmp := s.path + ".tmp"
	if err := os.WriteFile(tmp, raw, 0o644); err != nil {
		return err
	}
	return replaceFile(tmp, s.path)
}

func replaceFile(tmp, dest string) error {
	if err := os.Rename(tmp, dest); err == nil {
		return nil
	}
	_ = os.Remove(dest)
	return os.Rename(tmp, dest)
}