package server

import (
	"bufio"
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"
)

// The assistant: a chat in the panel that sees the workspace and suggests
// commands. The browser sends the talk and what it shows; the server asks
// the model - its address, kind and key are the admins', set in Настройки
// and never sent to the browser - and streams the answer back as text.

const (
	cfgAIType  = "ai_type" // gradio, openai or anthropic
	cfgAIURL   = "ai_url"
	cfgAIKey   = "ai_key"
	cfgAIModel = "ai_model"

	// DefaultAIURL is the proxy to the Qwen space the assistant uses by default.
	DefaultAIURL = "https://agentproxy.minomix.ru"
)

// aiUserKey is the config key of a user's own switch: the assistant is off
// for everyone until they turn it on in their profile.
func aiUserKey(id int64) string { return "ai_user_" + strconv.FormatInt(id, 10) }

type aiConfig struct{ Type, URL, Key, Model string }

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

func (s *Server) aiConfig() aiConfig {
	get := func(k, def string) string {
		if v, _ := s.db.GetConfig(k); v != "" {
			return v
		}
		return def
	}
	return aiConfig{Type: get(cfgAIType, "gradio"), URL: strings.TrimRight(get(cfgAIURL, DefaultAIURL), "/"), Key: get(cfgAIKey, ""), Model: get(cfgAIModel, "")}
}

var aiClient = &http.Client{Timeout: 10 * time.Minute}

// aiStream asks the model for the answer to msgs (the last one the user's)
// under the system prompt, handing emit the text as it comes.
func aiStream(ctx context.Context, cfg aiConfig, system string, msgs []aiMsg, emit func(string)) error {
	switch cfg.Type {
	case "openai":
		return aiOpenAI(ctx, cfg, system, msgs, emit)
	case "anthropic":
		return aiAnthropic(ctx, cfg, system, msgs, emit)
	default:
		return aiGradio(ctx, cfg, system, msgs, emit)
	}
}

// sseData calls fn with each "data:" payload of an event stream until fn
// says stop or the stream ends.
func sseData(r io.Reader, fn func(data string) (stop bool, err error)) error {
	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 64<<10), 16<<20)
	for sc.Scan() {
		line := sc.Text()
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		stop, err := fn(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
		if err != nil || stop {
			return err
		}
	}
	return sc.Err()
}

func aiPost(ctx context.Context, url string, body any, hdr map[string]string) (*http.Response, error) {
	b, _ := json.Marshal(body)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(b))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	for k, v := range hdr {
		req.Header.Set(k, v)
	}
	resp, err := aiClient.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode/100 != 2 {
		msg, _ := io.ReadAll(io.LimitReader(resp.Body, 400))
		resp.Body.Close()
		return nil, fmt.Errorf("модель ответила %d: %s", resp.StatusCode, strings.TrimSpace(string(msg)))
	}
	return resp, nil
}

func aiOpenAI(ctx context.Context, cfg aiConfig, system string, msgs []aiMsg, emit func(string)) error {
	all := append([]aiMsg{{Role: "system", Content: system}}, msgs...)
	hdr := map[string]string{}
	if cfg.Key != "" {
		hdr["Authorization"] = "Bearer " + cfg.Key
	}
	resp, err := aiPost(ctx, cfg.URL+"/chat/completions", map[string]any{"model": cfg.Model, "stream": true, "messages": all}, hdr)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	return sseData(resp.Body, func(d string) (bool, error) {
		if d == "[DONE]" {
			return true, nil
		}
		var m struct {
			Choices []struct {
				Delta struct {
					Content string `json:"content"`
				} `json:"delta"`
			} `json:"choices"`
			Error *struct {
				Message string `json:"message"`
			} `json:"error"`
		}
		if json.Unmarshal([]byte(d), &m) != nil {
			return false, nil
		}
		if m.Error != nil {
			return true, errors.New(m.Error.Message)
		}
		for _, c := range m.Choices {
			if c.Delta.Content != "" {
				emit(c.Delta.Content)
			}
		}
		return false, nil
	})
}

func aiAnthropic(ctx context.Context, cfg aiConfig, system string, msgs []aiMsg, emit func(string)) error {
	model := cfg.Model
	if model == "" {
		model = "claude-sonnet-5-5"
	}
	resp, err := aiPost(ctx, cfg.URL+"/v1/messages", map[string]any{"model": model, "max_tokens": 4096, "system": system, "messages": msgs, "stream": true},
		map[string]string{"x-api-key": cfg.Key, "anthropic-version": "2023-06-01"})
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	return sseData(resp.Body, func(d string) (bool, error) {
		var m struct {
			Type  string `json:"type"`
			Delta struct {
				Text string `json:"text"`
			} `json:"delta"`
			Error struct {
				Message string `json:"message"`
			} `json:"error"`
		}
		if json.Unmarshal([]byte(d), &m) != nil {
			return false, nil
		}
		switch m.Type {
		case "content_block_delta":
			emit(m.Delta.Text)
		case "message_stop":
			return true, nil
		case "error":
			return true, errors.New(m.Error.Message)
		}
		return false, nil
	})
}

