Agent-System/internal/mineru/client.go

package mineru

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

const (
	defaultEndpoint     = "https://mineru.net"
	defaultModelVersion = "vlm"
	defaultLanguage     = "ch"
	pollInterval        = 3 * time.Second
	submitRatePerMin    = 45 // stay under documented 50/min limit
)

// Client calls MinerU precision parse API (token-auth batch upload).
type Client struct {
	Endpoint     string
	APIKey       string
	ModelVersion string
	Language     string
	HTTP         *http.Client

	submitMu   sync.Mutex
	submitLog  []time.Time
}

// Config wires Client from application settings.
type Config struct {
	Endpoint     string
	APIKey       string
	ModelVersion string
	Language     string
	HTTPTimeout  time.Duration
}

// NewClient builds a MinerU client. Empty APIKey means Disabled.
func NewClient(cfg Config) *Client {
	endpoint := strings.TrimRight(strings.TrimSpace(cfg.Endpoint), "/")
	if endpoint == "" {
		endpoint = defaultEndpoint
	}
	model := strings.TrimSpace(cfg.ModelVersion)
	if model == "" {
		model = defaultModelVersion
	}
	lang := strings.TrimSpace(cfg.Language)
	if lang == "" {
		lang = defaultLanguage
	}
	timeout := cfg.HTTPTimeout
	if timeout <= 0 {
		timeout = 120 * time.Second
	}
	return &Client{
		Endpoint:     endpoint,
		APIKey:       strings.TrimSpace(cfg.APIKey),
		ModelVersion: model,
		Language:     lang,
		HTTP:         &http.Client{Timeout: timeout},
	}
}

// Enabled reports whether MinerU parsing can run.
// Precision API needs a key; Agent lightweight API works without one.
func (c *Client) Enabled() bool {
	return c != nil
}

// PreferPrecision reports whether a precision API key is configured.
func (c *Client) PreferPrecision() bool {
	return c != nil && c.APIKey != ""
}

// ParseResult is the extracted markdown and optional content_list JSON.
type ParseResult struct {
	BatchID     string
	Markdown    string
	ContentList json.RawMessage
	TotalPages  int
	EngineMode  string // "precision" | "agent"
}

// ParseFile uploads local bytes and waits for markdown.
// With API key: try precision API first; on auth failure fall back to Agent lightweight API.
// Without API key: use Agent lightweight API directly.
func (c *Client) ParseFile(ctx context.Context, fileName string, data []byte) (*ParseResult, error) {
	if c == nil {
		return nil, fmt.Errorf("mineru: client nil")
	}
	if len(data) == 0 {
		return nil, fmt.Errorf("mineru: empty file")
	}
	name := filepath.Base(fileName)
	if name == "" || name == "." {
		name = "upload.bin"
	}

	if c.PreferPrecision() {
		if err := c.waitSubmitSlot(ctx); err != nil {
			return nil, err
		}
		res, err := c.parsePrecision(ctx, name, data)
		if err == nil {
			res.EngineMode = "precision"
			return res, nil
		}
		if !isPrecisionAuthFailure(err) {
			return nil, err
		}
		// Token expired / invalid → Agent fallback so uploads still parse.
	}

	res, err := c.parseAgentFile(ctx, name, data)
	if err != nil {
		return nil, err
	}
	if res.EngineMode == "" {
		res.EngineMode = "agent"
	}
	return res, nil
}

func (c *Client) parsePrecision(ctx context.Context, name string, data []byte) (*ParseResult, error) {
	batchID, uploadURL, err := c.applyUploadURL(ctx, name)
	if err != nil {
		return nil, err
	}
	if err := c.putFile(ctx, uploadURL, data); err != nil {
		return nil, err
	}
	zipURL, err := c.pollBatch(ctx, batchID, name)
	if err != nil {
		return nil, err
	}
	zipBytes, err := c.download(ctx, zipURL)
	if err != nil {
		return nil, err
	}
	md, contentList, pages, err := ExtractFromZip(zipBytes)
	if err != nil {
		return nil, err
	}
	return &ParseResult{
		BatchID:     batchID,
		Markdown:    md,
		ContentList: contentList,
		TotalPages:  pages,
	}, nil
}

