qsdk/tushare/new.go

package tushare

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/jedib0t/go-pretty/v6/table"
)

// TushareClient Tushare API 客户端
type TushareClient struct {
	Token   string `json:"token"`    // API 访问令牌
	BaseUrl string `json:"base_url"` // API 基础 URL
	client  *http.Client
}

// TushareReq Tushare API 请求结构
type TushareReq struct {
	APIName string         `json:"api_name"` // API 接口名称
	Token   string         `json:"token"`    // API 访问令牌
	Params  map[string]any `json:"params"`   // 请求参数
	Fields  []string       `json:"fields"`   // 返回字段列表
}

// TushareResp Tushare API 响应结构
type TushareResp struct {
	RequestID string           `json:"request_id"` // 请求 ID
	Code      int              `json:"code"`       // 响应码 (0 表示成功)
	Data      *TushareRespData `json:"data"`       // 响应数据
	Msg       string           `json:"msg"`        // 响应消息
}

// TushareRespData Tushare API 响应数据结构
type TushareRespData struct {
	Fields  []string `json:"fields"`   // 字段列表
	Headers []string `json:"headers"`  // 表头（中文描述）
	Items   [][]any  `json:"items"`    // 数据项（二维数组）
	HasMore bool     `json:"has_more"` // 是否有更多数据
	Count   int      `json:"count"`    // 返回数据条数
}

// NewClient 创建并初始化 Tushare 客户端
// 返回配置好 Token 和基础 URL 的客户端实例
func NewClient(token string) *TushareClient {
	if token == "" {
		token = os.Getenv("TUSHARE_TOKEN")
	}
	return &TushareClient{
		Token:   token,
		BaseUrl: "http://api.tushare.pro",
		client:  &http.Client{Timeout: 30 * time.Second},
	}
}

func (c *TushareClient) httpDo() *http.Client {
	if c.client != nil {
		return c.client
	}
	return &http.Client{Timeout: 30 * time.Second}
}

// postTushare 发送 JSON 请求并解析为 TushareResp（不含业务码以外的处理）。
func (c *TushareClient) postTushare(reqBytes []byte) (*TushareResp, error) {
	resp, err := c.httpDo().Post(c.BaseUrl, "application/json", bytes.NewReader(reqBytes))
	if err != nil {
		return nil, fmt.Errorf("#100 发送请求失败：%w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("#101 读取响应失败：%w", err)
	}

	var tushareResp TushareResp
	if err := json.Unmarshal(body, &tushareResp); err != nil {
		return nil, fmt.Errorf("#101 响应解析失败：%v - %s", err, string(body))
	}
	if tushareResp.Code != 0 {
		return nil, fmt.Errorf("#102 API 返回错误：%s", tushareResp.Msg)
	}
	return &tushareResp, nil
}

// Do 执行 Tushare API 请求
// req: 请求参数，包含 API 名称、参数等
// fieldsVals: 字段配置列表，每个元素是一个 map，key 为字段名，value 为字段中文描述
// 返回：响应数据和错误信息
func (c *TushareClient) Do(req TushareReq, fieldsVals []map[string]string) (*TushareRespData, error) {
	// 构建请求体
	payload := TushareReq{
		APIName: req.APIName,
		Token:   c.Token,
		Params:  req.Params,
	}

	// 提取字段名和对应的中文表头
	n := 0
	for _, setting := range fieldsVals {
		n += len(setting)
	}
	fields := make([]string, 0, n)
	headers := make([]string, 0, n)
	for _, setting := range fieldsVals {
		for key, value := range setting {
			fields = append(fields, key)
			headers = append(headers, value)
		}
	}
	payload.Fields = fields

	reqBytes, err := json.Marshal(payload)
	if err != nil {
		return nil, fmt.Errorf("#100 请求序列化失败：%w", err)
	}

	tushareResp, err := c.postTushare(reqBytes)
	if err != nil {
		return nil, err
	}
	if tushareResp.Data != nil {
		tushareResp.Data.Headers = headers
	}
	return tushareResp.Data, nil
}

// Exec 执行 Tushare API 请求
// ApiName: API 名称
// fields: 字段列表，用逗号分隔
// Params: 请求参数
// 返回：响应数据和错误信息
func (c *TushareClient) Exec(ApiName string, fields string, Params map[string]any) (*TushareRespData, error) {
	fields = strings.TrimSpace(fields)
	if fields == "" {
		return nil, fmt.Errorf("#100 字段列表不能为空")
	}
	raw := strings.Split(fields, ",")
	fieldsList := make([]string, 0, len(raw))
	for _, f := range raw {
		f = strings.TrimSpace(f)
		if f != "" {
			fieldsList = append(fieldsList, f)
		}
	}
	if len(fieldsList) == 0 {
		return nil, fmt.Errorf("#100 字段列表不能为空")
	}

	payload := TushareReq{
		APIName: ApiName,
		Token:   c.Token,
		Params:  Params,
		Fields:  fieldsList,
	}

	reqBytes, err := json.Marshal(payload)
	if err != nil {
		return nil, fmt.Errorf("#100 请求序列化失败：%w", err)
	}

	tushareResp, err := c.postTushare(reqBytes)
	if err != nil {
		return nil, err
	}
	return tushareResp.Data, nil
}

// Map 将响应数据转换为 map 切片
// 每个 map 代表一行数据，key 为字段名，value 为对应的值
func (c *TushareRespData) Map() []map[string]any {
	if c == nil || len(c.Items) == 0 || len(c.Fields) == 0 {
		return make([]map[string]any, 0)
	}
	result := make([]map[string]any, 0, len(c.Items))
	for _, item := range c.Items {
		rowMap := make(map[string]any, len(c.Fields))
		for idx, fn := range c.Fields {
			if idx < len(item) {
				rowMap[fn] = item[idx]
			}
		}
		result = append(result, rowMap)
	}
	return result
}

// Json 将响应数据转换为 JSON 字节数组
// 返回：JSON 格式的字节数据和错误信息
func (c *TushareRespData) Json() ([]byte, error) {
	data := c.Map()
	return json.MarshalIndent(data, "", "  ")
}

// Output 将响应数据格式化为表格输出
// title: 表格标题
// 返回：格式化的表格写入器
func (c *TushareRespData) Output(title string) table.Writer {
	tw := table.NewWriter()
	tw.SetStyle(table.StyleLight)
	tw.SetTitle(title)
	if c == nil {
		return tw
	}

	n := len(c.Fields)
	if nh := len(c.Headers); nh < n {
		n = nh
	}
	headerRow := make(table.Row, 0, n)
	for idx := 0; idx < n; idx++ {
		headerRow = append(headerRow, c.Headers[idx]+"("+c.Fields[idx]+")")
	}
	tw.AppendHeader(headerRow)

	// 添加数据行
	for _, item := range c.Items {
		tw.AppendRow(item)
	}
	return tw
}