package gtm

import (
	"testing"
)

func TestParseEnrichmentJSON(t *testing.T) {
	tests := []struct {
		name    string
		input   string
		want    *CompanyEnrichment
		wantErr bool
	}{
		{
			name: "valid JSON",
			input: `{
				"company_name": "Acme AI",
				"headquarter": "San Francisco",
				"headquarter_country_code": "us",
				"vertical": "AI/ML",
				"business_model": "B2B",
				"approximate_annual_revenue": "$10M",
				"funding_series": "Series A",
				"ai_confidence_level": "high"
			}`,
			want: &CompanyEnrichment{
				CompanyName:            "Acme AI",
				Headquarter:            "San Francisco",
				HeadquarterCountryCode: "US",
				Vertical:               "AI/ML",
				BusinessModel:          "B2B",
				AnnualRevenue:          "$10M",
				FundingSeries:          "Series A",
				ConfidenceLevel:        "high",
			},
			wantErr: false,
		},
		{
			name: "JSON with markdown wrapper",
			input: "```json\n" + `{
				"company_name": "Test Corp",
				"headquarter": "NY",
				"vertical": "Tech",
				"business_model": "B2B",
				"approximate_annual_revenue": "$5M",
				"funding_series": "Seed",
				"ai_confidence_level": "medium"
			}` + "\n```",
			want: &CompanyEnrichment{
				CompanyName:     "Test Corp",
				Headquarter:     "NY",
				Vertical:        "Tech",
				BusinessModel:   "B2B",
				AnnualRevenue:   "$5M",
				FundingSeries:   "Seed",
				ConfidenceLevel: "medium",
			},
			wantErr: false,
		},
		{
			name:    "invalid JSON",
			input:   "{invalid}",
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := ParseEnrichmentJSON(tt.input)
			if (err != nil) != tt.wantErr {
				t.Errorf("ParseEnrichmentJSON() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !tt.wantErr && !compareEnrichment(got, tt.want) {
				t.Errorf("ParseEnrichmentJSON() = %+v, want %+v", got, tt.want)
			}
		})
	}
}

func TestParseEnrichmentText(t *testing.T) {
	tests := []struct {
		name  string
		input string
		want  *CompanyEnrichment
	}{
		{
			name: "text format with high confidence",
			input: `- Company Name: Acme Corp
- Headquarter: New York
- Vertical: Technology
- Business Model: B2B SaaS
- Approximate Annual Revenue: $50M
- Funding Series: Series B
AI confidence level: high`,
			want: &CompanyEnrichment{
				CompanyName:     "Acme Corp",
				Headquarter:     "New York",
				Vertical:        "Technology",
				BusinessModel:   "B2B SaaS",
				AnnualRevenue:   "$50M",
				FundingSeries:   "Series B",
				ConfidenceLevel: "high",
			},
		},
		{
			name: "text format with medium confidence",
			input: `- Company Name: Test Corp
- Vertical: AI
- Business Model: B2B
confidence level: medium`,
			want: &CompanyEnrichment{
				CompanyName:     "Test Corp",
				Vertical:        "AI",
				BusinessModel:   "B2B",
				ConfidenceLevel: "medium",
			},
		},
		{
			name: "text without confidence indicator",
			input: `Company Name: Test Inc
Vertical: AI/ML
Business Model: B2C`,
			want: &CompanyEnrichment{
				CompanyName:     "Test Inc",
				Vertical:        "AI/ML",
				BusinessModel:   "B2C",
				ConfidenceLevel: "low",
			},
		},
		{
			name: "text with 'high' in other context",
			input: `Company Name: High Tech Inc
Vertical: Technology with very high growth
Business Model: B2B
Confidence: low`,
			want: &CompanyEnrichment{
				CompanyName:     "High Tech Inc",
				Vertical:        "Technology with very high growth",
				BusinessModel:   "B2B",
				ConfidenceLevel: "low",
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := ParseEnrichmentText(tt.input)
			if !compareEnrichment(got, tt.want) {
				t.Errorf("ParseEnrichmentText() = %+v, want %+v", got, tt.want)
			}
		})
	}
}

func TestParseEnrichment(t *testing.T) {
	tests := []struct {
		name  string
		input string
		want  *CompanyEnrichment
	}{
		{
			name: "JSON input",
			input: `{
				"company_name": "JSON Corp",
				"headquarter": "SF",
				"vertical": "Tech",
				"business_model": "B2B",
				"approximate_annual_revenue": "$1M",
				"funding_series": "Seed",
				"ai_confidence_level": "high"
			}`,
			want: &CompanyEnrichment{
				CompanyName:     "JSON Corp",
				Headquarter:     "SF",
				Vertical:        "Tech",
				BusinessModel:   "B2B",
				AnnualRevenue:   "$1M",
				FundingSeries:   "Seed",
				ConfidenceLevel: "high",
			},
		},
		{
			name: "Text input (fallback)",
			input: `Company Name: Text Corp
Vertical: AI`,
			want: &CompanyEnrichment{
				CompanyName:     "Text Corp",
				Vertical:        "AI",
				ConfidenceLevel: "low",
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := ParseEnrichment(tt.input)
			if err != nil {
				t.Errorf("ParseEnrichment() error = %v", err)
				return
			}
			if !compareEnrichment(got, tt.want) {
				t.Errorf("ParseEnrichment() = %+v, want %+v", got, tt.want)
			}
		})
	}
}

func compareEnrichment(a, b *CompanyEnrichment) bool {
	if a == nil || b == nil {
		return a == b
	}
	return a.CompanyName == b.CompanyName &&
		a.Headquarter == b.Headquarter &&
		a.HeadquarterCountryCode == b.HeadquarterCountryCode &&
		a.Vertical == b.Vertical &&
		a.BusinessModel == b.BusinessModel &&
		a.AnnualRevenue == b.AnnualRevenue &&
		a.FundingSeries == b.FundingSeries &&
		a.ConfidenceLevel == b.ConfidenceLevel
}