// aiGradio talks to a Gradio chat app (chat_predict of the Qwen Omni demo):
// the history and the new text go in through the queue, and the chatbot's
// history comes back - whole, or as diffs - with the answer growing at its
// end.
func aiGradio(ctx context.Context, cfg aiConfig, system string, msgs []aiMsg, emit func(string)) error {
	hdr := map[string]string{}
	if cfg.Key != "" {
		hdr["X-Agent-Key"] = cfg.Key
	}
	get := func(url string) (*http.Response, error) {
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
		if err != nil {
			return nil, err
		}
		for k, v := range hdr {
			req.Header.Set(k, v)
		}
		resp, err := aiClient.Do(req)
		if err == nil && resp.StatusCode/100 != 2 {
			resp.Body.Close()
			return nil, fmt.Errorf("модель ответила %d", resp.StatusCode)
		}
		return resp, err
	}
	// Which function is chat_predict, from the app's config.
	resp, err := get(cfg.URL + "/config")
	if err != nil {
		return err
	}
	var conf struct {
		Dependencies []struct {
			ID      int     `json:"id"`
			APIName string  `json:"api_name"`
			Targets [][]any `json:"targets"`
		} `json:"dependencies"`
	}
	err = json.NewDecoder(resp.Body).Decode(&conf)
	resp.Body.Close()
	if err != nil {
		return errors.New("модель: не прочиталась её конфигурация")
	}
	fn, trigger := -1, any(nil)
	for _, d := range conf.Dependencies {
		if d.APIName == "chat_predict" {
			fn = d.ID
			if len(d.Targets) > 0 && len(d.Targets[0]) > 0 {
				trigger = d.Targets[0][0]
			}
		}
	}
	if fn < 0 {
		return errors.New("модель: нет функции chat_predict")
	}
	last := msgs[len(msgs)-1]
	b := make([]byte, 8)
	_, _ = rand.Read(b)
	session := hex.EncodeToString(b)
	jr, err := aiPost(ctx, cfg.URL+"/gradio_api/queue/join", map[string]any{
		"data":     []any{last.Content, nil, nil, nil, msgs[:len(msgs)-1], system, 0.7, 0.8, 20},
		"fn_index": fn, "trigger_id": trigger, "event_data": nil, "session_hash": session,
	}, hdr)
	if err != nil {
		return err
	}
	jr.Body.Close()
	resp, err = get(cfg.URL + "/gradio_api/queue/data?session_hash=" + session)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	sent := "" // what of the answer has gone out
	show := func(text string) {
		if strings.HasPrefix(text, sent) {
			emit(text[len(sent):])
			sent = text
		}
	}
	answer := "" // the answer as the chatbot has it now
	return sseData(resp.Body, func(d string) (bool, error) {
		var m struct {
			Msg     string `json:"msg"`
			Success *bool  `json:"success"`
			Output  struct {
				Data  []json.RawMessage `json:"data"`
				Error string            `json:"error"`
			} `json:"output"`
		}
		if json.Unmarshal([]byte(d), &m) != nil {
			return false, nil
		}
		switch m.Msg {
		case "process_generating", "process_completed":
		case "close_stream":
			return true, nil
		default:
			return false, nil
		}
		if m.Success != nil && !*m.Success {
			msg := m.Output.Error
			if msg == "" {
				msg = "модель не ответила"
			}
			return true, errors.New("модель: " + msg)
		}
		if len(m.Output.Data) > 4 {
			answer = gradioAnswer(m.Output.Data[4], answer)
			show(answer)
		}
		return m.Msg == "process_completed", nil
	})
}

// gradioAnswer is the assistant's last message after one update of the
// chatbot: a whole history, or diffs ("add", "append", "replace") to it.
func gradioAnswer(raw json.RawMessage, cur string) string {
	var whole []struct {
		Role    string `json:"role"`
		Content any    `json:"content"`
	}
	if json.Unmarshal(raw, &whole) == nil {
		if n := len(whole); n > 0 {
			if s, ok := whole[n-1].Content.(string); ok && whole[n-1].Role == "assistant" {
				return s
			}
		}
		return cur
	}
	var diffs [][]json.RawMessage
	if json.Unmarshal(raw, &diffs) != nil {
		return cur
	}
	for _, d := range diffs {
		if len(d) < 3 {
			continue
		}
		var op string
		_ = json.Unmarshal(d[0], &op)
		var path []any
		_ = json.Unmarshal(d[1], &path)
		switch {
		case op == "add" && len(path) == 1:
			var msg struct {
				Role    string `json:"role"`
				Content any    `json:"content"`
			}
			if json.Unmarshal(d[2], &msg) == nil && msg.Role == "assistant" {
				s, _ := msg.Content.(string)
				cur = s
			}
		case (op == "append" || op == "replace") && len(path) == 2 && path[1] == "content":
			var s string
			if json.Unmarshal(d[2], &s) == nil {
				if op == "append" {
					cur += s
				} else {
					cur = s
				}
			}
		}
	}
	return cur
}

