package geminiweb

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/go-rod/rod"
	"github.com/go-rod/rod/lib/proto"
)

const geminiURL = "https://gemini.google.com/app"

var modelSelectors = map[string]string{
	"gemini-2.0-flash":        "Flash",
	"gemini-2.5-pro":          "Pro",
	"gemini-2.5-pro-thinking": "Thinking",
}

func NormalizeModel(model string) string {
	if strings.HasPrefix(model, "gemini-") {
		return model
	}
	return "gemini-" + model
}

func GetModelDisplayName(model string) string {
	if name, ok := modelSelectors[model]; ok {
		return name
	}
	return "Flash"
}

func NavigateToGemini(page *rod.Page) error {
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	if err := page.Context(ctx).Navigate(geminiURL); err != nil {
		return fmt.Errorf("failed to navigate to gemini: %w", err)
	}
	return page.Context(ctx).WaitLoad()
}

func IsLoggedIn(page *rod.Page) bool {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	_, err := page.Context(ctx).Element(`[aria-label*="New chat"], [data-test-id*="new-chat"], button[aria-label*="chat"]`)
	return err == nil
}

func SelectModel(page *rod.Page, model string) error {
	displayName := GetModelDisplayName(model)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	modelSwitcher, err := page.Context(ctx).Element(`button[aria-label*="model"], [data-test-id="model-selector"], button[aria-haspopup="listbox"]`)
	if err != nil {
		return fmt.Errorf("model selector not found: %w", err)
	}

	if err := modelSwitcher.Click(proto.InputMouseButtonLeft, 1); err != nil {
		return fmt.Errorf("failed to click model selector: %w", err)
	}

	time.Sleep(500 * time.Millisecond)

	option, err := page.Context(ctx).Element(fmt.Sprintf(`[aria-label*="%s"], [data-value="%s"]`, displayName, displayName))
	if err != nil {
		return fmt.Errorf("model option %s not found: %w", displayName, err)
	}

	return option.Click(proto.InputMouseButtonLeft, 1)
}

func SendPrompt(page *rod.Page, prompt string) error {
	textarea, err := page.Element(`textarea[aria-label*="message"], textarea[placeholder*="message"], rich-textarea, .ql-editor, div[contenteditable="true"]`)
	if err != nil {
		return fmt.Errorf("input field not found: %w", err)
	}

	if err := textarea.Input(prompt); err != nil {
		return fmt.Errorf("failed to input prompt: %w", err)
	}

	time.Sleep(300 * time.Millisecond)

	sendBtn, err := page.Element(`button[aria-label*="Send"], button[aria-label*="submit"], button[type="submit"]`)
	if err != nil {
		return fmt.Errorf("send button not found: %w", err)
	}

	return sendBtn.Click(proto.InputMouseButtonLeft, 1)
}

func WaitForResponse(page *rod.Page, onChunk func(text string), onThinking func(thinking string), onComplete func()) error {
	lastText := ""
	lastThinking := ""
	responseComplete := false

	timeout := time.NewTimer(120 * time.Second)
	defer timeout.Stop()

	ticker := time.NewTicker(500 * time.Millisecond)
	defer ticker.Stop()

	for {
		select {
		case <-timeout.C:
			return fmt.Errorf("response timeout")
		case <-ticker.C:
			textChanged := false

			responseEls, err := page.Elements(`.response-text, message-content, .model-response, div[data-test-id="response"]`)
			if err == nil && len(responseEls) > 0 {
				for _, el := range responseEls {
					text, _ := el.Text()
					text = strings.TrimSpace(text)
					if text != "" && text != lastText {
						if strings.Contains(text, lastText) {
							newPart := strings.TrimPrefix(text, lastText)
							if newPart != "" {
								onChunk(newPart)
							}
						} else {
							onChunk(text)
						}
						lastText = text
						textChanged = true
					}
				}
			}

			thinkingEls, err := page.Elements(`.thinking-content, .thought-text, div[data-test-id="thinking"]`)
			if err == nil && len(thinkingEls) > 0 {
				for _, el := range thinkingEls {
					thinking, _ := el.Text()
					thinking = strings.TrimSpace(thinking)
					if thinking != "" && thinking != lastThinking {
						if strings.Contains(thinking, lastThinking) {
							newPart := strings.TrimPrefix(thinking, lastThinking)
							if newPart != "" {
								onThinking(newPart)
							}
						} else {
							onThinking(thinking)
						}
						lastThinking = thinking
						textChanged = true
					}
				}
			}

			doneBtn, err := page.Element(`button[aria-label*="stop"], button[aria-label*="regenerate"]`)
			if err == nil && doneBtn != nil {
				ariaLabel, _ := doneBtn.Attribute("aria-label")
				if ariaLabel != nil && (*ariaLabel == "Stop" || strings.Contains(*ariaLabel, "regenerate")) {
					if !responseComplete && lastText != "" {
						responseComplete = true
						onComplete()
						return nil
					}
				}
			}

			if !textChanged && responseComplete {
				return nil
			}
		}
	}
}

func IsRateLimited(page *rod.Page) bool {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	el, err := page.Context(ctx).Element(`[class*="rate-limit"], [class*="quota"], [data-test-id="rate-limited"]`)
	return err == nil && el != nil
}

func GetRateLimitMessage(page *rod.Page) string {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	el, err := page.Context(ctx).Element(`[class*="rate-limit"], [class*="quota"], [class*="error-message"]`)
	if err != nil || el == nil {
		return ""
	}

	text, _ := el.Text()
	return strings.TrimSpace(text)
}