package ai

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"

	"write-server/internal/config"
)

var AI_PROMPTS = map[string]string{
	"polish":  "\u4f60\u662f\u4e00\u4e2a\u7f16\u8f91\u3002\u6da6\u8272\u4ee5\u4e0b\u6587\u672c\uff0c\u4f7f\u5176\u66f4\u6d41\u7545\u3001\u66f4\u4e13\u4e1a\u3002\u4fdd\u6301\u539f\u610f\u548c\u8bed\u6c14\uff0c\u53ea\u8fd4\u56de\u6da6\u8272\u540e\u7684\u6587\u672c\uff0c\u4e0d\u8981\u89e3\u91ca\u3002",
	"expand":  "\u4f60\u662f\u4e00\u4e2a\u7f16\u8f91\u3002\u6269\u5c55\u4ee5\u4e0b\u6587\u672c\uff0c\u6dfb\u52a0\u66f4\u591a\u7ec6\u8282\u548c\u89e3\u91ca\uff0c\u540c\u65f6\u4fdd\u6301\u6e05\u6670\u548c\u5f15\u4eba\u5165\u80dc\u3002\u53ea\u8fd4\u56de\u6269\u5c55\u540e\u7684\u6587\u672c\uff0c\u4e0d\u8981\u89e3\u91ca\u3002",
	"shorten": "\u4f60\u662f\u4e00\u4e2a\u7f16\u8f91\u3002\u7f29\u77ed\u4ee5\u4e0b\u6587\u672c\uff0c\u4fdd\u6301\u6838\u5fc3\u610f\u601d\uff0c\u4f7f\u8868\u8fbe\u66f4\u7b80\u6d01\u3002\u53ea\u8fd4\u56de\u7f29\u77ed\u540e\u7684\u6587\u672c\uff0c\u4e0d\u8981\u89e3\u91ca\u3002",
	"fix":     "\u4f60\u662f\u4e00\u4e2a\u7f16\u8f91\u3002\u4fee\u590d\u4ee5\u4e0b\u6587\u672c\u4e2d\u7684\u8bed\u6cd5\u9519\u8bef\u3001\u62fc\u5199\u9519\u8bef\u548c\u4e0d\u81ea\u7136\u8868\u8fbe\uff0c\u4fdd\u6301\u539f\u610f\u3002\u53ea\u8fd4\u56de\u4fee\u6b63\u540e\u7684\u6587\u672c\uff0c\u4e0d\u8981\u89e3\u91ca\u3002",
}

type ChatMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type ChatRequest struct {
	Model    string        `json:"model"`
	Messages []ChatMessage `json:"messages"`
}

type ChatResponse struct {
	Choices []struct {
		Message ChatMessage `json:"message"`
	} `json:"choices"`
}

func AIChat(systemPrompt, userMessage string) (string, error) {
	if config.DEEPSEEK_API_KEY == "" {
		return "", fmt.Errorf("DEEPSEEK_API_KEY not configured")
	}

	req := ChatRequest{
		Model: config.DEEPSEEK_MODEL,
		Messages: []ChatMessage{
			{Role: "system", Content: systemPrompt},
			{Role: "user", Content: userMessage},
		},
	}

	data, _ := json.Marshal(req)
	httpReq, _ := http.NewRequest("POST", config.DEEPSEEK_BASE_URL+"/chat/completions", bytes.NewReader(data))
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Authorization", "Bearer "+config.DEEPSEEK_API_KEY)

	resp, err := http.DefaultClient.Do(httpReq)
	if err != nil {
		return "", fmt.Errorf("AI request failed: %v", err)
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("AI error: %s", string(body))
	}

	var chatResp ChatResponse
	if err := json.Unmarshal(body, &chatResp); err != nil {
		return "", err
	}
	if len(chatResp.Choices) == 0 {
		return "", fmt.Errorf("no response from AI")
	}
	return chatResp.Choices[0].Message.Content, nil
}