package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"

	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/client"
)

type Hosts struct {
	Hosts []Host `json:"hosts"`
}

type Host struct {
	Name       string      `json:"name"`
	IP         string      `json:"ip"`
	Port       string      `json:"port"`
	Containers []Container `json:"containers"`
}

type Containers struct {
	Containers []Container `json:"containers"`
}

type Container struct {
	Name   string `json:"names"`
	Port   int    `json:"port"`
	IP     string `json:"ip"`
	Image  string `json:"image"`
	ID     string `json:"id"`
	Status string `json:"status"`
	State  string `json:"state"`
}

type Details struct {
	Parent_model       string   `json:"parent_model"`
	Format             string   `json:"format"`
	Family             string   `json:"family"`
	Families           []string `json:"families"`
	Parameter_size     string   `json:"parameter_size"`
	Quantization_level string   `json:"quantization_level"`
}

type Model struct {
	Container   string  `json:"container"`
	Name        string  `json:"name"`
	Model       string  `json:"model"`
	Modified_at string  `json:"modified_at"`
	Size        int     `json:"size"`
	Details     Details `json:"details"`
	Ip          string  `json:"ip"`
	Port        int     `json:"port"`
	State       string  `json:"state"`
}

type Models struct {
	Models []Model `json:"models"`
}

func get_ollama_tags(hosts Hosts) (Models, error) {

	var ctr Containers

	for _, host := range hosts.Hosts {
		fmt.Printf("HOST: %s -> %s\n", host.Name, host.IP)
		ctr.Adds(get_ollama_containers(host, false))
		// ctr = append(ctr, get_ollama_containers(host)...)
		// fmt.Println(ctr)
	}

	retval := []Model{}

	for _, ollama := range ctr.Containers {
		// fmt.Println(ollama)
		if ollama.State != "running" {
			retval = append(retval, Model{Ip: ollama.IP, Port: ollama.Port, State: "stopped"})
		} else {
			url := fmt.Sprintf("http://%s:%d/api/tags", ollama.IP, ollama.Port)
			resp, err := http.Get(url)
			if err != nil {
				return Models{}, fmt.Errorf("failed to get %s: %v", url, err)
			}
			defer resp.Body.Close()

			body, err := io.ReadAll(resp.Body)
			if err != nil {
				return Models{}, fmt.Errorf("failed to read body: %v", err)
			}

			mods := Models{}
			err = json.Unmarshal(body, &mods)
			if err != nil {
				return Models{}, fmt.Errorf("failed to unmarshal JSON: %v", err)
			}
			fmt.Printf("%+v\n", mods)

			for n := range mods.Models {
				mods.Models[n].Ip = ollama.IP
				mods.Models[n].Port = ollama.Port
				mods.Models[n].State = "running"
				mods.Models[n].Container = ollama.Name
				fmt.Printf("mod: %s - %s:%d\n", mods.Models[n].Name, mods.Models[n].Ip, mods.Models[n].Port)
			}

			retval = append(retval, mods.Models...)
		}
	}

	return Models{Models: retval}, nil
}

func get_ollama_containers(host Host, all bool) Containers {

	ctr_list := Containers{}

	// fmt.Printf("%s -> %s\n", host.Name, host.IP)
	cli, err := client.NewClientWithOpts(client.WithHost(fmt.Sprintf("tcp://%s:%s", host.IP, host.Port)))
	if err != nil {
		fmt.Println(err)
		return ctr_list
	}
	defer cli.Close()

	containers, err := cli.ContainerList(context.Background(), container.ListOptions{All: all})
	if err != nil {
		fmt.Println(err)
		return ctr_list
	}

	for _, ctr := range containers {
		if strings.Contains(ctr.Names[0], "ollama") { //&& ctr.Image == "ollama/ollama" {
			c := Container{ID: ctr.ID, Name: ctr.Names[0], Image: ctr.Image, State: ctr.State, Status: ctr.Status}
			if ctr.State == "running" {
				// fmt.Println(ctr.Ports[0])
				c.Port = int(ctr.Ports[0].PublicPort)
			}
			c.IP = host.IP
			ctr_list.Add(c)
			// ctr_list = append(ctr_list, c)
		}
	}

	// fmt.Println(ctr_list)

	return ctr_list
}

func (h *Hosts) Add(name string, ip string, port string) {
	h.Hosts = append(h.Hosts, Host{Name: name, IP: ip, Port: port})
}

func (c *Containers) Add(ctr Container) {
	c.Containers = append(c.Containers, ctr)
}

func (c *Containers) Adds(ctrs Containers) {
	c.Containers = append(c.Containers, ctrs.Containers...)
}

func (h *Hosts) Init() {
	h.Add("thor", "192.168.55.13", "2375")
	h.Add("ironman", "192.168.55.10", "2375")
	h.Add("hela", "192.168.55.15", "2375")
}

func main() {

	port := "4000"

	hosts := Hosts{}
	hosts.Init()

	r := gin.Default()
	r.GET("/api/tags", func(c *gin.Context) {
		models, err := get_ollama_tags(hosts)
		if err != nil {
			panic(err)
		}
		c.JSON(200, models)
	})

	sapp := r.Group("/_app")
	sapp.Static("/", "./static/_app")

	static := r.Group("/static")
	static.Static("/", "./static/static")

	r.StaticFile("/", "./static/index.html")
	r.StaticFile("/favicon.png", "./static/favicon.png")

	r.ForwardedByClientIP = true

	fmt.Println("Starting server on port " + port)
	r.Run(":" + port)

	fmt.Print("END")
}