package gtm

import (
	"encoding/json"
	"fmt"
	"strings"
)

// ParseEnrichment parses LLM enrichment response, trying JSON first then falling back to text
func ParseEnrichment(rawResponse string) (*CompanyEnrichment, error) {
	// Try JSON parsing first
	enrichment, err := ParseEnrichmentJSON(rawResponse)
	if err == nil {
		return enrichment, nil
	}

	// Fall back to text parsing
	return ParseEnrichmentText(rawResponse), nil
}

// ParseEnrichmentJSON parses JSON format enrichment response
func ParseEnrichmentJSON(rawResponse string) (*CompanyEnrichment, error) {
	// Clean up potential markdown code blocks
	cleaned := strings.TrimSpace(rawResponse)
	cleaned = strings.TrimPrefix(cleaned, "```json")
	cleaned = strings.TrimPrefix(cleaned, "```")
	cleaned = strings.TrimSuffix(cleaned, "```")
	cleaned = strings.TrimSpace(cleaned)

	var result struct {
		CompanyName              string `json:"company_name"`
		Headquarter              string `json:"headquarter"`
		HeadquarterCountryCode   string `json:"headquarter_country_code"`
		Vertical                 string `json:"vertical"`
		BusinessModel            string `json:"business_model"`
		ApproximateAnnualRevenue string `json:"approximate_annual_revenue"`
		FundingSeries            string `json:"funding_series"`
		AIConfidenceLevel        string `json:"ai_confidence_level"`
	}

	err := json.Unmarshal([]byte(cleaned), &result)
	if err != nil {
		return nil, fmt.Errorf("failed to parse enrichment JSON: %w", err)
	}

	return &CompanyEnrichment{
		CompanyName:            result.CompanyName,
		Headquarter:            result.Headquarter,
		HeadquarterCountryCode: strings.ToUpper(strings.TrimSpace(result.HeadquarterCountryCode)),
		Vertical:               result.Vertical,
		BusinessModel:          result.BusinessModel,
		AnnualRevenue:          result.ApproximateAnnualRevenue,
		FundingSeries:          result.FundingSeries,
		ConfidenceLevel:        result.AIConfidenceLevel,
	}, nil
}

// ParseEnrichmentText parses text format enrichment response (fallback)
func ParseEnrichmentText(rawResponse string) *CompanyEnrichment {
	fields := parseLLMTriageFields(rawResponse)

	// Extract confidence level from text
	confidence := "low" // Default for text parsing
	lowerText := strings.ToLower(rawResponse)

	// Look for confidence level patterns
	if strings.Contains(lowerText, "confidence level: high") ||
		strings.Contains(lowerText, "confidence: high") ||
		strings.Contains(lowerText, "ai confidence level: high") {
		confidence = "high"
	} else if strings.Contains(lowerText, "confidence level: medium") ||
		strings.Contains(lowerText, "confidence: medium") ||
		strings.Contains(lowerText, "ai confidence level: medium") {
		confidence = "medium"
	} else if strings.Contains(lowerText, "confidence level: low") ||
		strings.Contains(lowerText, "confidence: low") ||
		strings.Contains(lowerText, "ai confidence level: low") {
		confidence = "low"
	}

	return &CompanyEnrichment{
		CompanyName:            fields["company name"],
		Headquarter:            fields["headquarter"],
		HeadquarterCountryCode: strings.ToUpper(strings.TrimSpace(fields["headquarter country code"])),
		Vertical:               fields["vertical"],
		BusinessModel:          fields["business model"],
		AnnualRevenue:          fields["approximate annual revenue"],
		FundingSeries:          fields["funding series"],
		ConfidenceLevel:        confidence,
	}
}

// parseLLMTriageFields parses key-value pairs from LLM text response
func parseLLMTriageFields(triage string) map[string]string {
	fields := map[string]string{}
	for _, line := range strings.Split(triage, "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		line = strings.TrimPrefix(line, "-")
		line = strings.TrimSpace(line)
		// Normalize full-width colon for key/value split
		line = strings.ReplaceAll(line, "：", ":")
		parts := strings.SplitN(line, ":", 2)
		if len(parts) != 2 {
			continue
		}
		key := strings.ToLower(strings.TrimSpace(parts[0]))
		val := strings.TrimSpace(parts[1])
		if key != "" && val != "" {
			fields[key] = val
		}
	}
	return fields
}