claude-statusline/internal/fetcher/fetcher.go

package fetcher

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"git.davoryn.de/calic/claude-statusline/internal/browser"
	"git.davoryn.de/calic/claude-statusline/internal/config"
)

const (
	apiBase   = "https://claude.ai"
	userAgent = "Mozilla/5.0 (X11; Linux x86_64; rv:135.0) Gecko/20100101 Firefox/135.0"
)

// ParsedUsage is the display-friendly usage data passed to callbacks.
type ParsedUsage struct {
	FiveHourPct      int
	FiveHourResetsAt string
	FiveHourResetsIn string
	SevenDayPct      int
	SevenDayResetsAt string
	SevenDayResetsIn string
	Error            string
}

// UpdateCallback is called when new usage data is available.
type UpdateCallback func(ParsedUsage)

// doRequest performs an authenticated HTTP GET to the Claude API.
// Includes any saved Cloudflare cookies from previous Chrome fallbacks.
func doRequest(url, sessionKey string) ([]byte, int, error) {
	client := &http.Client{Timeout: 10 * time.Second}
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, 0, err
	}

	cookie := "sessionKey=" + sessionKey
	// Append Cloudflare cookies if available (saved by Chrome fallback)
	if cfCookies := loadCFCookies(); cfCookies != "" {
		cookie += "; " + cfCookies
	}

	req.Header.Set("Cookie", cookie)
	req.Header.Set("User-Agent", userAgent)
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Referer", "https://claude.ai/")
	req.Header.Set("Origin", "https://claude.ai")

	resp, err := client.Do(req)
	if err != nil {
		return nil, 0, err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, resp.StatusCode, err
	}
	return body, resp.StatusCode, nil
}

// loadCFCookies reads saved Cloudflare cookies from the cf-cookies file.
func loadCFCookies() string {
	data, err := os.ReadFile(filepath.Join(config.ConfigDir(), "cf-cookies"))
	if err != nil {
		return ""
	}
	// File has one cookie per line (name=value), join with "; "
	lines := strings.Split(strings.TrimSpace(string(data)), "\n")
	var valid []string
	for _, l := range lines {
		l = strings.TrimSpace(l)
		if l != "" {
			valid = append(valid, l)
		}
	}
	return strings.Join(valid, "; ")
}

// fetchWithFallback tries a plain HTTP request first, then falls back to
// headless Chrome (which can solve Cloudflare JS challenges) on 403.
func fetchWithFallback(url, sessionKey string) ([]byte, error) {
	body, status, err := doRequest(url, sessionKey)
	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}
	if status == 200 {
		return body, nil
	}
	if status == 401 {
		return nil, fmt.Errorf("auth_expired")
	}
	if status == 403 {
		// Likely a Cloudflare JS challenge — fall back to headless Chrome
		log.Printf("HTTP 403 for %s, falling back to Chrome", url)
		chromeBody, chromeErr := browser.FetchViaChrome(url)
		if chromeErr != nil {
			log.Printf("Chrome fallback failed: %v", chromeErr)
			return nil, fmt.Errorf("auth_expired")
		}
		log.Printf("Chrome fallback succeeded (%d bytes)", len(chromeBody))
		return chromeBody, nil
	}
	return nil, fmt.Errorf("HTTP %d", status)
}

// DiscoverOrgID fetches the first organization UUID from the API.
func DiscoverOrgID(sessionKey string) (string, error) {
	body, err := fetchWithFallback(apiBase+"/api/organizations", sessionKey)
	if err != nil {
		return "", err
	}

	var orgs []struct {
		UUID string `json:"uuid"`
	}
	if err := json.Unmarshal(body, &orgs); err != nil {
		return "", fmt.Errorf("invalid response: %w", err)
	}
	if len(orgs) == 0 {
		return "", fmt.Errorf("no organizations found")
	}
	return orgs[0].UUID, nil
}

// FetchUsage fetches usage data from the API. If orgID is empty, discovers it.
// Returns the raw cache data and the resolved org ID.
func FetchUsage(sessionKey, orgID string) (*CacheData, string, error) {
	if orgID == "" {
		var err error
		orgID, err = DiscoverOrgID(sessionKey)
		if err != nil {
			if err.Error() == "auth_expired" {
				return &CacheData{Error: "auth_expired", Status: 401}, "", err
			}
			return nil, "", err
		}
	}

	url := fmt.Sprintf("%s/api/organizations/%s/usage", apiBase, orgID)
	body, err := fetchWithFallback(url, sessionKey)
	if err != nil {
		if err.Error() == "auth_expired" {
			return &CacheData{Error: "auth_expired", Status: 403}, orgID, err
		}
		return &CacheData{Error: "fetch_failed", Message: err.Error()}, orgID, err
	}

	var data CacheData
	if err := json.Unmarshal(body, &data); err != nil {
		return nil, orgID, fmt.Errorf("invalid JSON: %w", err)
	}
	return &data, orgID, nil
}

