package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strings"

	"github.com/alecthomas/kong"
	kongtoml "github.com/alecthomas/kong-toml"
)

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

type chatCompletionReq struct {
	Model           string        `json:"model"`
	Messages        []ChatMessage `json:"messages"`
	Temperature     *float64      `json:"temperature,omitempty"`
	ReasoningEffort string        `json:"reasoning_effort,omitempty"`
	Stream          bool          `json:"stream"`
}

type ResponseReasoning struct {
	Effort  string `json:"effort"`
	Summary string `json:"summary,omitempty"` // auto, concise, detailed
}
type chatResponseReq struct {
	Model       string            `json:"model"`
	Input       []ChatMessage     `json:"input"`
	Temperature *float64          `json:"temperature,omitempty"`
	Reasoning   ResponseReasoning `json:"reasoning,omitempty"`
	Stream      bool              `json:"stream"`
}

type responseStreamEvent struct {
	Type  string `json:"type"`
	Delta string `json:"delta"`
}

type streamEvent struct {
	Choices []struct {
		Delta struct {
			Content string `json:"content"`
		} `json:"delta"`
	} `json:"choices"`

	Error *struct {
		Message string `json:"message"`
		Type    string `json:"type"`
	} `json:"error,omitempty"`
}

func streamChatCompletions(
	ctx context.Context,
	client *http.Client,
	baseURL string,
	apiKey string,
	model string,
	reasoningEffort string,
	temperature *float64,
	msgs []ChatMessage,
) (io.ReadCloser, error) {
	if msgs == nil || len(msgs) == 0 {
		return nil, errors.New("missing messages")
	}
	endpoint := strings.TrimRight(baseURL, "/") + "/chat/completions"

	body := chatCompletionReq{
		Model:           model,
		Messages:        msgs,
		Temperature:     temperature,
		ReasoningEffort: reasoningEffort,
		Stream:          true,
	}

	payload, err := json.Marshal(body)
	if err != nil {
		return nil, err
	}

	pr, pw := io.Pipe()

	go func() {
		defer func(pw *io.PipeWriter) {
			err := pw.Close()
			if err != nil {
				log.Println(err)
			}
		}(pw)

		req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(payload))
		if err != nil {
			_ = pw.CloseWithError(err)
			return
		}

		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Accept", "text/event-stream")
		req.Header.Set("Authorization", "Bearer "+apiKey)

		resp, err := client.Do(req)
		if err != nil {
			_ = pw.CloseWithError(err)
			return
		}
		defer func(Body io.ReadCloser) {
			err := Body.Close()
			if err != nil {
				log.Println(err)
			}
		}(resp.Body)

		if resp.StatusCode < 200 || resp.StatusCode > 299 {
			b, _ := io.ReadAll(resp.Body)
			_ = pw.CloseWithError(fmt.Errorf("HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(b))))
			return
		}

		sc := bufio.NewScanner(resp.Body)
		sc.Buffer(make([]byte, 0, 64*1024), 2*1024*1024)

		var dataLines []string
		flushEvent := func() bool {
			if len(dataLines) == 0 {
				_ = pw.Close()
				return true
			}
			data := strings.Join(dataLines, "\n")
			dataLines = dataLines[:0]
			if data == "[DONE]" {
				_ = pw.Close()
				return false
			}

			var evt streamEvent
			if err := json.Unmarshal([]byte(data), &evt); err != nil {
				_ = pw.CloseWithError(fmt.Errorf("failed to unmarshal event: %w, data=%q", err, data))
				return false
			}
			if evt.Error != nil {
				_ = pw.CloseWithError(fmt.Errorf("api error: %+v", evt.Error))
				return false
			}
			if len(evt.Choices) > 0 {
				chunk := evt.Choices[0].Delta.Content
				if chunk != "" {
					if _, err := io.WriteString(pw, chunk); err != nil {
						_ = pw.CloseWithError(err)
						return false
					}
				}
			}
			return true
		}

		for sc.Scan() {
			line := sc.Text()

			if line == "" {
				if ok := flushEvent(); !ok {
					return
				}
				continue
			}

			if strings.HasPrefix(line, "data:") {
				v := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
				dataLines = append(dataLines, v)
			}
		}

		if len(dataLines) > 0 {
			_ = flushEvent()
			return
		}

		if err := sc.Err(); err != nil {
			_ = pw.CloseWithError(err)
			return
		}
	}()

	return pr, nil
}

func streamChatResponses(
	ctx context.Context,
	client *http.Client,
	baseURL string,
	apiKey string,
	model string,
	reasoningEffort string,
	temperature *float64,
	msgs []ChatMessage,
) (io.ReadCloser, error) {
	if msgs == nil || len(msgs) == 0 {
		return nil, errors.New("missing messages")
	}
	endpoint := strings.TrimRight(baseURL, "/") + "/responses"

	body := chatResponseReq{
		Model:       model,
		Input:       msgs,
		Temperature: temperature,
		Reasoning:   ResponseReasoning{Effort: reasoningEffort},
		Stream:      true,
	}

	payload, err := json.Marshal(body)
	if err != nil {
		return nil, err
	}

	pr, pw := io.Pipe()

	go func() {
		defer func(pw *io.PipeWriter) {
			err := pw.Close()
			if err != nil {
				log.Println(err)
			}
		}(pw)

		req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(payload))
		if err != nil {
			_ = pw.CloseWithError(err)
			return
		}

		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Accept", "text/event-stream")
		req.Header.Set("Authorization", "Bearer "+apiKey)

		resp, err := client.Do(req)
		if err != nil {
			_ = pw.CloseWithError(err)
			return
		}
		defer func(Body io.ReadCloser) {
			err := Body.Close()
			if err != nil {
				log.Println(err)
			}
		}(resp.Body)

		if resp.StatusCode < 200 || resp.StatusCode > 299 {
			b, _ := io.ReadAll(resp.Body)
			_ = pw.CloseWithError(fmt.Errorf("HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(b))))
			return
		}

		sc := bufio.NewScanner(resp.Body)
		sc.Buffer(make([]byte, 0, 64*1024), 2*1024*1024)

		var dataLines []string
		flushEvent := func() bool {
			if len(dataLines) == 0 {
				_ = pw.Close()
				return true
			}
			data := strings.Join(dataLines, "\n")
			dataLines = dataLines[:0]

			var evt responseStreamEvent
			if err := json.Unmarshal([]byte(data), &evt); err != nil {
				_ = pw.CloseWithError(fmt.Errorf("failed to unmarshal event: %w, data=%q", err, data))
				return false
			}
			if evt.Type == "response.output_text.delta" {
				if _, err := io.WriteString(pw, evt.Delta); err != nil {
					_ = pw.CloseWithError(err)
					return false
				}
			}
			return true
		}

		for sc.Scan() {
			line := sc.Text()

			if line == "" {
				if ok := flushEvent(); !ok {
					return
				}
				continue
			}

			if strings.HasPrefix(line, "data:") {
				v := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
				dataLines = append(dataLines, v)
			}
			if strings.HasPrefix(line, "event: response.completed") {
				break
			}
		}

		if len(dataLines) > 0 {
			_ = flushEvent()
			return
		}

		if err := sc.Err(); err != nil {
			_ = pw.CloseWithError(err)
			return
		}
	}()

	return pr, nil
}

// Translate src to lang
func Translate(
	ctx context.Context,
	client *http.Client,
	baseURL string,
	responseApi bool,
	apiKey string,
	model string,
	reasoningEffort string,
	temperature *float64,
	src string,
	lang string,
) (io.ReadCloser, error) {
	if client == nil {
		client = http.DefaultClient
	}
	msgs := []ChatMessage{
		{
			Role:    "system",
			Content: "You are a translation API. Please fully translate the user input to '" + lang + "'. Do not include any formatting changes/non-translation output. No explanation of the translation content. Do not translate proper nouns.",
		},
		{
			Role:    "user",
			Content: src,
		},
	}
	if responseApi {
		return streamChatResponses(ctx, client, baseURL, apiKey, model, reasoningEffort, temperature, msgs)
	}

	return streamChatCompletions(ctx, client, baseURL, apiKey, model, reasoningEffort, temperature, msgs)
}

var cli struct {
	Config          string   `short:"c" default:"~/.config/translate.toml" help:"Path to config file"`
	BaseURL         string   `short:"b" help:"LLM API base URL" default:"https://api.openai.com/v1"`
	ResponseApi     bool     `default:"false" help:"Use /v1/responses or /v1/chat/completions"`
	ApiKey          string   `short:"k" help:"LLM API Key"`
	Model           string   `short:"m" help:"LLM model" default:"gpt-5-nano"`
	ReasoningEffort string   `help:"LLM reasoning effort (note that some LLMs may not support certain settings)" default:"minimal"`
	Temperature     *float64 `help:"LLM Temperature"`
	Lang            string   `short:"l" help:"Target language" default:"zh-CN"`
	ReverseLang     string   `help:"Target language when use -r" default:"en-US"`
	Reverse         bool     `short:"r" help:"Set this flag to true to enable reverse translation"`
	Src             []string `arg:"" optional:"" name:"src" help:"Text to translate, or leave empty to use stdin"`
}

func main() {
	_ = kong.Parse(&cli)
	_ = kong.Parse(
		&cli,
		kong.Configuration(kongtoml.Loader, cli.Config),
	)

	lang := cli.Lang
	if cli.Reverse {
		lang = cli.ReverseLang
	}

	src := strings.Join(cli.Src, " ")

	if src == "" {
		stdin, err := io.ReadAll(os.Stdin)
		src = string(stdin)
		if err != nil {
			panic(err)
		}
	}

	stream, err := Translate(context.Background(), http.DefaultClient, cli.BaseURL, cli.ResponseApi, cli.ApiKey, cli.Model,
		cli.ReasoningEffort, cli.Temperature, src, lang)

	if err != nil {
		panic(err)
	}

	_, err = io.Copy(os.Stdout, stream)

	_, _ = os.Stdout.Write([]byte{'\n'})

	if err != nil {
		panic(err)
	}

	err = stream.Close()
	if err != nil {
		panic(err)
	}
}