package segments

import (
	"testing"

	"github.com/jandedobbeleer/oh-my-posh/src/cache"
	"github.com/jandedobbeleer/oh-my-posh/src/runtime/mock"
	"github.com/jandedobbeleer/oh-my-posh/src/segments/options"
	"github.com/jandedobbeleer/oh-my-posh/src/text"

	"github.com/stretchr/testify/assert"
)

func TestClaudeSegment(t *testing.T) {
	cases := []struct {
		Case                string
		ClaudeData          *ClaudeData
		ExpectedModel       string
		ExpectedSession     string
		ExpectedGitWorktree string
		ExpectedEnabled     bool
	}{
		{
			Case:            "No cache data",
			ClaudeData:      nil,
			ExpectedEnabled: false,
		},
		{
			Case: "Valid cache data with all fields",
			ClaudeData: &ClaudeData{
				SessionID: "abc123",
				Model: ClaudeModel{
					ID:          "claude-opus-4-1",
					DisplayName: "Opus",
				},
				Workspace: ClaudeWorkspace{
					CurrentDir:  "/repo/project/.worktrees/feature",
					ProjectDir:  "/repo/project",
					GitWorktree: "/repo/project/.worktrees/feature",
				},
				Cost: ClaudeCost{
					TotalCostUSD:       0.01,
					TotalDurationMS:    45000,
					TotalAPIDurationMS: 30000,
					TotalLinesAdded:    156,
					TotalLinesRemoved:  23,
				},
				ContextWindow: ClaudeContextWindow{
					TotalInputTokens:  15234,
					TotalOutputTokens: 4521,
					ContextWindowSize: 200000,
					CurrentUsage: &ClaudeCurrentUsage{
						InputTokens:  8500,
						OutputTokens: 1200,
					},
				},
				RateLimits: &ClaudeRateLimits{
					FiveHour: &ClaudeRateLimitWindow{
						UsedPercentage: new(24.5),
						ResetsAt:       new(int64(1711180800)),
					},
					SevenDay: &ClaudeRateLimitWindow{
						UsedPercentage: new(45.0),
						ResetsAt:       new(int64(1711612800)),
					},
				},
			},
			ExpectedEnabled:     true,
			ExpectedModel:       "Opus",
			ExpectedSession:     "abc123",
			ExpectedGitWorktree: "/repo/project/.worktrees/feature",
		},
		{
			Case: "Valid cache data with partial fields",
			ClaudeData: &ClaudeData{
				SessionID: "xyz789",
				Model: ClaudeModel{
					ID:          "claude-sonnet-3-5",
					DisplayName: "Sonnet 3.5",
				},
				ContextWindow: ClaudeContextWindow{
					TotalInputTokens:  1000,
					TotalOutputTokens: 500,
					ContextWindowSize: 100000,
				},
			},
			ExpectedEnabled: true,
			ExpectedModel:   "Sonnet 3.5",
			ExpectedSession: "xyz789",
		},
	}

	for _, tc := range cases {
		// Setup cache for test
		if tc.ClaudeData != nil {
			cache.Set(cache.Session, cache.CLAUDECACHE, *tc.ClaudeData, cache.INFINITE)
		} else {
			cache.Delete(cache.Session, cache.CLAUDECACHE)
		}

		env := new(mock.Environment)
		claude := &Claude{
			Base: Base{
				env:     env,
				options: options.Map{},
			},
		}

		enabled := claude.Enabled()
		assert.Equal(t, tc.ExpectedEnabled, enabled, tc.Case)

		if tc.ExpectedEnabled {
			assert.Equal(t, tc.ExpectedModel, claude.Model.DisplayName, tc.Case)
			assert.Equal(t, tc.ExpectedSession, claude.SessionID, tc.Case)
			assert.Equal(t, tc.ExpectedGitWorktree, claude.Workspace.GitWorktree, tc.Case)
		}
	}
}

