grokkit/internal/grok/client.go

package grok

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/fatih/color"
	"gmgauthier.com/grokkit/internal/logger"
)

// Client represents a client for interacting with the XAI API.
// It holds the API key and base URL for making requests.
type Client struct {
	APIKey  string
	BaseURL string
}

// NewClient creates and returns a new Client instance.
// It retrieves the API key from the XAI_API_KEY environment variable.
// If the variable is not set, it prints an error and exits the program.
func NewClient() *Client {
	key := os.Getenv("XAI_API_KEY")
	if key == "" {
		color.Red("Error: XAI_API_KEY environment variable not set")
		os.Exit(1)
	}
	return &Client{
		APIKey:  key,
		BaseURL: "https://api.x.ai/v1",
	}
}

// Stream sends a streaming chat completion request to the API and prints the response live to the terminal.
// It uses a default temperature of 0.7. This is intended for non-TUI commands.
// Returns the full response text.
func (c *Client) Stream(messages []map[string]string, model string) string {
	return c.StreamWithTemp(messages, model, 0.7)
}

// StreamWithTemp sends a streaming chat completion request to the API with a specified temperature
// and prints the response live to the terminal.
// Returns the full response text.
func (c *Client) StreamWithTemp(messages []map[string]string, model string, temperature float64) string {
	return c.streamInternal(messages, model, temperature, true)
}

// StreamSilent sends a streaming chat completion request to the API and returns the full response text
// without printing it live. It uses a default temperature of 0.7.
// This is intended for TUI and agent use.
func (c *Client) StreamSilent(messages []map[string]string, model string) string {
	return c.streamInternal(messages, model, 0.7, false)
}

func (c *Client) streamInternal(messages []map[string]string, model string, temperature float64, printLive bool) string {
	startTime := time.Now()
	url := c.BaseURL + "/chat/completions"

	logger.Debug("preparing API request",
		"model", model,
		"temperature", temperature,
		"message_count", len(messages),
		"stream", true)

	payload := map[string]interface{}{
		"model":       model,
		"messages":    messages,
		"temperature": temperature,
		"stream":      true,
	}

	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	// Manual cancel before os.Exit; otherwise defer is fine for the main path.
	defer cancel()

	body, err := json.Marshal(payload)
	if err != nil {
		logger.Error("failed to marshal API request", "error", err)
		color.Red("Failed to marshal request: %v", err)
		cancel()
		os.Exit(1)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		logger.Error("failed to create HTTP request", "error", err, "url", url)
		color.Red("Failed to create request: %v", err)
		cancel()
		os.Exit(1)
	}
	req.Header.Set("Authorization", "Bearer "+c.APIKey)
	req.Header.Set("Content-Type", "application/json")

	logger.Info("sending API request", "model", model, "url", url)

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		logger.Error("API request failed",
			"error", err,
			"model", model,
			"duration_ms", time.Since(startTime).Milliseconds())
		color.Red("Request failed: %v", err)
		cancel()
		os.Exit(1)
	}
	defer func() { _ = resp.Body.Close() }()

	logger.Debug("API response received",
		"status", resp.Status,
		"status_code", resp.StatusCode,
		"duration_ms", time.Since(startTime).Milliseconds())

	var fullReply strings.Builder
	scanner := bufio.NewScanner(resp.Body)
	chunkCount := 0
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "data: ") {
			data := strings.TrimPrefix(line, "data: ")
			if data == "[DONE]" {
				break
			}
			var chunk map[string]any
			if json.Unmarshal([]byte(data), &chunk) == nil {
				if choices, ok := chunk["choices"].([]any); ok && len(choices) > 0 {
					if delta, ok := choices[0].(map[string]any)["delta"].(map[string]any); ok {
						if content, ok := delta["content"].(string); ok && content != "" {
							chunkCount++
							fullReply.WriteString(content)
							if printLive {
								fmt.Print(content)
							}
						}
					}
				}
			}
		}
	}
	if printLive {
		fmt.Println()
	}

	responseLength := fullReply.Len()
	duration := time.Since(startTime)

	logger.Info("API request completed",
		"model", model,
		"response_length", responseLength,
		"chunks_received", chunkCount,
		"duration_ms", duration.Milliseconds(),
		"duration", duration.String())

	return fullReply.String()
}

// CleanCodeResponse removes Markdown code fences from the input text and returns the pure code content.
// It removes opening fences (with optional language tags) and closing fences, then trims leading and trailing whitespace.
// Internal blank lines are preserved as they may be intentional in code.
func CleanCodeResponse(text string) string {
	fence := "```"

	// Remove any line that starts with the fence (opening fence, possibly with language tag)
	text = regexp.MustCompile(`(?m)^`+regexp.QuoteMeta(fence)+`.*$`).ReplaceAllString(text, "")

	// Remove any line that is just the fence (closing fence)
	text = regexp.MustCompile(`(?m)^`+regexp.QuoteMeta(fence)+`\s*$`).ReplaceAllString(text, "")

	// Trim only leading and trailing whitespace.
	// Do NOT collapse internal blank lines — they are intentional in code.
	text = strings.TrimSpace(text)

	return text
}