package main

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

// OllamaClient communicates with the Ollama API.
type OllamaClient struct {
    baseURL    string
    httpClient *http.Client
}

// NewOllamaClient creates a new client for the given Ollama server URL.
func NewOllamaClient(baseURL string) *OllamaClient {
    return &OllamaClient{
        baseURL: strings.TrimRight(baseURL, "/"),
        httpClient: &http.Client{
            Timeout: 30 * time.Second,
        },
    }
}

// ChatMessage represents a message in the chat conversation.
type ChatMessage struct {
    Role    string `json:"role"`
    Content string `json:"content"`
}

// ChatRequest is sent to Ollama's /api/chat endpoint.
type ChatRequest struct {
    Model    string         `json:"model"`
    Messages []ChatMessage  `json:"messages"`
    Stream   bool           `json:"stream"`
    Options  map[string]any `json:"options,omitempty"`
}

// ChatResponse is the non-streaming response from Ollama.
type ChatResponse struct {
    Message         ChatMessage `json:"message"`
    Done            bool        `json:"done"`
    EvalCount       int         `json:"eval_count"`
    PromptEvalCount int         `json:"prompt_eval_count"`
}

// ListModelsResponse is the response from /api/tags.
type ListModelsResponse struct {
    Models []struct {
        Name string `json:"name"`
    } `json:"models"`
}

// ChatCompletion sends a chat request and returns the full response.
func (c *OllamaClient) ChatCompletion(req *ChatRequest) (*ChatResponse, error) {
    req.Stream = false
    body, err := json.Marshal(req)
    if err != nil {
        return nil, fmt.Errorf("marshal request: %w", err)
    }

    httpReq, err := http.NewRequest("POST", c.baseURL+"/api/chat", strings.NewReader(string(body)))
    if err != nil {
        return nil, fmt.Errorf("create request: %w", err)
    }
    httpReq.Header.Set("Content-Type", "application/json")

    resp, err := c.httpClient.Do(httpReq)
    if err != nil {
        return nil, fmt.Errorf("http request: %w", err)
    }
    defer resp.Body.Close()

    if resp.StatusCode != http.StatusOK {
        respBody, _ := io.ReadAll(resp.Body)
        return nil, fmt.Errorf("ollama returned status %d: %s", resp.StatusCode, strings.TrimSpace(string(respBody)))
    }

    var chatResp ChatResponse
    if err := json.NewDecoder(resp.Body).Decode(&chatResp); err != nil {
        return nil, fmt.Errorf("decode response: %w", err)
    }

    return &chatResp, nil
}

// ListModels fetches the list of available models from Ollama.
func (c *OllamaClient) ListModels() ([]string, error) {
    httpReq, err := http.NewRequest("GET", c.baseURL+"/api/tags", nil)
    if err != nil {
        return nil, fmt.Errorf("create request: %w", err)
    }

    resp, err := c.httpClient.Do(httpReq)
    if err != nil {
        return nil, fmt.Errorf("http request: %w", err)
    }
    defer resp.Body.Close()

    if resp.StatusCode != http.StatusOK {
        respBody, _ := io.ReadAll(resp.Body)
        return nil, fmt.Errorf("ollama returned status %d: %s", resp.StatusCode, strings.TrimSpace(string(respBody)))
    }

    var listResp ListModelsResponse
    if err := json.NewDecoder(resp.Body).Decode(&listResp); err != nil {
        return nil, fmt.Errorf("decode response: %w", err)
    }

    models := make([]string, len(listResp.Models))
    for i, m := range listResp.Models {
        models[i] = m.Name
    }
    return models, nil
}