package agent

import (
	"agent/src/common"
	"agent/src/llm"
	"agent/src/utils/log"
	"agent/src/workflow"
	"context"
	"fmt"
	"io"
	"net/http"
	"sync"
)

const AgentWorkers = 4

type Agent struct {
	ctx      context.Context
	apicheck *keyChecker
	clients  *common.ThreadSafeMap[*client, struct{}]
}

var once sync.Once

func NewAgent(ctx context.Context, ollamaUrls []string) (*Agent, error) {
	if ollamaUrls == nil {
		ollamaUrls = []string{llm.DefaultOllamaUrl}
	}
	err := llm.InitLocalModel(ollamaUrls)
	if err != nil {
		return nil, err
	}
	return &Agent{
		ctx:     ctx,
		clients: new(common.ThreadSafeMap[*client, struct{}]).Init(nil, false),
	}, nil
}

func (a *Agent) Start(port uint64, crtDir *string) {
	workflow.InitTaskPool(a.ctx)

	go func() {
		sm := http.NewServeMux()
		sm.HandleFunc("/sse", a.serveSSE)

		addr := fmt.Sprintf(":%d", port)
		log.Info("[agent] start SSE server", "addr", addr)

		var err error
		if crtDir == nil {
			err = http.ListenAndServe(addr, sm)
		} else {
			err = http.ListenAndServeTLS(addr, fmt.Sprintf("%s/server.crt", *crtDir), fmt.Sprintf("%s/server.key", *crtDir), sm)
		}

		if err != nil {
			log.Crit("ListenAndServe: ", "err", err)
		}
	}()
}

func (a *Agent) serveSSE(w http.ResponseWriter, r *http.Request) {
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
		return
	}

	// 创建新 client
	client := newClient(w, flusher, a)
	a.addClient(client)

	// 设置 SSE 必备的响应头
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Access-Control-Allow-Origin", "*")

	// 启动写消息循环
	client.run()

	// 读取请求体
	msg, err := io.ReadAll(r.Body)
	if err != nil {
		client.SendText("Invalid request body")
		return
	}
	defer r.Body.Close()

	// 处理请求
	client.processSSERequest(msg)

	// 阻塞直到 client 关闭
	<-client.GetCtx().Done()
}

func (a *Agent) serveTestPage(w http.ResponseWriter, r *http.Request) {
	log.Info("serveHome", "url", r.URL)
	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	http.ServeFile(w, r, "home.html")
}

func (a *Agent) serveDoc(w http.ResponseWriter, r *http.Request) {
	log.Info("serveHome", "url", r.URL)
	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	http.ServeFile(w, r, "./assets/agent服务协议.html")
}

func (a *Agent) serveAssets(w http.ResponseWriter, r *http.Request) {
	log.Info("serveAssets", "url", r.URL)
	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	fmt.Println(r.URL.Path[1:])
	http.ServeFile(w, r, "./"+r.URL.Path[1:])
}

// 给Agent添加addClient方法
func (a *Agent) addClient(c *client) {
	a.clients.Set(c, struct{}{}) // 放进去
	log.Info("client added", "connected", a.clients.Size())
}

func (a *Agent) removeClient(c *client) {
	a.clients.Delete(c)
}
