package main

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

	tea "charm.land/bubbletea/v2"
)

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

type chatRequest struct {
	Model    string        `json:"model"`
	Messages []chatMessage `json:"messages"`
	Stream   bool          `json:"stream"`
}

type sseChunk struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	Model   string `json:"model"`
	Choices []struct {
		Index        int         `json:"index"`
		Delta        chatMessage `json:"delta"`
		FinishReason *string     `json:"finish_reason"`
	} `json:"choices"`
}

type tagModel struct {
	Name string `json:"name"`
}

type listTagsResp struct {
	Models []tagModel `json:"models"`
}

type streamChunkMsg string
type streamDoneMsg struct{}
type streamErrMsg struct{ err error }

func fetchModels(baseURL string) ([]string, error) {
	resp, err := http.Get(strings.TrimRight(baseURL, "/") + "/api/tags")
	if err != nil {
		return nil, fmt.Errorf("failed to fetch models: %w", err)
	}
	defer resp.Body.Close()

	var r listTagsResp
	if err := json.NewDecoder(resp.Body).Decode(&r); err != nil {
		return nil, fmt.Errorf("failed to decode models: %w", err)
	}

	models := make([]string, len(r.Models))
	for i, m := range r.Models {
		models[i] = m.Name
	}
	return models, nil
}

func startStream(ctx context.Context, msgs []chatMessage, model, baseURL string) (<-chan tea.Msg, tea.Cmd) {
	ch := make(chan tea.Msg, 16)
	go runStream(ctx, ch, msgs, model, baseURL)
	return ch, readNext(ch)
}

func readNext(ch <-chan tea.Msg) tea.Cmd {
	return func() tea.Msg {
		msg, ok := <-ch
		if !ok {
			return streamDoneMsg{}
		}
		return msg
	}
}

func runStream(ctx context.Context, ch chan<- tea.Msg, msgs []chatMessage, model, baseURL string) {
	defer close(ch)

	reqBody := chatRequest{
		Model:    model,
		Messages: msgs,
		Stream:   true,
	}

	body, err := json.Marshal(reqBody)
	if err != nil {
		ch <- streamErrMsg{err: fmt.Errorf("marshal request: %w", err)}
		return
	}

	url := strings.TrimRight(baseURL, "/") + "/v1/chat/completions"
	req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
	if err != nil {
		ch <- streamErrMsg{err: fmt.Errorf("create request: %w", err)}
		return
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "text/event-stream")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		if ctx.Err() != nil {
			return
		}
		ch <- streamErrMsg{err: fmt.Errorf("http request: %w", err)}
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		errBody, _ := io.ReadAll(resp.Body)
		ch <- streamErrMsg{err: fmt.Errorf("api %d: %s", resp.StatusCode, string(errBody))}
		return
	}

	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		if ctx.Err() != nil {
			return
		}

		line := strings.TrimSpace(scanner.Text())

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

		data := strings.TrimSpace(line[5:])
		if data == "[DONE]" {
			break
		}

		var chunk sseChunk
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			continue
		}

		if len(chunk.Choices) > 0 && chunk.Choices[0].Delta.Content != "" {
			ch <- streamChunkMsg(chunk.Choices[0].Delta.Content)
		}
	}

	if err := scanner.Err(); err != nil {
		if ctx.Err() != nil {
			return
		}
		ch <- streamErrMsg{err: fmt.Errorf("scanner: %w", err)}
	}
}