func (c *Client) waitSubmitSlot(ctx context.Context) error {
	c.submitMu.Lock()
	defer c.submitMu.Unlock()
	for {
		now := time.Now()
		cutoff := now.Add(-time.Minute)
		kept := c.submitLog[:0]
		for _, t := range c.submitLog {
			if t.After(cutoff) {
				kept = append(kept, t)
			}
		}
		c.submitLog = kept
		if len(c.submitLog) < submitRatePerMin {
			c.submitLog = append(c.submitLog, now)
			return nil
		}
		wait := time.Minute - now.Sub(c.submitLog[0]) + 50*time.Millisecond
		c.submitMu.Unlock()
		timer := time.NewTimer(wait)
		select {
		case <-ctx.Done():
			timer.Stop()
			c.submitMu.Lock()
			return ctx.Err()
		case <-timer.C:
		}
		c.submitMu.Lock()
	}
}

type apiEnvelope struct {
	Code    any             `json:"code"`
	Msg     string          `json:"msg"`
	TraceID string          `json:"trace_id"`
	Data    json.RawMessage `json:"data"`
}

func (c *Client) applyUploadURL(ctx context.Context, fileName string) (batchID, uploadURL string, err error) {
	body := map[string]any{
		"files": []map[string]any{
			{"name": fileName},
		},
		"model_version": c.ModelVersion,
		"language":      c.Language,
		"enable_table":  true,
		"enable_formula": true,
	}
	raw, err := c.postJSON(ctx, c.Endpoint+"/api/v4/file-urls/batch", body)
	if err != nil {
		return "", "", err
	}
	var env apiEnvelope
	if err := json.Unmarshal(raw, &env); err != nil {
		return "", "", fmt.Errorf("mineru: decode apply url: %w", err)
	}
	if !codeOK(env.Code) {
		return "", "", classifyAPIError(env.Code, env.Msg)
	}
	var data struct {
		BatchID  string   `json:"batch_id"`
		FileURLs []string `json:"file_urls"`
	}
	if err := json.Unmarshal(env.Data, &data); err != nil {
		return "", "", fmt.Errorf("mineru: decode apply data: %w", err)
	}
	if data.BatchID == "" || len(data.FileURLs) == 0 {
		return "", "", fmt.Errorf("mineru: empty batch_id or file_urls")
	}
	return data.BatchID, data.FileURLs[0], nil
}

func (c *Client) putFile(ctx context.Context, uploadURL string, data []byte) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodPut, uploadURL, bytes.NewReader(data))
	if err != nil {
		return err
	}
	// MinerU docs: do not set Content-Type on signed PUT.
	resp, err := c.HTTP.Do(req)
	if err != nil {
		return fmt.Errorf("mineru: upload put: %w", err)
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return fmt.Errorf("mineru: upload put http %s", resp.Status)
	}
	return nil
}

func (c *Client) pollBatch(ctx context.Context, batchID, fileName string) (zipURL string, err error) {
	url := fmt.Sprintf("%s/api/v4/extract-results/batch/%s", c.Endpoint, batchID)
	ticker := time.NewTicker(pollInterval)
	defer ticker.Stop()
	for {
		raw, err := c.getJSON(ctx, url)
		if err != nil {
			return "", err
		}
		var env apiEnvelope
		if err := json.Unmarshal(raw, &env); err != nil {
			return "", fmt.Errorf("mineru: decode poll: %w", err)
		}
		if !codeOK(env.Code) {
			return "", classifyAPIError(env.Code, env.Msg)
		}
		var data struct {
			BatchID       string `json:"batch_id"`
			ExtractResult []struct {
				FileName    string `json:"file_name"`
				State       string `json:"state"`
				FullZipURL  string `json:"full_zip_url"`
				ErrMsg      string `json:"err_msg"`
			} `json:"extract_result"`
		}
		if err := json.Unmarshal(env.Data, &data); err != nil {
			return "", fmt.Errorf("mineru: decode poll data: %w", err)
		}
		item := pickResult(data.ExtractResult, fileName)
		if item == nil {
			// still waiting for file to appear
		} else {
			switch item.State {
			case "done":
				if item.FullZipURL == "" {
					return "", fmt.Errorf("mineru: done without zip url")
				}
				return item.FullZipURL, nil
			case "failed":
				return "", fmt.Errorf("mineru: parse failed: %s", strings.TrimSpace(item.ErrMsg))
			}
		}
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case <-ticker.C:
		}
	}
}

