package agent

import (
	"agent/src/common"
	"agent/src/llm"
	"agent/src/utils/log"
	"agent/src/workflow"
	"context"
	"fmt"
	"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("/simpletest2025", a.serveTestPage)
		// sm.HandleFunc("/doc", a.serveDoc)
		// sm.HandleFunc("/assets/", a.serveAssets)
		sm.HandleFunc("/ws", a.serveWs)

		addr := fmt.Sprintf(":%d", port)
		log.Info("[agent] start websocket 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) 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:])
}

func (a *Agent) serveWs(w http.ResponseWriter, r *http.Request) {
	// TODO IP连接数限制
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Error("[agent] serveWs", "err", err)
		return
	}
	log.Info("[agent] new connection", "remote", conn.RemoteAddr())
	c := newClient(conn, a)

	// check api key
	if a.apicheck.check(c, r) {
		c.run()
		a.clients.Set(c, struct{}{})
	} else {
		go func() {
			c.SendText("HTTP/1.1 401 Unauthorized\r\n\r\n")
			conn.Close()
		}()
	}
}

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


