package main

import (
  "encoding/json"
  "fmt"
  "net/http"
  "net/http/httptest"
  "strings"
  "testing"
)

func TestHunyuanServiceEndpoint(t *testing.T) {
  svc := NewHunyuanService("https://example.com/gradio/")
  if svc.endpoint != "https://example.com/gradio" {
    t.Errorf("expected trailing slash stripped, got %q", svc.endpoint)
  }

  models := svc.ListModels()
  if len(models) != 1 || models[0].ID != "hy3" {
    t.Fatalf("expected model ID 'hy3', got %v", models)
  }
}

func TestParseGradioStreamData(t *testing.T) {
  sampleJSON := `[["Hello world", "Thinking process...", null]]`
  c, r, tc, ok := parseGradioStreamData(sampleJSON)
  if !ok {
    t.Fatalf("expected parseGradioStreamData to succeed")
  }
  if c != "Hello world" {
    t.Errorf("content=%q, want 'Hello world'", c)
  }
  if r != "Thinking process..." {
    t.Errorf("reasoning=%q, want 'Thinking process...'", r)
  }
  if len(tc) != 0 {
    t.Errorf("expected 0 tool calls, got %d", len(tc))
  }
}

func TestHunyuanMockedCompletion(t *testing.T) {
  var receivedCallURL, receivedStreamURL string

  handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if strings.HasSuffix(r.URL.Path, "/gradio_api/call/chat") && r.Method == http.MethodPost {
      receivedCallURL = r.URL.Path
      w.Header().Set("Content-Type", "application/json")
      fmt.Fprint(w, `{"event_id":"test-event-123"}`)
      return
    }
    if strings.Contains(r.URL.Path, "/gradio_api/call/chat/test-event-123") && r.Method == http.MethodGet {
      receivedStreamURL = r.URL.Path
      w.Header().Set("Content-Type", "text/event-stream")
      fmt.Fprint(w, "data: [[\"Mocked answer\", \"Mocked reasoning\", null]]\n\n")
      return
    }
    http.NotFound(w, r)
  })

  server := httptest.NewServer(handler)
  defer server.Close()

  svc := NewHunyuanService(server.URL)

  rec := httptest.NewRecorder()
  req, _ := http.NewRequest("POST", "/v1/chat/completions", nil)

  chatReq := ChatCompletionRequest{
    Model: "hy3",
    Messages: []ChatMessage{
      {Role: "user", Content: "Hello!"},
    },
    Stream: false,
  }

  err := svc.Chat(rec, req, chatReq)
  if err != nil {
    t.Fatalf("unexpected Chat error: %v", err)
  }

  if receivedCallURL != "/gradio_api/call/chat" {
    t.Errorf("expected call URL '/gradio_api/call/chat', got %q", receivedCallURL)
  }
  if receivedStreamURL != "/gradio_api/call/chat/test-event-123" {
    t.Errorf("expected stream URL '/gradio_api/call/chat/test-event-123', got %q", receivedStreamURL)
  }

  var resp ChatCompletionResponse
  if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
    t.Fatalf("failed to decode completion response: %v", err)
  }

  if len(resp.Choices) != 1 {
    t.Fatalf("expected 1 choice, got %d", len(resp.Choices))
  }
  if resp.Choices[0].Message.GetContentString() != "Mocked answer" {
    t.Errorf("content=%q, want 'Mocked answer'", resp.Choices[0].Message.GetContentString())
  }
  if resp.Choices[0].Message.ReasoningContent != "Mocked reasoning" {
    t.Errorf("reasoning=%q, want 'Mocked reasoning'", resp.Choices[0].Message.ReasoningContent)
  }
}

