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"`
}

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 {
	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"`
}

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("%s -> %s\n", host.Name, host.IP)
		ctr.Adds(get_ollama_containers(host))
		// ctr = append(ctr, get_ollama_containers(host)...)
		fmt.Println(ctr)
	}

	retval := []Model{}

	for _, ollama := range ctr.Containers {
		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
		}

		fmt.Printf("mod: %+v\n", mods)

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

	return Models{Models: retval}, nil
}

func get_ollama_containers(host Host) 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 {
		panic(err)
	}
	defer cli.Close()

	containers, err := cli.ContainerList(context.Background(), container.ListOptions{All: false})
	if err != nil {
		panic(err)
	}

	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}
			fmt.Println(ctr.Ports[0])
			c.IP = host.IP
			c.Port = int(ctr.Ports[0].PublicPort)
			ctr_list.Add(c)
			// ctr_list = append(ctr_list, c)
		}
	}

	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")
}

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)
	})

	fmt.Println("liste")
	r.Run(":" + port)

	fmt.Print("END")
}