package grok

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

	"github.com/fatih/color"
)

type Client struct {
	APIKey  string
	BaseURL string
}

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",
	}
}

func (c *Client) Stream(messages []map[string]string, model string) string {
	url := c.BaseURL + "/chat/completions"
	payload := map[string]interface{}{
		"model":       model,
		"messages":    messages,
		"temperature": 0.7,
		"stream":      true,
	}

	body, _ := json.Marshal(payload)
	req, _ := http.NewRequest("POST", url, bytes.NewReader(body))
	req.Header.Set("Authorization", "Bearer "+c.APIKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		color.Red("Request failed: %v", err)
		os.Exit(1)
	}
	defer func() {
		if cerr := resp.Body.Close(); cerr != nil {
			color.Yellow("Warning: failed to close response body: %v", cerr)
		}
	}()

	var fullReply strings.Builder
	scanner := bufio.NewScanner(resp.Body)
	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 != "" {
							fmt.Print(content)
							fullReply.WriteString(content)
						}
					}
				}
			}
		}
	}
	fmt.Println()
	return fullReply.String()
}

// CleanCodeResponse removes markdown fences and returns pure code content
func CleanCodeResponse(text string) string {
	text = strings.ReplaceAll(text, "```", "")
	text = strings.TrimSpace(text)
	return text
}