package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"

	"fyne.io/fyne/v2"
	"fyne.io/fyne/v2/app"
	"fyne.io/fyne/v2/container"
	"fyne.io/fyne/v2/data/validation"
	"fyne.io/fyne/v2/dialog"
	"fyne.io/fyne/v2/theme"
	"fyne.io/fyne/v2/widget"
)

// ---------- Структуры для работы с Ollama ----------
type Message struct {
	Role      string `json:"role"`
	Content   string `json:"content"`
	Model     string `json:"model,omitempty"`
	Timestamp string `json:"timestamp"`
	IsBot     bool   `json:"is_bot"`
	Round     int    `json:"round,omitempty"`
}

type ChatRequest struct {
	Model    string    `json:"model"`
	Messages []Message `json:"messages"`
	Stream   bool      `json:"stream"`
	Options  struct {
		NumPredict int `json:"num_predict,omitempty"`
		Seed       int `json:"seed,omitempty"`
	} `json:"options,omitempty"`
}

type ChatResponse struct {
	Model     string  `json:"model"`
	Message   Message `json:"message"`
	Done      bool    `json:"done"`
	CreatedAt string  `json:"created_at"`
}

// ModelContext хранит системное сообщение и историю для каждой модели
type ModelContext struct {
	Role          string
	SystemMessage Message
	DialogHistory []Message
}

// ConversationState управляет состоянием диалога
type ConversationState struct {
	Model1        string
	Model2        string
	MaxRounds     int
	TimeoutMin    int
	InitialPrompt string
	History       []Message
	CurrentRound  int
	IsActive      bool
	Model1Context *ModelContext
	Model2Context *ModelContext
}

// guiApp объединяет все элементы GUI и состояние
type guiApp struct {
	window       fyne.Window
	status       *widget.Label
	model1Select *widget.Select
	model2Select *widget.Select
	role1Entry   *widget.Entry
	role2Entry   *widget.Entry
	topicEntry   *widget.Entry
	roundsEntry  *widget.Entry
	timeoutEntry *widget.Entry
	startBtn     *widget.Button
	resetBtn     *widget.Button
	chatRichText *widget.RichText
	roundLabel   *widget.Label

	state ConversationState
	mu    sync.Mutex // защита state и chatRichText.Segments
	wg    sync.WaitGroup
}

