package ollama

import (
    "bytes"
    "encoding/json"
    "io"
    "net/http"
    "net/http/httptest"
    "os"
    "strings"
    "testing"

    "gitea.miguelmuniz.com/rogueking/junk2jive-server/internal/response"
)

// Mock HTTP client for testing
type mockHTTPClient struct {
    DoFunc func(req *http.Request) (*http.Response, error)
}

func (m *mockHTTPClient) Do(req *http.Request) (*http.Response, error) {
    return m.DoFunc(req)
}

// Helper function to patch the HTTP client for testing
func patchHTTPClient(client interface{}) func() {
    originalClient := http.DefaultClient
    http.DefaultClient = client.(*http.Client)
    return func() {
        http.DefaultClient = originalClient
    }
}

// TestNewOllama tests the NewOllama function
func TestNewOllama(t *testing.T) {
    // Save original env and restore after test
    originalKey := os.Getenv("OLLAMA_API_KEY")
    defer os.Setenv("OLLAMA_API_KEY", originalKey)

    testKey := "test-api-key"
    os.Setenv("OLLAMA_API_KEY", testKey)

    ollama := NewOllama()
    if ollama == nil {
        t.Fatal("Expected NewOllama to return a non-nil value")
    }

    if ollama.OllamaKey != testKey {
        t.Errorf("Expected OllamaKey to be %q, got %q", testKey, ollama.OllamaKey)
    }
}

// TestSendRequest tests the SendRequest method
func TestSendRequest(t *testing.T) {
    tests := []struct {
        name           string
        userMessage    string
        fileID         string
        setupMock      func() *mockHTTPClient
        expectedResult string
        expectError    bool
    }{
        {
            name:        "Successful request without file",
            userMessage: "Hello, world!",
            fileID:      "",
            setupMock: func() *mockHTTPClient {
                return &mockHTTPClient{
                    DoFunc: func(req *http.Request) (*http.Response, error) {
                        // Verify request has correct content and headers
                        if req.Header.Get("Content-Type") != "application/json" {
                            t.Errorf("Expected Content-Type header to be application/json")
                        }
                        if req.Header.Get("Authorization") != "Bearer test-key" {
                            t.Errorf("Expected Authorization header to be Bearer test-key")
                        }

                        // Check request body
                        body, _ := io.ReadAll(req.Body)
                        var reqBody ChatCompletionRequest
                        if err := json.Unmarshal(body, &reqBody); err != nil {
                            t.Errorf("Failed to unmarshal request body: %v", err)
                        }
                        if reqBody.Model != "gemma3:12b" {
                            t.Errorf("Expected model to be gemma3:12b, got %s", reqBody.Model)
                        }
                        if len(reqBody.Messages) != 1 || reqBody.Messages[0].Content != "Hello, world!" {
                            t.Errorf("Expected message content to be 'Hello, world!'")
                        }
                        if len(reqBody.Files) != 0 {
                            t.Errorf("Expected no files, got %d", len(reqBody.Files))
                        }

                        // Return successful response
                        return &http.Response{
                            StatusCode: http.StatusOK,
                            Body:       io.NopCloser(bytes.NewBufferString(`{"response": "Hello from Ollama!"}`)),
                        }, nil
                    },
                }
            },
            expectedResult: `{"response": "Hello from Ollama!"}`,
            expectError:    false,
        },
        {
            name:        "Successful request with file",
            userMessage: "Analyze this file",
            fileID:      "file123",
            setupMock: func() *mockHTTPClient {
                return &mockHTTPClient{
                    DoFunc: func(req *http.Request) (*http.Response, error) {
                        // Check request body to ensure file was included
                        body, _ := io.ReadAll(req.Body)
                        var reqBody ChatCompletionRequest
                        json.Unmarshal(body, &reqBody)
                        
                        if len(reqBody.Files) != 1 || reqBody.Files[0].ID != "file123" {
                            t.Errorf("Expected file ID to be file123")
                        }

                        // Return successful response
                        return &http.Response{
                            StatusCode: http.StatusOK,
                            Body:       io.NopCloser(bytes.NewBufferString(`{"response": "File analysis complete"}`)),
                        }, nil
                    },
                }
            },
            expectedResult: `{"response": "File analysis complete"}`,
            expectError:    false,
        },
        {
            name:        "HTTP error response",
            userMessage: "Hello, world!",
            fileID:      "",
            setupMock: func() *mockHTTPClient {
                return &mockHTTPClient{
                    DoFunc: func(req *http.Request) (*http.Response, error) {
                        return &http.Response{
                            StatusCode: http.StatusBadRequest,
                            Body:       io.NopCloser(bytes.NewBufferString(`{"error": "Bad request"}`)),
                        }, nil
                    },
                }
            },
            expectedResult: "",
            expectError:    true,
        },
    }

    for _, tc := range tests {
        t.Run(tc.name, func(t *testing.T) {
            // Create mock client
            mockClient := tc.setupMock()
            
            // Create a custom HTTP client and replace the default
            client := &http.Client{
                Transport: &mockTransport{mockClient: mockClient},
            }
            
            // Initialize ollama with test key
            ollama := &Ollama{OllamaKey: "test-key"}
            
            // Mock response.RespondWithError to prevent nil pointer dereference
            originalRespondWithError := response.RespondWithError
            response.RespondWithError = func(w http.ResponseWriter, r *http.Request, status int, message string, err error) {}
            defer func() { response.RespondWithError = originalRespondWithError }()
            
            // Call the function
            result, err := ollama.SendRequest(tc.userMessage, tc.fileID)
            
            // Check results
            if tc.expectError && err == nil {
                t.Error("Expected an error but got none")
            }
            
            if !tc.expectError && err != nil {
                t.Errorf("Did not expect an error but got: %v", err)
            }
            
            if !tc.expectError && result != tc.expectedResult {
                t.Errorf("Expected result %q, got %q", tc.expectedResult, result)
            }
        })
    }
}