func TestClaudeWorkspaceGitWorktree(t *testing.T) {
	t.Cleanup(func() {
		cache.Delete(cache.Session, cache.CLAUDECACHE)
	})

	cases := []struct {
		Case                string
		Workspace           ClaudeWorkspace
		ExpectedGitWorktree string
	}{
		{
			Case: "Inside a linked worktree",
			Workspace: ClaudeWorkspace{
				CurrentDir:  "/repo/project/.worktrees/feature",
				ProjectDir:  "/repo/project",
				GitWorktree: "/repo/project/.worktrees/feature",
			},
			ExpectedGitWorktree: "/repo/project/.worktrees/feature",
		},
		{
			Case: "Outside a linked worktree",
			Workspace: ClaudeWorkspace{
				CurrentDir: "/repo/project",
				ProjectDir: "/repo/project",
			},
			ExpectedGitWorktree: "",
		},
	}

	for _, tc := range cases {
		cache.Set(cache.Session, cache.CLAUDECACHE, ClaudeData{Workspace: tc.Workspace}, cache.INFINITE)

		env := new(mock.Environment)
		claude := &Claude{
			Base: Base{
				env:     env,
				options: options.Map{},
			},
		}

		assert.True(t, claude.Enabled(), tc.Case)
		assert.Equal(t, tc.ExpectedGitWorktree, claude.Workspace.GitWorktree, tc.Case)
	}
}

func TestClaudeTokenUsagePercent(t *testing.T) {
	cases := []struct {
		UsedPercentage           *int
		Case                     string
		InputTokens              int
		OutputTokens             int
		CurrentInput             int
		CacheCreationInputTokens int
		CacheReadInputTokens     int
		ContextWindow            int
		ExpectedPercent          text.Percentage
		HasCurrentUsage          bool
	}{
		{
			Case:            "Uses UsedPercentage when available",
			UsedPercentage:  new(42),
			ContextWindow:   200000,
			ExpectedPercent: 42,
		},
		{
			Case:            "UsedPercentage capped at 100",
			UsedPercentage:  new(150),
			ContextWindow:   200000,
			ExpectedPercent: 100,
		},
		{
			Case:            "UsedPercentage zero is valid",
			UsedPercentage:  new(0),
			ContextWindow:   200000,
			ExpectedPercent: 0,
		},
		{
			Case:            "Context reset - both UsedPercentage and CurrentUsage nil",
			UsedPercentage:  nil,
			HasCurrentUsage: false,
			InputTokens:     50000, // High cumulative total - should be ignored
			OutputTokens:    50000,
			ContextWindow:   200000,
			ExpectedPercent: 0, // Should return 0 after reset, not fallback to total
		},
		{
			Case:            "Zero context window (no UsedPercentage)",
			HasCurrentUsage: true,
			InputTokens:     1000,
			OutputTokens:    500,
			ContextWindow:   0,
			ExpectedPercent: 0,
		},
		{
			Case:            "10% usage (fallback to total)",
			HasCurrentUsage: true,
			InputTokens:     8000,
			OutputTokens:    2000,
			ContextWindow:   100000,
			ExpectedPercent: 10,
		},
		{
			Case:            "50% usage (fallback to total)",
			HasCurrentUsage: true,
			InputTokens:     50000,
			OutputTokens:    50000,
			ContextWindow:   200000,
			ExpectedPercent: 50,
		},
		{
			Case:            "Over 100% usage (capped)",
			HasCurrentUsage: true,
			InputTokens:     100000,
			OutputTokens:    50000,
			ContextWindow:   100000,
			ExpectedPercent: 100,
		},
		{
			Case:            "Uses CurrentUsage input tokens",
			HasCurrentUsage: true,
			InputTokens:     100000, // High cumulative total
			OutputTokens:    50000,
			CurrentInput:    20000, // Current context input
			ContextWindow:   200000,
			ExpectedPercent: 10, // Should use current input (20000/200000 = 10%)
		},
		{
			Case:                     "Uses CurrentUsage with cache tokens",
			HasCurrentUsage:          true,
			InputTokens:              100000, // High cumulative total
			OutputTokens:             50000,
			CurrentInput:             10000,
			CacheCreationInputTokens: 5000,
			CacheReadInputTokens:     5000,
			ContextWindow:            200000,
			ExpectedPercent:          10, // (10000+5000+5000)/200000 = 10%
		},
		{
			Case:            "Uses CurrentUsage after compact (low current, high total)",
			HasCurrentUsage: true,
			InputTokens:     100000, // High cumulative total
			OutputTokens:    50000,
			CurrentInput:    6000, // Low current context (after compact)
			ContextWindow:   200000,
			ExpectedPercent: 3, // Should use current (6000/200000 = 3%)
		},
		{
			Case:            "Fallback to total when CurrentUsage is zero",
			HasCurrentUsage: true,
			InputTokens:     20000,
			OutputTokens:    10000,
			CurrentInput:    0,
			ContextWindow:   100000,
			ExpectedPercent: 30, // Should fallback to total (30000/100000 = 30%)
		},
	}

	for _, tc := range cases {
		claude := &Claude{}
		claude.ContextWindow.TotalInputTokens = tc.InputTokens
		claude.ContextWindow.TotalOutputTokens = tc.OutputTokens
		if tc.HasCurrentUsage {
			claude.ContextWindow.CurrentUsage = &ClaudeCurrentUsage{
				InputTokens:              tc.CurrentInput,
				CacheCreationInputTokens: tc.CacheCreationInputTokens,
				CacheReadInputTokens:     tc.CacheReadInputTokens,
			}
		}
		claude.ContextWindow.UsedPercentage = tc.UsedPercentage
		claude.ContextWindow.ContextWindowSize = tc.ContextWindow

		percent := claude.TokenUsagePercent()
		assert.Equal(t, tc.ExpectedPercent, percent, tc.Case)
	}
}

