package llm

import (
	"bufio"
	"encoding/json"
	"net/http"
	"strings"

	config "git.estatecloud.org/radumaco/souvenir/config"
	"git.estatecloud.org/radumaco/souvenir/model"
)

type StreamEvent struct {
	Reasoning string
	Content   string
	Messages  []model.Message
	Done      bool
	Err       error
}

type ChatStreamResponse struct {
	Choices []StreamChoice `json:"choices"`
}

type StreamChoice struct {
	Delta        StreamDelta `json:"delta"`
	FinishReason *string     `json:"finish_reason"`
}

type StreamDelta struct {
	Role             string          `json:"role,omitempty"`
	Content          string          `json:"content,omitempty"`
	ReasoningContent string          `json:"reasoning_content,omitempty"`
	ToolCalls        []ToolCallDelta `json:"tool_calls,omitempty"`
}

type ToolCallDelta struct {
	Index    int    `json:"index"`
	Id       string `json:"id,omitempty"`
	Type     string `json:"type,omitempty"`
	Function struct {
		Name      string `json:"name,omitempty"`
		Arguments string `json:"arguments,omitempty"`
	} `json:"function"`
}

func (cl ChatClient) QueryStream(messages []model.Message) (<-chan StreamEvent, error) {
	requestBody := ChatRequest{
		Model:    cl.Cfg.Llm.Model,
		Messages: messages,
		Tools:    config.AvailableTools(),
		Stream:   true,
		Kwargs:   Kwargs{Thinking: false},
	}
	return cl.CallStream(requestBody)
}

func (cl ChatClient) CallStream(requestBody ChatRequest) (<-chan StreamEvent, error) {
	resp, err := cl.CallApi(requestBody)
	if err != nil {
		cl.logger.Error("There was an error during the API call", "error", err)
		return nil, err
	}
	events := make(chan StreamEvent)
	go cl.readStream(resp, events)
	return events, nil
}

func (cl ChatClient) readStream(resp *http.Response, events chan<- StreamEvent) {
	cl.logger.Debug("Started readStrea")
	defer resp.Body.Close()
	defer close(events)

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

	var content strings.Builder
	toolAcc := map[int]*model.ToolCall{}
	var order []int

	for scanner.Scan() {
		line := scanner.Text()
		cl.logger.Debug("Stream got new line", "line", line)

		if !strings.HasPrefix(line, "data:") {
			continue
		}

		payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if payload == "[DONE]" {
			cl.logger.Debug("Stream done")
			break
		}

		var chunk ChatStreamResponse
		if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
			cl.logger.Error("Error parsing stream chunk", "payload", payload, "error", err)
			events <- StreamEvent{Err: err}
			return
		}

		if len(chunk.Choices) == 0 {
			continue
		}

		delta := chunk.Choices[0].Delta
		if delta.ReasoningContent != "" {
			cl.logger.Debug("Stream got reasoning content", "content", delta.ReasoningContent)
			events <- StreamEvent{Reasoning: delta.ReasoningContent}
		}
		if delta.Content != "" {
			cl.logger.Debug("Stream got content", "content", delta.Content)
			content.WriteString(delta.Content)
			events <- StreamEvent{Content: delta.Content}
		}

		for _, tc := range delta.ToolCalls {
			acc, ok := toolAcc[tc.Index]
			if !ok {
				acc = &model.ToolCall{}
				toolAcc[tc.Index] = acc
				order = append(order, tc.Index)
			}
			if tc.Id != "" {
				acc.Id = tc.Id
			}
			if tc.Function.Name != "" {
				acc.FunctionName = tc.Function.Name
			}
			acc.FunctionArguments += tc.Function.Arguments
		}
	}

	if err := scanner.Err(); err != nil {
		cl.logger.Error("error reading stream", "error", err)
		events <- StreamEvent{Err: err}
		return
	}

	final := model.Message{Role: "assistant", Content: content.String()}
	for _, idx := range order {
		final.ToolCalls = append(final.ToolCalls, *toolAcc[idx])
	}

	events <- StreamEvent{Done: true, Messages: []model.Message{final}}
}