"""Bucket company sizes into standard segments."""

from typing import Dict, Any, Union, Optional
import re

__transform_id__ = "company_size_bucket"
__version__ = "1.1.0"
__updated__ = "2024-09-20"

# Compile regex patterns at module level for performance
PLUS_PATTERN = re.compile(r"^(\d+(?:,\d{3})*|\d+[kKmM]?)\+")
APPROX_PATTERN = re.compile(r"^(~|approximately)\s*", re.IGNORECASE)
NUMBER_PATTERN = re.compile(r"(\d+(?:,\d{3})*|\d+[kKmM]?)")
RANGE_SEPARATOR_PATTERN = re.compile(r"[-–]")
NUMBER_EXTRACT_PATTERN = re.compile(r"\d+[kKmM]|\d+(?:,\d{3})*")


def parse_number_with_suffix(text: str) -> Optional[int]:
    """Parse numbers with k/K/m/M suffixes."""
    text = text.strip()
    if not text:
        return None

    # Remove commas and spaces
    text = text.replace(",", "").replace(" ", "")

    # Handle k/K/m/M suffixes
    multiplier = 1
    if text.lower().endswith("k"):
        multiplier = 1000
        text = text[:-1]
    elif text.lower().endswith("m"):
        multiplier = 1000000
        text = text[:-1]

    try:
        return int(float(text) * multiplier)
    except ValueError:
        return None


def _extract_numeric_part(part: str) -> Optional[int]:
    """Extract a numeric value from a range segment."""
    num = parse_number_with_suffix(part)
    if num is not None:
        return num

    literal_suffix = re.search(r"\d+(?:\.\d+)?[kKmM]", part.replace(" ", ""))
    if literal_suffix:
        num = parse_number_with_suffix(literal_suffix.group(0))
        if num is not None:
            return num

    cleaned = re.sub(r"[^0-9,]", "", part)
    if cleaned:
        num = parse_number_with_suffix(cleaned)
        if num is not None:
            return num

    match = NUMBER_EXTRACT_PATTERN.search(part)
    if match:
        return parse_number_with_suffix(match.group(0))
    return None


def company_size_bucket(value: Union[str, int, float]) -> Dict[str, Any]:
    """
    Bucket company employee counts into standard RevOps segments.

    Now handles:
    - "10K+" format
    - "~5,000 employees"
    - "10k-15k" ranges

    Standard B2B segments:
        1-10: "Startup"
        11-50: "Small Business"
        51-200: "SMB"
        201-500: "Mid-Market"
        501-1000: "Commercial"
        1001-5000: "Enterprise"
        5000+: "Strategic/Global"

    Examples:
        15 -> "Small Business"
        250 -> "Mid-Market"
        5000 -> "Enterprise"
        "10-50 employees" -> "Small Business"
        "10K+" -> "Strategic/Global"
        "~5,000 employees" -> "Enterprise"
    """
    # Handle None/empty
    if value is None or value == "":
        return {"value": None, "numeric": None}

    # Handle numeric input directly
    if isinstance(value, (int, float)):
        employee_count = int(value)
    else:
        s = str(value).strip()
        employee_count = None

        if "-" not in s and "–" not in s and "—" not in s:
            suffix_match = re.search(r"\d+[kKmM]", s.replace(" ", ""))
            if suffix_match:
                candidate = parse_number_with_suffix(suffix_match.group(0))
                if candidate:
                    employee_count = candidate

        if employee_count is None:
            # Pattern 1: X+ (e.g., "10K+", "5000+")
            plus_match = PLUS_PATTERN.match(s)
            if plus_match:
                num = parse_number_with_suffix(plus_match.group(1))
                if num:
                    employee_count = num
                    # For X+, we use X as the baseline

            # Pattern 2: ~X or approximately X
            elif s.startswith("~") or s.lower().startswith("approximately"):
                s_clean = APPROX_PATTERN.sub("", s)
                # Extract first number
                num_match = NUMBER_PATTERN.search(s_clean)
                if num_match:
                    employee_count = parse_number_with_suffix(num_match.group(1))

            # Pattern 3: X-Y range (e.g., "10k-15k", "100-250")
            elif "-" in s or "–" in s or "—" in s:
                parts = RANGE_SEPARATOR_PATTERN.split(s)
                if len(parts) == 2:
                    low = _extract_numeric_part(parts[0])
                    high = _extract_numeric_part(parts[1])
                    if low and high:
                        employee_count = (low + high) // 2

            # Pattern 4: Standard number extraction
            else:
                numbers = NUMBER_EXTRACT_PATTERN.findall(s)
                if numbers:
                    # Take the first number or average of range
                    if len(numbers) >= 2:
                        low = parse_number_with_suffix(numbers[0])
                        high = parse_number_with_suffix(numbers[1])
                        if low and high:
                            employee_count = (low + high) // 2
                    else:
                        employee_count = parse_number_with_suffix(numbers[0])

    # Apply bucketing logic
    if employee_count is None:
        return {"value": None, "numeric": None}

    # Define buckets with min/max ranges
    buckets = [
        (1, 10, "Startup", "startup"),
        (11, 50, "Small Business", "smb"),
        (51, 200, "SMB", "smb"),
        (201, 500, "Mid-Market", "mid_market"),
        (501, 1000, "Commercial", "commercial"),
        (1001, 5000, "Enterprise", "enterprise"),
        (5001, 10000, "Large Enterprise", "large_enterprise"),
        (10001, float("inf"), "Strategic/Global", "strategic"),
    ]

    # Find the appropriate bucket
    for min_size, max_size, label, segment_code in buckets:
        if min_size <= employee_count <= max_size:
            result = {
                "value": label,
                "numeric": employee_count,
                "segment_code": segment_code,
                "min": min_size,
                "max": max_size if max_size != float("inf") else None,
            }

            # Add tier for sales routing
            if employee_count <= 50:
                result["tier"] = "T4"  # Self-serve/Inside Sales
            elif employee_count <= 500:
                result["tier"] = "T3"  # SMB Sales
            elif employee_count <= 5000:
                result["tier"] = "T2"  # Enterprise Sales
            else:
                result["tier"] = "T1"  # Strategic Accounts

            # Add ARR potential estimate (rough B2B SaaS benchmarks)
            if employee_count <= 10:
                result["arr_potential"] = "$10K-$50K"
            elif employee_count <= 50:
                result["arr_potential"] = "$50K-$250K"
            elif employee_count <= 200:
                result["arr_potential"] = "$250K-$1M"
            elif employee_count <= 500:
                result["arr_potential"] = "$1M-$5M"
            elif employee_count <= 1000:
                result["arr_potential"] = "$5M-$10M"
            elif employee_count <= 5000:
                result["arr_potential"] = "$10M-$50M"
            elif employee_count <= 10000:
                result["arr_potential"] = "$50M-$100M"
            else:
                result["arr_potential"] = "$100M+"

            return result

    # Fallback for edge cases (0 or negative)
    return {
        "value": None,
        "numeric": employee_count,
        "segment_code": "unknown",
        "tier": "T4",
    }
