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

// StreamEvent represents a single event from a streaming chat response.
type StreamEvent struct {
    Token    string // delta content from the model
    Done     bool   // true for the final event
    Error    error  // non-nil if the stream failed
}

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

// streamingChunk represents a single line from Ollama's ndjson stream.
type streamingChunk struct {
    Message struct {
        Content string `json:"content"`
    } `json:"message"`
    Done            bool `json:"done"`
    EvalCount       int  `json:"eval_count"`
    PromptEvalCount int  `json:"prompt_eval_count"`
}

// ChatCompletionStream sends a streaming chat request and returns a channel
// of StreamEvents. The caller must drain the channel until Done or Error.
func (c *OllamaClient) ChatCompletionStream(req *ChatRequest) (<-chan StreamEvent, error) {
    req.Stream = true
    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)
    }

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

    ch := make(chan StreamEvent)
    go func() {
        defer resp.Body.Close()
        defer close(ch)

        decoder := json.NewDecoder(resp.Body)
        for {
            var chunk streamingChunk
            if err := decoder.Decode(&chunk); err != nil {
                if err == io.EOF {
                    return
                }
                ch <- StreamEvent{Error: fmt.Errorf("decode chunk: %w", err)}
                return
            }

            ch <- StreamEvent{
                Token: chunk.Message.Content,
                Done:  chunk.Done,
            }

            if chunk.Done {
                return
            }
        }
    }()

    return ch, 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
}