func TestClaudeFormattedCost(t *testing.T) {
	cases := []struct {
		Case         string
		ExpectedCost string
		CostUSD      float64
	}{
		{
			Case:         "Very small cost",
			CostUSD:      0.0012,
			ExpectedCost: "$0.0012",
		},
		{
			Case:         "Small cost",
			CostUSD:      0.0099,
			ExpectedCost: "$0.0099",
		},
		{
			Case:         "Regular cost",
			CostUSD:      0.15,
			ExpectedCost: "$0.15",
		},
		{
			Case:         "Large cost",
			CostUSD:      12.34,
			ExpectedCost: "$12.34",
		},
	}

	for _, tc := range cases {
		claude := &Claude{}
		claude.Cost.TotalCostUSD = tc.CostUSD

		formatted := claude.FormattedCost()
		assert.Equal(t, tc.ExpectedCost, formatted, tc.Case)
	}
}

func TestClaudeFormattedDuration(t *testing.T) {
	cases := []struct {
		Case     string
		Expected string
		MS       DurationMS
	}{
		{Case: "Zero", MS: 0, Expected: "0m 0s"},
		{Case: "Seconds only", MS: 45000, Expected: "0m 45s"},
		{Case: "Minutes and seconds", MS: 125000, Expected: "2m 5s"},
		{Case: "Exact minute", MS: 60000, Expected: "1m 0s"},
	}

	for _, tc := range cases {
		claude := &Claude{}
		claude.Cost.TotalDurationMS = tc.MS
		assert.Equal(t, tc.Expected, claude.FormattedDuration(), tc.Case)
	}
}

func TestClaudeFormattedAPIDuration(t *testing.T) {
	cases := []struct {
		Case     string
		Expected string
		MS       DurationMS
	}{
		{Case: "Zero", MS: 0, Expected: "0m 0s"},
		{Case: "Seconds only", MS: 30000, Expected: "0m 30s"},
		{Case: "Minutes and seconds", MS: 90000, Expected: "1m 30s"},
	}

	for _, tc := range cases {
		claude := &Claude{}
		claude.Cost.TotalAPIDurationMS = tc.MS
		assert.Equal(t, tc.Expected, claude.FormattedAPIDuration(), tc.Case)
	}
}