func TestParseGradioStreamDataWithToolCalls(t *testing.T) {
  sampleJSON := `[["", "Let me check the weather.", [{"id":"call_999","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"Tokyo\"}"}}]]]`
  c, r, tc, ok := parseGradioStreamData(sampleJSON)
  if !ok {
    t.Fatalf("expected parseGradioStreamData to succeed")
  }
  if c != "" {
    t.Errorf("content=%q, want empty", c)
  }
  if r != "Let me check the weather." {
    t.Errorf("reasoning=%q, want 'Let me check the weather.'", r)
  }
  if len(tc) != 1 {
    t.Fatalf("expected 1 tool call, got %d", len(tc))
  }
  if tc[0].ID != "call_999" || tc[0].Function.Name != "get_weather" || tc[0].Function.Arguments != `{"location":"Tokyo"}` {
    t.Errorf("unexpected tool call payload: %+v", tc[0])
  }
}

func TestHunyuanToolCallingNonStreaming(t *testing.T) {
  handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if strings.HasSuffix(r.URL.Path, "/gradio_api/call/chat") && r.Method == http.MethodPost {
      var body map[string]interface{}
      json.NewDecoder(r.Body).Decode(&body)
      dataArr, _ := body["data"].([]interface{})
      if len(dataArr) < 9 || dataArr[8] == nil || dataArr[8] == "" {
        t.Errorf("expected functionsJSONStr in data[8], got %v", dataArr)
      }

      w.Header().Set("Content-Type", "application/json")
      fmt.Fprint(w, `{"event_id":"evt-tool-1"}`)
      return
    }
    if strings.Contains(r.URL.Path, "/gradio_api/call/chat/evt-tool-1") && r.Method == http.MethodGet {
      w.Header().Set("Content-Type", "text/event-stream")
      fmt.Fprint(w, "data: [[\"\", \"I should call the search tool.\", [{\"id\":\"call_abc\",\"type\":\"function\",\"function\":{\"name\":\"search\",\"arguments\":\"{\\\"query\\\":\\\"golang\\\"}\"}}]]]\n\n")
      return
    }
    http.NotFound(w, r)
  })

  server := httptest.NewServer(handler)
  defer server.Close()

  svc := NewHunyuanService(server.URL)
  rec := httptest.NewRecorder()
  req, _ := http.NewRequest("POST", "/v1/chat/completions", nil)

  chatReq := ChatCompletionRequest{
    Model: "hy3",
    Messages: []ChatMessage{
      {Role: "user", Content: "Search for golang news"},
    },
    Tools: []Tool{
      {
        Type: "function",
        Function: map[string]interface{}{
          "name":        "search",
          "description": "Search web",
        },
      },
    },
    Stream: false,
  }

  err := svc.Chat(rec, req, chatReq)
  if err != nil {
    t.Fatalf("unexpected Chat error: %v", err)
  }

  var resp ChatCompletionResponse
  if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
    t.Fatalf("failed to decode response: %v", err)
  }

  if len(resp.Choices) != 1 {
    t.Fatalf("expected 1 choice, got %d", len(resp.Choices))
  }
  choice := resp.Choices[0]
  if choice.FinishReason != "tool_calls" {
    t.Errorf("finish_reason=%q, want 'tool_calls'", choice.FinishReason)
  }
  if choice.Message.ReasoningContent != "I should call the search tool." {
    t.Errorf("reasoning_content=%q, want 'I should call the search tool.'", choice.Message.ReasoningContent)
  }
  if len(choice.Message.ToolCalls) != 1 {
    t.Fatalf("expected 1 tool call, got %d", len(choice.Message.ToolCalls))
  }
  tc := choice.Message.ToolCalls[0]
  if tc.ID != "call_abc" || tc.Function.Name != "search" || tc.Function.Arguments != `{"query":"golang"}` {
    t.Errorf("unexpected tool call: %+v", tc)
  }
}

