package main

import (
	"context"
	"fmt"
	"os"
	"time"

	"github.com/tidbcloud/devops/pkg/infra/utils/pointer"
	"google.golang.org/genai"
)

// Test script for LLM email domain verification
// Usage: GEMINI_API_KEY=your_key go run test_llm_prompt.go <email_domain>
// Example: GEMINI_API_KEY=your_key go run test_llm_prompt.go salesforce.com

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: GEMINI_API_KEY=your_key go run test_llm_prompt.go <email_domain>")
		fmt.Println("Example: GEMINI_API_KEY=your_key go run test_llm_prompt.go salesforce.com")
		os.Exit(1)
	}

	emailDomain := os.Args[1]

	// Get API key from environment
	apiKey := os.Getenv("GEMINI_API_KEY")
	if apiKey == "" {
		fmt.Println("Error: GEMINI_API_KEY environment variable not set")
		fmt.Println("\nUsage: GEMINI_API_KEY=your_key go run test_llm_prompt.go <email_domain>")
		os.Exit(1)
	}

	// Initialize Gemini client
	geminiClient, err := genai.NewClient(
		context.Background(),
		&genai.ClientConfig{
			Backend: genai.BackendGeminiAPI,
			APIKey:  apiKey,
		},
	)
	if err != nil {
		fmt.Printf("Error creating Gemini client: %v\n", err)
		os.Exit(1)
	}

	// Build the prompt (same as in production code)
	prompt := fmt.Sprintf("I just received an email from %s .", emailDomain) + `Search the web to find the company using this exact email domain:

IMPORTANT VERIFICATION STEPS:
1. Search for companies associated with the domain "` + emailDomain + `"
2. For each potential match, visit the company's official website
3. Check the webpage content (contact page, about page, footer, etc.) for email addresses or domain references
4. ONLY return information if the webpage content explicitly confirms they use the email domain "` + emailDomain + `"
5. If you find multiple companies with similar domain names, use webpage content to identify the correct one
6. If the webpage shows a DIFFERENT email domain, discard that result and continue searching
7. If no company webpage confirms using "` + emailDomain + `", return "Unknown" for all fields

Respond in the following structured fields only, each in one short phrase (no long paragraphs):

- Company Name:
- Headquarter:
- Vertical:
- Business Model:
- Approximate Annual Revenue:
- Funding Series:
(Please notice that these info are search from the internet and validate as needed. AI confidence level: xx%)

If information is unavailable or cannot be verified from webpage content, write 'Unknown' rather than inventing facts.
Responses must exactly follow the above format, no additional explanation.`

	fmt.Printf("Testing LLM prompt for email domain: %s\n", emailDomain)
	fmt.Println("=================================================")
	fmt.Println()

	// Set 2 minute timeout for LLM search (same as production)
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()

	// Call Gemini (same as llm.SimpleSearch)
	resp, err := geminiClient.Models.GenerateContent(
		ctx,
		"gemini-flash-latest",
		genai.Text(prompt),
		&genai.GenerateContentConfig{
			Temperature: pointer.Pointer(float32(0.0)),
			Tools: []*genai.Tool{
				{
					GoogleSearch: &genai.GoogleSearch{},
				},
			},
		},
	)
	if err != nil {
		fmt.Printf("Error calling Gemini: %v\n", err)
		os.Exit(1)
	}

	result := resp.Text()

	fmt.Println("LLM Response:")
	fmt.Println("=================================================")
	fmt.Println(result)
	fmt.Println("=================================================")
	fmt.Println()

	// Parse and display structured fields
	fmt.Println("Parsed Fields:")
	fmt.Println("=================================================")
	fields := parseFields(result)
	for key, value := range fields {
		fmt.Printf("%-30s: %s\n", key, value)
	}
}

func parseFields(response string) map[string]string {
	fields := make(map[string]string)
	lines := splitLines(response)

	for _, line := range lines {
		if line == "" {
			continue
		}
		// Remove leading dash
		if len(line) > 0 && line[0] == '-' {
			line = line[1:]
		}
		line = trim(line)

		// Split by colon
		parts := splitN(line, ":", 2)
		if len(parts) == 2 {
			key := trim(parts[0])
			value := trim(parts[1])
			if key != "" && value != "" {
				fields[key] = value
			}
		}
	}

	return fields
}

// Helper functions to avoid importing strings package
func splitLines(s string) []string {
	var result []string
	start := 0
	for i := 0; i < len(s); i++ {
		if s[i] == '\n' {
			result = append(result, s[start:i])
			start = i + 1
		}
	}
	if start < len(s) {
		result = append(result, s[start:])
	}
	return result
}

func splitN(s, sep string, n int) []string {
	var result []string
	start := 0
	count := 0

	for i := 0; i <= len(s)-len(sep) && count < n-1; i++ {
		if s[i:i+len(sep)] == sep {
			result = append(result, s[start:i])
			start = i + len(sep)
			count++
		}
	}

	if start < len(s) {
		result = append(result, s[start:])
	}

	return result
}

func trim(s string) string {
	start := 0
	end := len(s)

	// Trim leading whitespace
	for start < end && (s[start] == ' ' || s[start] == '\t' || s[start] == '\n' || s[start] == '\r') {
		start++
	}

	// Trim trailing whitespace
	for end > start && (s[end-1] == ' ' || s[end-1] == '\t' || s[end-1] == '\n' || s[end-1] == '\r') {
		end--
	}

	return s[start:end]
}