func TestClaudeFormattedTokens(t *testing.T) {
	cases := []struct {
		Case                     string
		ExpectedFormat           string
		InputTokens              int
		OutputTokens             int
		CurrentInput             int
		CacheCreationInputTokens int
		CacheReadInputTokens     int
		HasCurrentUsage          bool
	}{
		{
			Case:            "Small token count (fallback to total)",
			HasCurrentUsage: true,
			InputTokens:     300,
			OutputTokens:    200,
			ExpectedFormat:  "500",
		},
		{
			Case:            "Thousands (fallback to total)",
			HasCurrentUsage: true,
			InputTokens:     8500,
			OutputTokens:    1500,
			ExpectedFormat:  "10.0K",
		},
		{
			Case:            "Tens of thousands (fallback to total)",
			HasCurrentUsage: true,
			InputTokens:     50000,
			OutputTokens:    25000,
			ExpectedFormat:  "75.0K",
		},
		{
			Case:            "Millions (fallback to total)",
			HasCurrentUsage: true,
			InputTokens:     1500000,
			OutputTokens:    500000,
			ExpectedFormat:  "2.0M",
		},
		{
			Case:            "Uses CurrentUsage input tokens",
			HasCurrentUsage: true,
			InputTokens:     100000, // High cumulative total
			OutputTokens:    50000,
			CurrentInput:    10000, // Current context input
			ExpectedFormat:  "10.0K",
		},
		{
			Case:                     "Uses CurrentUsage with cache tokens",
			HasCurrentUsage:          true,
			InputTokens:              100000, // High cumulative total
			OutputTokens:             50000,
			CurrentInput:             5000,
			CacheCreationInputTokens: 2500,
			CacheReadInputTokens:     2500,
			ExpectedFormat:           "10.0K", // 5000+2500+2500 = 10000
		},
		{
			Case:            "Uses CurrentUsage after compact (low current)",
			HasCurrentUsage: true,
			InputTokens:     500000, // High cumulative total
			OutputTokens:    200000,
			CurrentInput:    500, // Low current context (after compact)
			ExpectedFormat:  "500",
		},
		{
			Case:            "Fallback to total when CurrentUsage is zero",
			HasCurrentUsage: true,
			InputTokens:     50000,
			OutputTokens:    25000,
			CurrentInput:    0,
			ExpectedFormat:  "75.0K", // Should fallback to total
		},
		{
			Case:            "Nil CurrentUsage falls back to total",
			HasCurrentUsage: false,
			InputTokens:     50000,
			OutputTokens:    25000,
			ExpectedFormat:  "75.0K", // Should fallback to total
		},
	}

	for _, tc := range cases {
		claude := &Claude{}
		claude.ContextWindow.TotalInputTokens = tc.InputTokens
		claude.ContextWindow.TotalOutputTokens = tc.OutputTokens
		if tc.HasCurrentUsage {
			claude.ContextWindow.CurrentUsage = &ClaudeCurrentUsage{
				InputTokens:              tc.CurrentInput,
				CacheCreationInputTokens: tc.CacheCreationInputTokens,
				CacheReadInputTokens:     tc.CacheReadInputTokens,
			}
		}

		formatted := claude.FormattedTokens()
		assert.Equal(t, tc.ExpectedFormat, formatted, tc.Case)
	}
}

func TestClaudeRateLimitUsage(t *testing.T) {
	cases := []struct {
		RateLimits    *ClaudeRateLimits
		Case          string
		ExpectedFive  text.Percentage
		ExpectedSeven text.Percentage
	}{
		{
			Case:          "Nil RateLimits",
			RateLimits:    nil,
			ExpectedFive:  0,
			ExpectedSeven: 0,
		},
		{
			Case: "Nil FiveHour window",
			RateLimits: &ClaudeRateLimits{
				SevenDay: &ClaudeRateLimitWindow{UsedPercentage: new(50.0)},
			},
			ExpectedFive:  0,
			ExpectedSeven: 50,
		},
		{
			Case: "Nil SevenDay window",
			RateLimits: &ClaudeRateLimits{
				FiveHour: &ClaudeRateLimitWindow{UsedPercentage: new(25.0)},
			},
			ExpectedFive:  25,
			ExpectedSeven: 0,
		},
		{
			Case: "Nil UsedPercentage",
			RateLimits: &ClaudeRateLimits{
				FiveHour: &ClaudeRateLimitWindow{UsedPercentage: nil},
				SevenDay: &ClaudeRateLimitWindow{UsedPercentage: nil},
			},
			ExpectedFive:  0,
			ExpectedSeven: 0,
		},
		{
			Case: "Valid percentages",
			RateLimits: &ClaudeRateLimits{
				FiveHour: &ClaudeRateLimitWindow{UsedPercentage: new(42.7)},
				SevenDay: &ClaudeRateLimitWindow{UsedPercentage: new(75.3)},
			},
			ExpectedFive:  43,
			ExpectedSeven: 75,
		},
		{
			Case: "Value over 100 capped",
			RateLimits: &ClaudeRateLimits{
				FiveHour: &ClaudeRateLimitWindow{UsedPercentage: new(150.0)},
				SevenDay: &ClaudeRateLimitWindow{UsedPercentage: new(200.0)},
			},
			ExpectedFive:  100,
			ExpectedSeven: 100,
		},
		{
			Case: "Zero percentages",
			RateLimits: &ClaudeRateLimits{
				FiveHour: &ClaudeRateLimitWindow{UsedPercentage: new(0.0)},
				SevenDay: &ClaudeRateLimitWindow{UsedPercentage: new(0.0)},
			},
			ExpectedFive:  0,
			ExpectedSeven: 0,
		},
	}

	for _, tc := range cases {
		claude := &Claude{}
		claude.RateLimits = tc.RateLimits

		assert.Equal(t, tc.ExpectedFive, claude.FiveHourUsage(), tc.Case+" (FiveHour)")
		assert.Equal(t, tc.ExpectedSeven, claude.SevenDayUsage(), tc.Case+" (SevenDay)")
	}
}