func TestHunyuanReasoningAndToolCallingStreaming(t *testing.T) {
  handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if strings.HasSuffix(r.URL.Path, "/gradio_api/call/chat") && r.Method == http.MethodPost {
      w.Header().Set("Content-Type", "application/json")
      fmt.Fprint(w, `{"event_id":"evt-stream-1"}`)
      return
    }
    if strings.Contains(r.URL.Path, "/gradio_api/call/chat/evt-stream-1") && r.Method == http.MethodGet {
      w.Header().Set("Content-Type", "text/event-stream")
      flusher, _ := w.(http.Flusher)

      fmt.Fprint(w, "data: [[\"\", \"Step 1 reasoning.\", null]]\n\n")
      if flusher != nil {
        flusher.Flush()
      }

      fmt.Fprint(w, "data: [[\"Hello \", \"Step 1 reasoning. Step 2 reasoning.\", null]]\n\n")
      if flusher != nil {
        flusher.Flush()
      }

      fmt.Fprint(w, "data: [[\"Hello \", \"Step 1 reasoning. Step 2 reasoning.\", [{\"id\":\"call_xyz\",\"type\":\"function\",\"function\":{\"name\":\"calculator\",\"arguments\":\"{\\\"expr\\\":\\\"2+2\\\"}\"}}]]]\n\n")
      if flusher != nil {
        flusher.Flush()
      }

      return
    }
    http.NotFound(w, r)
  })

  server := httptest.NewServer(handler)
  defer server.Close()

  svc := NewHunyuanService(server.URL)
  rec := httptest.NewRecorder()
  req, _ := http.NewRequest("POST", "/v1/chat/completions", nil)

  chatReq := ChatCompletionRequest{
    Model: "hy3",
    Messages: []ChatMessage{
      {Role: "user", Content: "Calculate 2+2"},
    },
    Stream: true,
  }

  err := svc.Chat(rec, req, chatReq)
  if err != nil {
    t.Fatalf("unexpected Chat error: %v", err)
  }

  bodyStr := rec.Body.String()
  lines := strings.Split(bodyStr, "\n")

  var receivedDeltas []StreamDelta
  var finishReasons []string
  var foundDone bool

  for _, line := range lines {
    line = strings.TrimSpace(line)
    if strings.HasPrefix(line, "data: ") {
      payload := strings.TrimPrefix(line, "data: ")
      if payload == "[DONE]" {
        foundDone = true
        continue
      }
      var streamResp StreamResponse
      if err := json.Unmarshal([]byte(payload), &streamResp); err == nil && len(streamResp.Choices) > 0 {
        choice := streamResp.Choices[0]
        receivedDeltas = append(receivedDeltas, choice.Delta)
        if choice.FinishReason != nil {
          finishReasons = append(finishReasons, *choice.FinishReason)
        }
      }
    }
  }

  if !foundDone {
    t.Errorf("expected 'data: [DONE]' in SSE stream")
  }

  if len(receivedDeltas) < 1 || receivedDeltas[0].Role != "assistant" {
    t.Errorf("expected first delta to set role='assistant', got %+v", receivedDeltas[0])
  }

  var combinedReasoning, combinedContent string
  var toolCallsEmitted []ToolCall
  for _, d := range receivedDeltas {
    combinedReasoning += d.ReasoningContent
    combinedContent += d.Content
    if len(d.ToolCalls) > 0 {
      toolCallsEmitted = append(toolCallsEmitted, d.ToolCalls...)
    }
  }

  if combinedReasoning != "Step 1 reasoning. Step 2 reasoning." {
    t.Errorf("combined reasoning=%q, want 'Step 1 reasoning. Step 2 reasoning.'", combinedReasoning)
  }
  if combinedContent != "Hello " {
    t.Errorf("combined content=%q, want 'Hello '", combinedContent)
  }
  if len(toolCallsEmitted) == 0 {
    t.Fatalf("expected tool call deltas to be emitted")
  }
  if len(finishReasons) != 1 || finishReasons[0] != "tool_calls" {
    t.Errorf("finish_reasons=%v, want ['tool_calls']", finishReasons)
  }
}