const aiSystem = `Ты - ассистент панели LPRA (удалённый доступ к машинам: терминалы, файлы, порты). Отвечай по-русски, коротко и по делу.
Ниже - что пользователь сейчас видит в панели: его машины, вкладки, активная панель, текст терминала или список файлов.
Действия - отдельными блоками, панель выполнит их с правами пользователя:
` + "```run машина\nкоманда\n```" + ` - команда в терминале этой машины (машину можно не указывать - тогда активная);
` + "```read машина:/путь\n```" + ` - прочитать файл;
` + "```ls машина:/путь\n```" + ` - список папки;
` + "```write машина:/путь\nполное новое содержимое файла\n```" + ` - записать файл целиком.
Имя машины - как в списке машин. Одна команда на блок, без пояснений внутри блока. Результат действия придёт следующим сообщением - не выдумывай его.
Опасные команды (удаление, перезапуск, изменение системы) - только с предупреждением. Перед правкой файла сначала прочитай его.`

const aiAgent = `
Режим агента: ты сам ведёшь задачу шаг за шагом. В каждом ответе - коротко, что делаешь, и ровно одно действие. Получив результат, решай следующий шаг.
Когда задача решена (или решить нельзя) - ответь итогом без блоков действий.`

// POST /api/panel/ai/chat - {messages, context}: the answer, streamed as
// plain text. A failure after the start ends the text with "[ошибка: …]".
func (s *Server) handleAIChat(w http.ResponseWriter, r *http.Request) {
	a := s.currentAdmin(r)
	if on, _ := s.db.GetConfig(aiUserKey(a.ID)); on != "1" {
		http.Error(w, "ассистент выключен в профиле", http.StatusForbidden)
		return
	}
	var in struct {
		Messages []aiMsg `json:"messages"`
		Context  string  `json:"context"`
		Agent    bool    `json:"agent"`
	}
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&in); err != nil || len(in.Messages) == 0 {
		http.Error(w, "плохой запрос", http.StatusBadRequest)
		return
	}
	// The last 30 messages are enough to go on.
	if n := len(in.Messages); n > 30 {
		in.Messages = in.Messages[n-30:]
	}
	for i := range in.Messages {
		if in.Messages[i].Role != "assistant" {
			in.Messages[i].Role = "user"
		}
	}
	if len(in.Context) > 60000 {
		in.Context = in.Context[len(in.Context)-60000:]
		for !utf8.ValidString(in.Context) && in.Context != "" {
			in.Context = in.Context[1:]
		}
	}
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.Header().Set("Cache-Control", "no-store")
	w.Header().Set("X-Accel-Buffering", "no")
	fl, _ := w.(http.Flusher)
	wrote := false
	emit := func(t string) {
		if t == "" {
			return
		}
		wrote = true
		_, _ = io.WriteString(w, t)
		if fl != nil {
			fl.Flush()
		}
	}
	system := aiSystem
	if in.Agent {
		system += aiAgent
	}
	err := aiStream(r.Context(), s.aiConfig(), system+"\n\n--- Панель сейчас ---\n"+in.Context, in.Messages, emit)
	if err != nil && r.Context().Err() == nil {
		if !wrote {
			w.WriteHeader(http.StatusBadGateway)
		}
		_, _ = io.WriteString(w, "\n[ошибка: "+err.Error()+"]")
	}
}

// GET /api/panel/ai - whether the assistant is on for this user.
func (s *Server) handleAIState(w http.ResponseWriter, r *http.Request) {
	on, _ := s.db.GetConfig(aiUserKey(s.currentAdmin(r).ID))
	writeJSON(w, map[string]any{"on": on == "1"})
}

// POST /api/panel/ai/toggle - {on}: the user's own switch.
func (s *Server) handleAIToggle(w http.ResponseWriter, r *http.Request) {
	var in struct {
		On bool `json:"on"`
	}
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<10)).Decode(&in); err != nil {
		writeJSON(w, map[string]any{"ok": false, "error": "плохой запрос"})
		return
	}
	v := ""
	if in.On {
		v = "1"
	}
	_ = s.db.SetConfig(aiUserKey(s.currentAdmin(r).ID), v)
	writeJSON(w, map[string]any{"ok": true})
}

// POST /settings/ai - the model: kind, address, model name and key (an
// empty key keeps the one there).
func (s *Server) handleSetAI(w http.ResponseWriter, r *http.Request) {
	if !checkCSRF(r) {
		http.Error(w, "bad csrf", http.StatusForbidden)
		return
	}
	typ := r.FormValue("ai_type")
	if typ != "openai" && typ != "anthropic" {
		typ = "gradio"
	}
	_ = s.db.SetConfig(cfgAIType, typ)
	_ = s.db.SetConfig(cfgAIURL, strings.TrimSpace(r.FormValue("ai_url")))
	_ = s.db.SetConfig(cfgAIModel, strings.TrimSpace(r.FormValue("ai_model")))
	if k := strings.TrimSpace(r.FormValue("ai_key")); k != "" {
		_ = s.db.SetConfig(cfgAIKey, k)
	}
	if r.FormValue("ai_key_clear") == "1" {
		_ = s.db.SetConfig(cfgAIKey, "")
	}
	http.Redirect(w, r, "/settings#ai", http.StatusSeeOther)
}