// Custom transport that uses our mockClient
type mockTransport struct {
    mockClient *mockHTTPClient
}

func (m *mockTransport) RoundTrip(req *http.Request) (*http.Response, error) {
    return m.mockClient.Do(req)
}

// TestOllamaRequest tests the OllamaRequest HTTP handler
func TestOllamaRequest(t *testing.T) {
    tests := []struct {
        name               string
        requestBody        string
        setupMock          func() (*Ollama, func())
        expectedStatusCode int
        expectedResponse   string
    }{
        {
            name:        "Valid request",
            requestBody: `{"message": "Hello", "fileId": ""}`,
            setupMock: func() (*Ollama, func()) {
                // Create a mock Ollama that returns a successful response
                mockOllama := &Ollama{}
                
                // Save the original NewOllama function and replace it
                originalNewOllama := NewOllama
                NewOllama = func() *Ollama {
                    return mockOllama
                }
                
                // Save the original SendRequest method and replace it
                originalSendRequest := mockOllama.SendRequest
                mockOllama.SendRequest = func(userMessage, fileID string) (string, error) {
                    return `{"content": "Hello from Ollama!"}`, nil
                }
                
                // Return the mock and a cleanup function
                return mockOllama, func() {
                    NewOllama = originalNewOllama
                    mockOllama.SendRequest = originalSendRequest
                }
            },
            expectedStatusCode: http.StatusOK,
            expectedResponse:   `{"response":{"content": "Hello from Ollama!"}}`,
        },
        {
            name:        "Invalid JSON request",
            requestBody: `{"message": "Hello"`, // Missing closing brace
            setupMock: func() (*Ollama, func()) {
                mockOllama := &Ollama{}
                originalNewOllama := NewOllama
                NewOllama = func() *Ollama {
                    return mockOllama
                }
                return mockOllama, func() {
                    NewOllama = originalNewOllama
                }
            },
            expectedStatusCode: http.StatusBadRequest,
            expectedResponse:   "", // We don't care about the exact error message
        },
        {
            name:        "API error",
            requestBody: `{"message": "Error", "fileId": ""}`,
            setupMock: func() (*Ollama, func()) {
                mockOllama := &Ollama{}
                originalNewOllama := NewOllama
                NewOllama = func() *Ollama {
                    return mockOllama
                }
                originalSendRequest := mockOllama.SendRequest
                mockOllama.SendRequest = func(userMessage, fileID string) (string, error) {
                    return "", &httpError{message: "API error", statusCode: http.StatusInternalServerError}
                }
                return mockOllama, func() {
                    NewOllama = originalNewOllama
                    mockOllama.SendRequest = originalSendRequest
                }
            },
            expectedStatusCode: http.StatusInternalServerError,
            expectedResponse:   "", // We don't care about the exact error message
        },
    }

    for _, tc := range tests {
        t.Run(tc.name, func(t *testing.T) {
            // Setup mock
            _, cleanup := tc.setupMock()
            defer cleanup()

            // Create request and response recorder
            req := httptest.NewRequest("POST", "/api/ollama", strings.NewReader(tc.requestBody))
            req.Header.Set("Content-Type", "application/json")
            w := httptest.NewRecorder()

            // Call the handler
            OllamaRequest(w, req)

            // Check status code
            if w.Code != tc.expectedStatusCode {
                t.Errorf("Expected status code %d, got %d", tc.expectedStatusCode, w.Code)
            }

            // For successful responses, check the response body
            if tc.expectedStatusCode == http.StatusOK {
                // Remove whitespace for comparison
                expectedNoSpace := strings.ReplaceAll(tc.expectedResponse, " ", "")
                actualNoSpace := strings.ReplaceAll(w.Body.String(), " ", "")
                actualNoSpace = strings.ReplaceAll(actualNoSpace, "\n", "")
                
                if expectedNoSpace != actualNoSpace {
                    t.Errorf("Expected response %q, got %q", tc.expectedResponse, w.Body.String())
                }
            }
        })
    }
}

// Helper error type for testing
type httpError struct {
    message    string
    statusCode int
}

func (e *httpError) Error() string {
    return e.message
}