func TestClaudeGaugeMethods(t *testing.T) {
	cases := []struct {
		RateLimits             *ClaudeRateLimits
		Case                   string
		MarkedChar             string
		UnmarkedChar           string
		ExpectedTokenGauge     string
		ExpectedTokenGaugeUsed string
		ExpectedFiveHourGauge  string
		ExpectedSevenDayGauge  string
		UsedPercentage         int
	}{
		{
			Case:                   "Default chars (▰▱) at 40% used",
			MarkedChar:             "▰",
			UnmarkedChar:           "▱",
			UsedPercentage:         40,
			ExpectedTokenGauge:     "▰▰▰▱▱", // 60% remaining = 3 blocks
			ExpectedTokenGaugeUsed: "▰▰▱▱▱", // 40% used = 2 blocks
		},
		{
			Case:                   "Custom chars (█░) at 40% used",
			MarkedChar:             "█",
			UnmarkedChar:           "░",
			UsedPercentage:         40,
			ExpectedTokenGauge:     "███░░", // 60% remaining = 3 blocks
			ExpectedTokenGaugeUsed: "██░░░", // 40% used = 2 blocks
		},
		{
			Case:           "Custom chars with rate limits",
			MarkedChar:     "█",
			UnmarkedChar:   "░",
			UsedPercentage: 0,
			RateLimits: &ClaudeRateLimits{
				FiveHour: &ClaudeRateLimitWindow{UsedPercentage: new(60.0)},
				SevenDay: &ClaudeRateLimitWindow{UsedPercentage: new(20.0)},
			},
			ExpectedFiveHourGauge: "███░░", // 60% used = 3 blocks
			ExpectedSevenDayGauge: "█░░░░", // 20% used = 1 block
		},
	}

	for _, tc := range cases {
		t.Run(tc.Case, func(t *testing.T) {
			usedPct := tc.UsedPercentage
			claude := &Claude{
				markedChar:   tc.MarkedChar,
				unmarkedChar: tc.UnmarkedChar,
			}
			claude.ContextWindow.UsedPercentage = &usedPct
			claude.RateLimits = tc.RateLimits

			if tc.ExpectedTokenGauge != "" {
				assert.Equal(t, tc.ExpectedTokenGauge, claude.TokenGauge(), tc.Case+" (TokenGauge)")
			}

			if tc.ExpectedTokenGaugeUsed != "" {
				assert.Equal(t, tc.ExpectedTokenGaugeUsed, claude.TokenGaugeUsed(), tc.Case+" (TokenGaugeUsed)")
			}

			if tc.ExpectedFiveHourGauge != "" {
				assert.Equal(t, tc.ExpectedFiveHourGauge, claude.FiveHourGauge(), tc.Case+" (FiveHourGauge)")
			}

			if tc.ExpectedSevenDayGauge != "" {
				assert.Equal(t, tc.ExpectedSevenDayGauge, claude.SevenDayGauge(), tc.Case+" (SevenDayGauge)")
			}
		})
	}
}

func TestClaudeGaugeOptionsReadInEnabled(t *testing.T) {
	t.Cleanup(func() {
		cache.Delete(cache.Session, cache.CLAUDECACHE)
	})

	cache.Set(cache.Session, cache.CLAUDECACHE, ClaudeData{
		Model: ClaudeModel{DisplayName: "Opus"},
	}, cache.INFINITE)

	env := new(mock.Environment)
	claude := &Claude{
		Base: Base{
			env: env,
			options: options.Map{
				gaugeMarkedChar:   "█",
				gaugeUnmarkedChar: "░",
			},
		},
	}

	assert.True(t, claude.Enabled())
	assert.Equal(t, "█", claude.markedChar)
	assert.Equal(t, "░", claude.unmarkedChar)
}