// ParseUsage converts raw cache data into display-friendly format.
func ParseUsage(data *CacheData) ParsedUsage {
	if data == nil {
		return ParsedUsage{Error: "no data"}
	}
	if data.Error != "" {
		msg := data.Error
		if msg == "auth_expired" {
			msg = "session expired"
		}
		return ParsedUsage{Error: msg}
	}

	p := ParsedUsage{}
	if data.FiveHour != nil {
		p.FiveHourPct = int(math.Round(data.FiveHour.Utilization))
		p.FiveHourResetsAt = data.FiveHour.ResetsAt
		p.FiveHourResetsIn = formatResetsIn(data.FiveHour.ResetsAt)
	}
	if data.SevenDay != nil {
		p.SevenDayPct = int(math.Round(data.SevenDay.Utilization))
		p.SevenDayResetsAt = data.SevenDay.ResetsAt
		p.SevenDayResetsIn = formatResetsIn(data.SevenDay.ResetsAt)
	}
	return p
}

// formatResetsIn converts an ISO 8601 timestamp to a human-readable duration.
func formatResetsIn(isoStr string) string {
	if isoStr == "" {
		return ""
	}
	t, err := time.Parse(time.RFC3339, isoStr)
	if err != nil {
		return ""
	}
	total := int(math.Max(0, time.Until(t).Seconds()))
	days := total / 86400
	hours := (total % 86400) / 3600
	minutes := (total % 3600) / 60

	if days > 0 {
		return fmt.Sprintf("%dd %dh", days, hours)
	}
	if hours > 0 {
		return fmt.Sprintf("%dh %dm", hours, minutes)
	}
	return fmt.Sprintf("%dm", minutes)
}

// BackgroundFetcher runs periodic usage fetches in a goroutine.
type BackgroundFetcher struct {
	onUpdate UpdateCallback
	mu       sync.Mutex
	interval time.Duration
	orgID    string
	stopCh   chan struct{}
	forceCh  chan struct{}
}

// NewBackgroundFetcher creates a new background fetcher.
func NewBackgroundFetcher(onUpdate UpdateCallback) *BackgroundFetcher {
	cfg := config.Load()
	return &BackgroundFetcher{
		onUpdate: onUpdate,
		interval: time.Duration(cfg.RefreshInterval) * time.Second,
		orgID:    cfg.OrgID,
		stopCh:   make(chan struct{}),
		forceCh:  make(chan struct{}, 1),
	}
}

// Start begins the background fetch loop.
func (bf *BackgroundFetcher) Start() {
	go bf.loop()
}

// Stop signals the background fetcher to stop.
func (bf *BackgroundFetcher) Stop() {
	close(bf.stopCh)
}

// Refresh forces an immediate fetch.
func (bf *BackgroundFetcher) Refresh() {
	select {
	case bf.forceCh <- struct{}{}:
	default:
	}
}

// SetInterval changes the refresh interval.
func (bf *BackgroundFetcher) SetInterval(seconds int) {
	bf.mu.Lock()
	bf.interval = time.Duration(seconds) * time.Second
	bf.mu.Unlock()

	cfg := config.Load()
	cfg.RefreshInterval = seconds
	_ = config.Save(cfg)

	bf.Refresh()
}

// Interval returns the current refresh interval in seconds.
func (bf *BackgroundFetcher) Interval() int {
	bf.mu.Lock()
	defer bf.mu.Unlock()
	return int(bf.interval.Seconds())
}

func (bf *BackgroundFetcher) loop() {
	// Load from cache immediately for instant display
	if data, _ := ReadCache(); data != nil {
		bf.onUpdate(ParseUsage(data))
	}

	// Initial fetch
	bf.doFetch(false)

	for {
		bf.mu.Lock()
		interval := bf.interval
		bf.mu.Unlock()

		timer := time.NewTimer(interval)
		select {
		case <-bf.stopCh:
			timer.Stop()
			return
		case <-bf.forceCh:
			timer.Stop()
			bf.doFetch(true)
		case <-timer.C:
			bf.doFetch(false)
		}
	}
}

func (bf *BackgroundFetcher) doFetch(force bool) {
	bf.mu.Lock()
	halfInterval := bf.interval / 2
	bf.mu.Unlock()

	if !force {
		if cached := ReadCacheIfFresh(halfInterval); cached != nil && cached.Error == "" {
			bf.onUpdate(ParseUsage(cached))
			return
		}
	}

	sessionKey := config.GetSessionKey()
	if sessionKey == "" {
		errData := &CacheData{Error: "no_session_key"}
		_ = WriteCache(errData)
		bf.onUpdate(ParseUsage(errData))
		return
	}

	data, orgID, _ := FetchUsage(sessionKey, bf.orgID)
	if data == nil {
		return
	}

	_ = WriteCache(data)

	if orgID != "" && orgID != bf.orgID {
		bf.mu.Lock()
		bf.orgID = orgID
		bf.mu.Unlock()
		cfg := config.Load()
		cfg.OrgID = orgID
		_ = config.Save(cfg)
	}

	bf.onUpdate(ParseUsage(data))
}