package main

import (
	"bytes"
	"encoding/json"
	"github.com/gorilla/websocket"
	"github.com/ollama/ollama/api"
	"html/template"
	"io"
	"log"
	"net"
	"net/http"
	"strings"
)

var upgrader = websocket.Upgrader{}

func runAsWeb() {
	log.Println("Starting web server")
	fs := http.FileServer(http.Dir("static"))
	http.Handle("/static/", http.StripPrefix("/static/", fs))
	http.HandleFunc("/", servePage)
	http.HandleFunc("/ws", serveWebSocket)

	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		log.Fatal(err)
	}
	log.Printf("Listening on http://127.0.0.1:%d", listener.Addr().(*net.TCPAddr).Port)

	err = http.Serve(listener, nil)
	if err != nil {
		log.Fatal(err)
	}
}

func executeTemplate(w io.Writer, name string, data any) error {
	t, err := template.
		New("templates").
		Funcs(template.FuncMap{
			"trim": strings.TrimSpace,
		}).
		ParseGlob("templates/*.gohtml")
	if err != nil {
		return err
	}
	err = t.ExecuteTemplate(w, name, data)
	if err != nil {
		return err
	}
	return nil
}

func servePage(w http.ResponseWriter, _ *http.Request) {
	wm := WebModel{}
	wm.AddMessages(conversation)
	err := executeTemplate(w, "index", wm)
	if err != nil {
		log.Println(err)
	}
}

func serveWebSocket(w http.ResponseWriter, r *http.Request) {
	log.Println("Websocket connected")
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("upgrade: ", err)
		return
	}

	defer c.Close()
	for {
		// Read the prompt
		mt, message, err := c.ReadMessage()
		if err != nil {
			log.Println("Read error: ", err)
			break
		}
		formInput := make(map[string]any)
		err = json.Unmarshal(message, &formInput)
		if err != nil {
			log.Println("Unmarshal: ", err)
			break
		}
		userRequest, ok := formInput["message"].(string)
		if !ok {
			log.Println("Invalid user input: ", err)
			break
		}

		err = func() error {
			var content bytes.Buffer
			err := executeTemplate(&content, "message", WebMessage{
				Id:      len(conversation) + 1,
				New:     true,
				Replace: false,
				Content: userRequest,
			})
			if err != nil {
				return err
			}
			return c.WriteMessage(mt, content.Bytes())
		}()
		if err != nil {
			log.Println("Sending initial response: ", err)
			break
		}

		// Send the request to Ollama
		fullResponse := ""
		firstResponse := true
		sendPromptInput(userRequest, func(response api.ChatResponse) error {
			fullResponse = fullResponse + response.Message.Content
			var content bytes.Buffer
			err = executeTemplate(&content, "message", WebMessage{
				Id:      len(conversation) + 1,
				New:     firstResponse,
				Replace: !firstResponse,
				Content: fullResponse,
			})
			if err != nil {
				return err
			}
			firstResponse = false
			return c.WriteMessage(mt, content.Bytes())
		})
	}
}