// callOllamaAPI отправляет запрос к локальному Ollama и возвращает ответ
func callOllamaAPI(ctx context.Context, model string, messages []Message) (string, error) {
	cleanMessages := []Message{}
	for _, msg := range messages {
		if msg.Content != "" && strings.TrimSpace(msg.Content) != "" {
			cleanMessages = append(cleanMessages, msg)
		}
	}
	if len(cleanMessages) == 0 {
		return "", fmt.Errorf("нет валидных сообщений для отправки в модель")
	}

	requestBody := ChatRequest{
		Model:    model,
		Messages: cleanMessages,
		Stream:   false,
	}
	requestBody.Options.NumPredict = 1024

	jsonData, err := json.Marshal(requestBody)
	if err != nil {
		return "", fmt.Errorf("ошибка маршалинга запроса: %v", err)
	}

	req, err := http.NewRequestWithContext(ctx, "POST",
		"http://localhost:11434/api/chat",
		bytes.NewBuffer(jsonData))
	if err != nil {
		return "", fmt.Errorf("ошибка создания запроса: %v", err)
	}
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{
		Timeout: 0,
	}
	resp, err := client.Do(req)
	if err != nil {
		return "", fmt.Errorf("ошибка HTTP запроса: %v", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("ошибка чтения ответа: %v", err)
	}

	var response ChatResponse
	if err := json.Unmarshal(body, &response); err != nil {
		return "", fmt.Errorf("ошибка парсинга JSON: %v", err)
	}
	if strings.TrimSpace(response.Message.Content) == "" {
		return "", fmt.Errorf("модель вернула пустой ответ")
	}
	return response.Message.Content, nil
}

// updateDialogContext добавляет сообщение в историю контекста модели
func updateDialogContext(context *ModelContext, message Message, isOwnMessage bool) {
	if strings.TrimSpace(message.Content) == "" {
		return
	}
	msgToAdd := message
	if isOwnMessage {
		msgToAdd.Role = "assistant"
	} else {
		msgToAdd.Role = "user"
	}
	context.DialogHistory = append(context.DialogHistory, msgToAdd)
	if len(context.DialogHistory) > 20 {
		context.DialogHistory = context.DialogHistory[len(context.DialogHistory)-20:]
	}
}

// buildContext формирует слайс сообщений для отправки в модель
func buildContext(context *ModelContext) []Message {
	messages := make([]Message, 0)
	if context.SystemMessage.Content != "" {
		messages = append(messages, context.SystemMessage)
	}
	messages = append(messages, context.DialogHistory...)
	return messages
}

// appendMessage добавляет сообщение в RichText
func (g *guiApp) appendMessage(msg Message) {
	g.mu.Lock()
	defer g.mu.Unlock()

	timestamp := msg.Timestamp
	if timestamp == "" {
		timestamp = time.Now().Format("15:04:05")
	}
	sender := ""
	if msg.IsBot {
		sender = msg.Model
	} else if msg.Role == "error" {
		sender = "Ошибка"
	} else {
		sender = "Система"
	}
	header := fmt.Sprintf("[%s] %s", timestamp, sender)

	var colorName fyne.ThemeColorName
	if msg.IsBot {
		if msg.Model == g.model1Select.Selected {
			colorName = theme.ColorNameWarning
		} else if msg.Model == g.model2Select.Selected {
			colorName = theme.ColorNameSuccess
		} else {
			colorName = theme.ColorNameForeground
		}
	} else if msg.Role == "error" {
		colorName = theme.ColorNameError
	} else {
		colorName = theme.ColorNameForeground
	}

	headerSegment := &widget.TextSegment{
		Text: header + "\n",
		Style: widget.RichTextStyle{
			TextStyle: fyne.TextStyle{Bold: true},
			ColorName: colorName,
		},
	}
	msgSegment := &widget.TextSegment{
		Text: msg.Content + "\n\n",
		Style: widget.RichTextStyle{
			ColorName: colorName,
		},
	}

	g.chatRichText.Segments = append(g.chatRichText.Segments, headerSegment, msgSegment)
	g.chatRichText.Refresh()
}

// loadModels получает список моделей через ollama list и обновляет селекты
func (g *guiApp) loadModels() {
	cmd := exec.Command("ollama", "list")
	output, err := cmd.Output()
	if err != nil {
		log.Printf("Ошибка выполнения ollama list: %v", err)
		fyne.Do(func() {
			dialog.ShowInformation("Предупреждение",
				"Не удалось получить список моделей от Ollama.\n"+
					"Убедитесь, что Ollama установлен и запущен.\n", g.window)
		})
		defaultModels := []string{
			"Модели Ollama не найдены.",
		}
		fyne.Do(func() {
			g.model1Select.Options = defaultModels
			g.model2Select.Options = defaultModels
			g.model1Select.Refresh()
			g.model2Select.Refresh()
			g.status.SetText("Ollama не запущен")
		})
		return
	}

	lines := strings.Split(string(output), "\n")
	models := []string{}
	for i, line := range lines {
		if i == 0 || strings.TrimSpace(line) == "" {
			continue
		}
		columns := strings.Fields(line)
		if len(columns) > 0 {
			modelName := columns[0]
			if modelName != "" {
				models = append(models, modelName)
			}
		}
	}
	if len(models) == 0 {
		log.Println("Модели Ollama не найдены.")
		models = []string{
			"Модели Ollama не найдены.",
		}
		fyne.Do(func() {
			g.model1Select.Options = models
			g.model2Select.Options = models
			g.model1Select.Refresh()
			g.model2Select.Refresh()
			g.status.SetText("Нет моделей Ollama")
		})
	} else {
		fyne.Do(func() {
			g.model1Select.Options = models
			g.model2Select.Options = models
			g.model1Select.Refresh()
			g.model2Select.Refresh()
			g.status.SetText("Готов")
		})
	}
	log.Printf("Загружено %d моделей", len(models))
}

// runConversation выполняет диалог в отдельной горутине
func (g *guiApp) runConversation(params struct {
	Model1, Model2, Role1, Role2, InitialPrompt string
	MaxRounds, TimeoutMin                       int
}) {
	defer g.wg.Done()

	g.mu.Lock()
	if g.state.IsActive {
		g.mu.Unlock()
		fyne.Do(func() {
			g.appendMessage(Message{
				Role:      "system",
				Content:   "Диалог уже запущен, дождитесь завершения",
				Timestamp: time.Now().Format("15:04:05"),
				IsBot:     false,
			})
		})
		return
	}
	g.state = ConversationState{
		Model1:        params.Model1,
		Model2:        params.Model2,
		MaxRounds:     params.MaxRounds,
		TimeoutMin:    params.TimeoutMin,
		InitialPrompt: params.InitialPrompt,
		History:       make([]Message, 0),
		CurrentRound:  0,
		IsActive:      true,
		Model1Context: &ModelContext{
			Role:          params.Role1,
			DialogHistory: make([]Message, 0),
		},
		Model2Context: &ModelContext{
			Role:          params.Role2,
			DialogHistory: make([]Message, 0),
		},
	}
	systemMsg1 := params.InitialPrompt
	if params.Role1 != "" {
		systemMsg1 = fmt.Sprintf("Ты играешь роль: %s. %s", params.Role1, params.InitialPrompt)
	}
	systemMsg2 := params.InitialPrompt
	if params.Role2 != "" {
		systemMsg2 = fmt.Sprintf("Ты играешь роль: %s. %s", params.Role2, params.InitialPrompt)
	}
	g.state.Model1Context.SystemMessage = Message{
		Role:    "system",
		Content: systemMsg1,
		Model:   params.Model1,
	}
	g.state.Model2Context.SystemMessage = Message{
		Role:    "system",
		Content: systemMsg2,
		Model:   params.Model2,
	}
	g.mu.Unlock()

	fyne.Do(func() {
		g.appendMessage(Message{
			Role:      "system",
			Content:   fmt.Sprintf("Диалог начался: %s vs %s", params.Model1, params.Model2),
			Timestamp: time.Now().Format("15:04:05"),
			IsBot:     false,
		})
		g.appendMessage(Message{
			Role:      "user",
			Content:   "Давайте начнем наш диалог. " + params.InitialPrompt,
			Timestamp: time.Now().Format("15:04:05"),
			IsBot:     false,
		})
	})

	for round := 1; round <= params.MaxRounds; round++ {
		g.mu.Lock()
		if !g.state.IsActive {
			g.mu.Unlock()
			break
		}
		g.state.CurrentRound = round
		g.mu.Unlock()

		fyne.Do(func() {
			g.mu.Lock()
			roundNum := g.state.CurrentRound
			maxRounds := g.state.MaxRounds
			g.mu.Unlock()
			g.roundLabel.SetText(fmt.Sprintf("Раунд: %d/%d", roundNum, maxRounds))
		})

		log.Printf("--- Раунд %d ---", round)

		// Ход первой модели
		g.mu.Lock()
		messages1 := buildContext(g.state.Model1Context)
		g.mu.Unlock()

		ctx1, cancel1 := context.WithTimeout(context.Background(), time.Duration(params.TimeoutMin)*time.Minute)
		response1, err := callOllamaAPI(ctx1, params.Model1, messages1)
		cancel1()
		if err != nil {
			fyne.Do(func() {
				g.appendMessage(Message{
					Role:      "error",
					Content:   fmt.Sprintf("Ошибка от %s: %v", params.Model1, err),
					Timestamp: time.Now().Format("15:04:05"),
					IsBot:     false,
				})
			})
			break
		}
		msg1 := Message{
			Role:      "assistant",
			Content:   response1,
			Model:     params.Model1,
			Timestamp: time.Now().Format("15:04:05"),
			IsBot:     true,
			Round:     round,
		}

		g.mu.Lock()
		g.state.History = append(g.state.History, msg1)
		updateDialogContext(g.state.Model1Context, msg1, true)
		updateDialogContext(g.state.Model2Context, Message{
			Role:    "user",
			Content: response1,
			Model:   params.Model1,
		}, false)
		g.mu.Unlock()

		fyne.Do(func() {
			g.appendMessage(msg1)
		})

		g.mu.Lock()
		active := g.state.IsActive
		g.mu.Unlock()
		if !active {
			break
		}

		// Ход второй модели
		g.mu.Lock()
		messages2 := buildContext(g.state.Model2Context)
		g.mu.Unlock()

		ctx2, cancel2 := context.WithTimeout(context.Background(), time.Duration(params.TimeoutMin)*time.Minute)
		response2, err := callOllamaAPI(ctx2, params.Model2, messages2)
		cancel2()
		if err != nil {
			fyne.Do(func() {
				g.appendMessage(Message{
					Role:      "error",
					Content:   fmt.Sprintf("Ошибка от %s: %v", params.Model2, err),
					Timestamp: time.Now().Format("15:04:05"),
					IsBot:     false,
				})
			})
			break
		}
		msg2 := Message{
			Role:      "assistant",
			Content:   response2,
			Model:     params.Model2,
			Timestamp: time.Now().Format("15:04:05"),
			IsBot:     true,
			Round:     round,
		}

		g.mu.Lock()
		g.state.History = append(g.state.History, msg2)
		updateDialogContext(g.state.Model2Context, msg2, true)
		updateDialogContext(g.state.Model1Context, Message{
			Role:    "user",
			Content: response2,
			Model:   params.Model2,
		}, false)
		g.mu.Unlock()

		fyne.Do(func() {
			g.appendMessage(msg2)
		})

		time.Sleep(1 * time.Second)
	}

	g.mu.Lock()
	g.state.IsActive = false
	finalRound := g.state.CurrentRound
	g.mu.Unlock()

	fyne.Do(func() {
		g.appendMessage(Message{
			Role:      "system",
			Content:   fmt.Sprintf("Диалог завершен. Всего раундов: %d", finalRound),
			Timestamp: time.Now().Format("15:04:05"),
			IsBot:     false,
		})
		g.startBtn.Enable()
	})
	log.Println("Диалог завершен")
}

// startConversation вызывается по нажатию кнопки "Начать диалог"
func (g *guiApp) startConversation() {
	g.mu.Lock()
	if g.state.IsActive {
		g.mu.Unlock()
		dialog.ShowInformation("Ошибка", "Диалог уже запущен", g.window)
		return
	}
	g.mu.Unlock()

	if g.model1Select.Selected == "" || g.model2Select.Selected == "" {
		dialog.ShowInformation("Ошибка", "Выберите обе модели", g.window)
		return
	}
	if g.model1Select.Selected == g.model2Select.Selected {
		dialog.ShowInformation("Ошибка", "Модели должны быть разными", g.window)
		return
	}
	topic := g.topicEntry.Text
	if topic == "" {
		dialog.ShowInformation("Ошибка", "Введите тему диалога", g.window)
		return
	}
	rounds, err := strconv.Atoi(g.roundsEntry.Text)
	if err != nil || rounds < 1 {
		dialog.ShowInformation("Ошибка", "Некорректное количество раундов", g.window)
		return
	}
	timeout, err := strconv.Atoi(g.timeoutEntry.Text)
	if err != nil || timeout < 1 {
		dialog.ShowInformation("Ошибка", "Некорректный таймаут", g.window)
		return
	}

	params := struct {
		Model1, Model2, Role1, Role2, InitialPrompt string
		MaxRounds, TimeoutMin                       int
	}{
		Model1:        g.model1Select.Selected,
		Model2:        g.model2Select.Selected,
		Role1:         g.role1Entry.Text,
		Role2:         g.role2Entry.Text,
		InitialPrompt: topic,
		MaxRounds:     rounds,
		TimeoutMin:    timeout,
	}

	g.startBtn.Disable()
	g.wg.Add(1)
	go g.runConversation(params)
}

// resetConversation выполняет экстренную остановку и возврат к исходным настройкам
func (g *guiApp) resetConversation() {
	// Блокируем кнопки на время сброса
	g.startBtn.Disable()
	g.resetBtn.Disable()

	// Показываем диалог ожидания
	loadingDialog := dialog.NewInformation("Сброс", "Остановка диалога...", g.window)
	loadingDialog.Show()

	go func() {
		// Останавливаем диалог
		g.mu.Lock()
		wasActive := g.state.IsActive
		if wasActive {
			g.state.IsActive = false
		}
		g.mu.Unlock()

		if wasActive {
			g.wg.Wait() // ждем полной остановки
		}

		// Очищаем состояние и историю
		g.mu.Lock()
		g.state = ConversationState{
			History:       make([]Message, 0),
			CurrentRound:  0,
			IsActive:      false,
			Model1Context: nil,
			Model2Context: nil,
		}
		g.mu.Unlock()

		// Сбрасываем все поля ввода к исходным значениям
		fyne.Do(func() {
			// Очищаем выбор моделей
			g.model1Select.ClearSelected()
			g.model2Select.ClearSelected()
			// Очищаем поля ролей и темы
			g.role1Entry.SetText("")
			g.role2Entry.SetText("")
			g.topicEntry.SetText("")
			// Восстанавливаем значения по умолчанию
			g.roundsEntry.SetText("2")
			g.timeoutEntry.SetText("10")
			// Очищаем чат
			g.chatRichText.Segments = nil
			g.chatRichText.Refresh()
			// Сбрасываем индикатор раунда
			g.roundLabel.SetText("Раунд: 0/0")
			// Разблокируем кнопки
			g.startBtn.Enable()
			g.resetBtn.Enable()
			// Закрываем диалог ожидания
			loadingDialog.Hide()
			// Показываем сообщение о сбросе
			g.appendMessage(Message{
				Role:      "system",
				Content:   "Состояние сброшено. Настройки восстановлены по умолчанию.",
				Timestamp: time.Now().Format("15:04:05"),
				IsBot:     false,
			})
		})
	}()
}

// buildGUI создаёт и размещает все элементы интерфейса
func (g *guiApp) buildGUI() {
	g.status = widget.NewLabel("Готов")
	g.status.TextStyle = fyne.TextStyle{Bold: true}

	g.model1Select = widget.NewSelect([]string{}, func(s string) {})
	g.model1Select.PlaceHolder = "Выберите модель 1"
	g.model2Select = widget.NewSelect([]string{}, func(s string) {})
	g.model2Select.PlaceHolder = "Выберите модель 2"
	g.role1Entry = widget.NewEntry()
	g.role1Entry.SetPlaceHolder("Роль модели 1 (необязательно)")
	g.role2Entry = widget.NewEntry()
	g.role2Entry.SetPlaceHolder("Роль модели 2 (необязательно)")
	g.topicEntry = widget.NewMultiLineEntry()
	g.topicEntry.SetPlaceHolder("Тема диалога")
	g.topicEntry.Wrapping = fyne.TextWrapWord
	g.roundsEntry = widget.NewEntry()
	g.roundsEntry.SetText("2")
	g.roundsEntry.Validator = validation.NewRegexp(`^[0-9]+$`, "только число")
	g.timeoutEntry = widget.NewEntry()
	g.timeoutEntry.SetText("10")
	g.timeoutEntry.Validator = validation.NewRegexp(`^[0-9]+$`, "только число")

	g.startBtn = widget.NewButtonWithIcon("Начать диалог", theme.MediaPlayIcon(), g.startConversation)
	g.resetBtn = widget.NewButtonWithIcon("Сброс", theme.ViewRefreshIcon(), g.resetConversation)

	settingsForm := widget.NewForm(
		widget.NewFormItem("Модель 1", g.model1Select),
		widget.NewFormItem("Роль 1", g.role1Entry),
		widget.NewFormItem("Модель 2", g.model2Select),
		widget.NewFormItem("Роль 2", g.role2Entry),
		widget.NewFormItem("Тема диалога", g.topicEntry),
		widget.NewFormItem("Количество раундов", g.roundsEntry),
		widget.NewFormItem("Таймаут (мин)", g.timeoutEntry),
	)
	buttons := container.NewGridWithColumns(2, g.startBtn, g.resetBtn)
	settingsPanel := container.NewBorder(
		container.NewVBox(widget.NewLabelWithStyle("Настройки диалога", fyne.TextAlignLeading, fyne.TextStyle{Bold: true})),
		buttons, nil, nil,
		settingsForm,
	)

	g.chatRichText = widget.NewRichText()
	g.chatRichText.Wrapping = fyne.TextWrapWord
	chatScroll := container.NewScroll(g.chatRichText)
	g.roundLabel = widget.NewLabel("Раунд: 0/0")
	chatHeader := container.NewBorder(nil, nil, nil, g.roundLabel, widget.NewLabelWithStyle("Диалог моделей", fyne.TextAlignLeading, fyne.TextStyle{Bold: true}))
	chatPanel := container.NewBorder(chatHeader, nil, nil, nil, chatScroll)

	split := container.NewHSplit(settingsPanel, chatPanel)
	split.SetOffset(0.35)

	statusBox := container.NewHBox(widget.NewLabel("Статус:"), g.status)
	topBar := container.NewBorder(nil, nil, nil, nil,
		container.NewVBox(statusBox, widget.NewSeparator()))
	content := container.NewBorder(topBar, nil, nil, nil, split)
	g.window.SetContent(content)

	// Загружаем модели в фоне
	go func() {
		g.loadModels()
	}()
}

func main() {
	a := app.New()
	w := a.NewWindow("Диалог LLM моделей Ollama")
	w.Resize(fyne.NewSize(1000, 700))

	gui := &guiApp{
		window: w,
	}
	gui.buildGUI()

	w.ShowAndRun()
}