func pickResult(items []struct {
	FileName   string `json:"file_name"`
	State      string `json:"state"`
	FullZipURL string `json:"full_zip_url"`
	ErrMsg     string `json:"err_msg"`
}, fileName string) *struct {
	FileName   string `json:"file_name"`
	State      string `json:"state"`
	FullZipURL string `json:"full_zip_url"`
	ErrMsg     string `json:"err_msg"`
} {
	if len(items) == 0 {
		return nil
	}
	base := strings.ToLower(filepath.Base(fileName))
	for i := range items {
		if strings.ToLower(filepath.Base(items[i].FileName)) == base {
			return &items[i]
		}
	}
	return &items[0]
}

func (c *Client) download(ctx context.Context, url string) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.HTTP.Do(req)
	if err != nil {
		return nil, fmt.Errorf("mineru: download zip: %w", err)
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 64<<20))
	if err != nil {
		return nil, err
	}
	if resp.StatusCode >= 300 {
		return nil, fmt.Errorf("mineru: download zip http %s", resp.Status)
	}
	return raw, nil
}

func (c *Client) postJSON(ctx context.Context, url string, body any) ([]byte, error) {
	b, _ := json.Marshal(body)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(b))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+c.APIKey)
	resp, err := c.HTTP.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
	if resp.StatusCode == http.StatusTooManyRequests {
		return nil, &RetryableError{Msg: "mineru: rate limited (429)"}
	}
	if resp.StatusCode >= 300 {
		return nil, fmt.Errorf("mineru: http %s: %s", resp.Status, truncate(string(raw), 300))
	}
	return raw, nil
}

func (c *Client) getJSON(ctx context.Context, url string) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+c.APIKey)
	resp, err := c.HTTP.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
	if resp.StatusCode == http.StatusTooManyRequests {
		return nil, &RetryableError{Msg: "mineru: rate limited (429)"}
	}
	if resp.StatusCode >= 300 {
		return nil, fmt.Errorf("mineru: http %s: %s", resp.Status, truncate(string(raw), 300))
	}
	return raw, nil
}

func codeOK(code any) bool {
	switch v := code.(type) {
	case float64:
		return v == 0
	case int:
		return v == 0
	case string:
		return v == "0" || strings.EqualFold(v, "ok")
	case nil:
		return true
	default:
		return fmt.Sprint(v) == "0"
	}
}

// RetryableError marks transient MinerU failures (queue full / rate limit).
type RetryableError struct {
	Msg string
}

func (e *RetryableError) Error() string { return e.Msg }

// PermanentError marks non-retryable failures (auth / quota / bad file).
type PermanentError struct {
	Code string
	Msg  string
}

func (e *PermanentError) Error() string {
	if e.Code == "" {
		return e.Msg
	}
	return fmt.Sprintf("%s: %s", e.Code, e.Msg)
}

func classifyAPIError(code any, msg string) error {
	codeStr := strings.TrimSpace(fmt.Sprint(code))
	msg = strings.TrimSpace(msg)
	switch codeStr {
	case "A0202", "A0211":
		return &PermanentError{Code: codeStr, Msg: msg}
	case "-60018", "-60019":
		return &PermanentError{Code: codeStr, Msg: msg}
	case "-60009":
		return &RetryableError{Msg: fmt.Sprintf("%s: %s", codeStr, msg)}
	default:
		if strings.Contains(strings.ToLower(msg), "queue") {
			return &RetryableError{Msg: msg}
		}
		return &PermanentError{Code: codeStr, Msg: msg}
	}
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}

// SupportsExt reports whether MinerU precision API accepts the file type.
func SupportsExt(fileName string) bool {
	ext := strings.ToLower(filepath.Ext(fileName))
	switch ext {
	case ".pdf", ".doc", ".docx", ".ppt", ".pptx", ".xls", ".xlsx",
		".png", ".jpg", ".jpeg", ".jp2", ".webp", ".gif", ".bmp", ".html":
		return true
	default:
		return false
	}
}