package main

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"regexp"
	"strings"
)

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

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

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

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

	var ctr Containers

	for _, host := range hosts.Hosts {
		fmt.Printf("get ollama containers: %s -> %s\n", host.IP, host.Port)
		ctr.Adds(get_ollama_containers(host, false))
	}

	retval := []Model{}

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

			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 (h *Hosts) Add(s string) {
	ip_port := strings.Split(s, ":")
	h.Hosts = append(h.Hosts, Host{IP: ip_port[0], Port: ip_port[1]})
}

func (h *Hosts) Init(server string) error {

	if debug {
		fmt.Printf("server: %s\n", server)
	}

	servers := strings.Split(server, ",")

	if len(servers) == 0 {
		return fmt.Errorf("no servers specified in config file: '%s'", server)
	}

	if len(servers) == 1 && servers[0] == "" {
		return fmt.Errorf("no servers specified in config file: '%s'", server)
	}

	for _, server := range servers {
		if !regexp.MustCompile(`^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}:\d{1,5}$`).MatchString(server) {
			return fmt.Errorf("invalid server specified: '%s'", server)
		}
		h.Add(server)
	}

	return nil
}