core/address_processor.py

"""
Address Processor - Smart address handling for location.apply and state.apply.

Handles:
- HQ address extraction with priority rules
- US address selection when multiple locations exist
- State abbreviation ↔ full name mapping
- Hiring concentration analysis
- Non-US location filtering
"""

import re
from typing import Optional, Tuple, Dict


# US State abbreviations and full names mapping
US_STATES = {
    "alabama": "Alabama", "alaska": "Alaska", "arizona": "Arizona",
    "arkansas": "Arkansas", "california": "California", "colorado": "Colorado",
    "connecticut": "Connecticut", "delaware": "Delaware", "florida": "Florida",
    "georgia": "Georgia", "hawaii": "Hawaii", "idaho": "Idaho",
    "illinois": "Illinois", "indiana": "Indiana", "iowa": "Iowa",
    "kansas": "Kansas", "kentucky": "Kentucky", "louisiana": "Louisiana",
    "maine": "Maine", "maryland": "Maryland", "massachusetts": "Massachusetts",
    "michigan": "Michigan", "minnesota": "Minnesota", "mississippi": "Mississippi",
    "missouri": "Missouri", "montana": "Montana", "nebraska": "Nebraska",
    "nevada": "Nevada", "new hampshire": "New Hampshire", "new jersey": "New Jersey",
    "new mexico": "New Mexico", "new york": "New York",
    "north carolina": "North Carolina", "north dakota": "North Dakota",
    "ohio": "Ohio", "oklahoma": "Oklahoma", "oregon": "Oregon",
    "pennsylvania": "Pennsylvania", "rhode island": "Rhode Island",
    "south carolina": "South Carolina", "south dakota": "South Dakota",
    "tennessee": "Tennessee", "texas": "Texas", "utah": "Utah",
    "vermont": "Vermont", "virginia": "Virginia", "washington": "Washington",
    "west virginia": "West Virginia", "wisconsin": "Wisconsin", "wyoming": "Wyoming",
    "district of columbia": "District of Columbia", "washington dc": "District of Columbia",
    "dc": "District of Columbia",
}

STATE_ABBREV = {
    "AL": "Alabama", "AK": "Alaska", "AZ": "Arizona", "AR": "Arkansas",
    "CA": "California", "CO": "Colorado", "CT": "Connecticut", "DE": "Delaware",
    "FL": "Florida", "GA": "Georgia", "HI": "Hawaii", "ID": "Idaho",
    "IL": "Illinois", "IN": "Indiana", "IA": "Iowa", "KS": "Kansas",
    "KY": "Kentucky", "LA": "Louisiana", "ME": "Maine", "MD": "Maryland",
    "MA": "Massachusetts", "MI": "Michigan", "MN": "Minnesota", "MS": "Mississippi",
    "MO": "Missouri", "MT": "Montana", "NE": "Nebraska", "NV": "Nevada",
    "NH": "New Hampshire", "NJ": "New Jersey", "NM": "New Mexico", "NY": "New York",
    "NC": "North Carolina", "ND": "North Dakota", "OH": "Ohio", "OK": "Oklahoma",
    "OR": "Oregon", "PA": "Pennsylvania", "RI": "Rhode Island",
    "SC": "South Carolina", "SD": "South Dakota", "TN": "Tennessee",
    "TX": "Texas", "UT": "Utah", "VT": "Vermont", "VA": "Virginia",
    "WA": "Washington", "WV": "West Virginia", "WI": "Wisconsin", "WY": "Wyoming",
    "DC": "District of Columbia",
}

# Reverse mapping: full name → abbreviation
STATE_TO_ABBREV = {v: k for k, v in STATE_ABBREV.items()}
# Add lowercase versions
STATE_TO_ABBREV.update({v.lower(): k for k, v in STATE_ABBREV.items()})


class AddressProcessor:
    """Smart address processing for location.apply and state.apply columns."""

    # Known non-US locations (country/city names)
    NON_US_LOCATIONS = {
        "india", "china", "japan", "korea", "singapore", "hong kong",
        "canada", "uk", "united kingdom", "england", "germany", "france",
        "australia", "brazil", "mexico", "spain", "italy", "netherlands",
        "sweden", "switzerland", "ireland", "dubai", "uae", "bangalore",
        "bengaluru", "mumbai", "delhi", "toronto", "vancouver", "london",
        "berlin", "paris", "sydney", "tokyo", "seoul", "dublin", "amsterdam",
        "stockholm", "zurich", "hyderabad", "pune", "chennai", "gurgaon",
        "noida", "manila", "jakarta", "bangkok", "ho chi minh", "saigon",
        "kuala lumpur", "taipei", "shanghai", "beijing", "shenzhen",
        "warsaw", "lisbon", "madrid", "barcelona", "milan", "rome",
        "vienna", "brussels", "oslo", "helsinki", "copenhagen",
        "auckland", "wellington", "tel aviv", "istanbul", "moscow",
        "cape town", "johannesburg", "nairobi", "lagos", "cairo",
        "buenos aires", "sao paulo", "santiago", "bogota", "lima",
    }

    def determine_location_apply(
        self,
        wellfound_location: str,
        ai_contact_data: Dict,
        scraped_addresses: list,
    ) -> Tuple[Optional[str], Optional[str]]:
        """
        Determine location.apply (city) and state.apply (state).

        Priority:
        1. US HQ address from AI contact extraction
        2. US office from AI contact extraction
        3. Hiring focus location from AI extraction
        4. Wellfound location (if US)
        5. First US address found on website
        """
        location_apply = None
        state_apply = None

        # Priority 1: HQ address if it's in the US
        hq_city = ai_contact_data.get("headquarters_city")
        hq_state = ai_contact_data.get("headquarters_state")
        hq_country = ai_contact_data.get("headquarters_country", "").lower()

        if hq_city and self._is_us_location(hq_state, hq_country):
            location_apply = self._normalize_city(hq_city)
            state_apply = self._normalize_state(hq_state)

        # Priority 2: US office
        if not location_apply:
            us_city = ai_contact_data.get("us_office_city")
            us_state = ai_contact_data.get("us_office_state")
            if us_city:
                state_apply = self._normalize_state(us_state)
                location_apply = self._normalize_city(us_city)

        # Priority 3: Hiring focus location (if US)
        if not location_apply:
            hiring_loc = ai_contact_data.get("hiring_focus_location", "")
            city, state = self._parse_location_string(hiring_loc)
            if city and state and self._is_us_state(state):
                location_apply = self._normalize_city(city)
                state_apply = self._normalize_state(state)

        # Priority 4: Parse wellfound location for US parts
        if not location_apply and wellfound_location:
            city, state = self._extract_us_from_wellfound_location(wellfound_location)
            if city:
                location_apply = self._normalize_city(city)
                if state:
                    state_apply = self._normalize_state(state)

        # Priority 5: First US address from scraping
        if not location_apply and scraped_addresses:
            for addr in scraped_addresses:
                city, state = self._parse_address_string(addr)
                if city and state and self._is_us_state(state):
                    location_apply = self._normalize_city(city)
                    state_apply = self._normalize_state(state)
                    break

        return location_apply, state_apply

    def _is_us_location(self, state: Optional[str], country: Optional[str]) -> bool:
        """Check if a location is in the US."""
        if country and country.lower() in ("us", "usa", "united states", "united states of america"):
            return True
        if state and self._is_us_state(state):
            return True
        return False

    def _is_us_state(self, state: str) -> bool:
        """Check if a string is a valid US state."""
        s = state.strip().lower()
        if s in US_STATES:
            return True
        if s.upper() in STATE_ABBREV:
            return True
        return False

    def _normalize_state(self, state: str) -> Optional[str]:
        """Convert state to full name. Returns None if input is empty or unrecognized."""
        if not state:
            return None
        s = state.strip()
        if s.upper() in STATE_ABBREV:
            return STATE_ABBREV[s.upper()]
        s_lower = s.lower()
        if s_lower in US_STATES:
            return US_STATES[s_lower]
        return s

    def _normalize_city(self, city: str) -> str:
        """Normalize city name — strip trailing punctuation and split on comma."""
        if not city:
            return ""
        city = city.strip().strip(",.;")
        # Handle "City, State" format — keep only the city part
        if "," in city:
            city = city.split(",")[0].strip()
        return city

    def _parse_location_string(self, location: str) -> Tuple[Optional[str], Optional[str]]:
        """Parse 'City, State' or 'City, State, Country' format."""
        if not location:
            return None, None
        parts = [p.strip() for p in location.split(",")]
        if len(parts) >= 2:
            city = parts[0]
            for part in parts[1:]:
                if self._is_us_state(part):
                    return city, part
            return city, None
        return None, None

    def _extract_us_from_wellfound_location(self, location: str) -> Tuple[Optional[str], Optional[str]]:
        """Extract US city/state from wellfound location string like 'San Francisco , New York'."""
        if not location:
            return None, None

        parts = [p.strip() for p in location.split(",")]

        # Try to find a US state abbreviation among the parts
        for part in parts:
            part_stripped = part.strip()
            if self._is_us_state(part_stripped):
                # This is a state — look for preceding city
                idx = parts.index(part_stripped)
                if idx > 0:
                    city = parts[idx - 1].strip()
                    # Verify city is not a non-US location
                    if city.lower() not in self.NON_US_LOCATIONS:
                        return city, part_stripped
                # State name as city? (e.g., "New York")
                if part_stripped.lower() in ("new york", "washington"):
                    return part_stripped, part_stripped

        # Check if any part is a known US city
        us_cities = {
            "san francisco": "California", "new york": "New York",
            "new york city": "New York", "los angeles": "California",
            "chicago": "Illinois", "boston": "Massachusetts",
            "seattle": "Washington", "austin": "Texas",
            "denver": "Colorado", "portland": "Oregon",
            "san diego": "California", "san jose": "California",
            "palo alto": "California", "mountain view": "California",
            "menlo park": "California", "redwood city": "California",
            "sunnyvale": "California", "santa clara": "California",
            "oakland": "California", "berkeley": "California",
            "miami": "Florida", "atlanta": "Georgia",
            "dallas": "Texas", "houston": "Texas",
            "philadelphia": "Pennsylvania", "phoenix": "Arizona",
            "minneapolis": "Minnesota", "detroit": "Michigan",
            "salt lake city": "Utah", "raleigh": "North Carolina",
            "charlotte": "North Carolina", "nashville": "Tennessee",
            "boulder": "Colorado", "cambridge": "Massachusetts",
            "pittsburgh": "Pennsylvania", "ann arbor": "Michigan",
            "madison": "Wisconsin", "irvine": "California",
            "santa monica": "California", "cupertino": "California",
            "bellevue": "Washington", "arlington": "Virginia",
            "reston": "Virginia", "mclean": "Virginia",
            "bethesda": "Maryland", "columbia": "Maryland",
            "st louis": "Missouri", "kansas city": "Missouri",
            "indianapolis": "Indiana", "columbus": "Ohio",
            "cleveland": "Ohio", "cincinnati": "Ohio",
            "tampa": "Florida", "orlando": "Florida",
            "durham": "North Carolina", "chapel hill": "North Carolina",
            "remote": None,
        }

        for part in parts:
            part_lower = part.strip().lower()
            if part_lower in us_cities and us_cities[part_lower]:
                return part.strip(), us_cities[part_lower]

        return None, None

    def _parse_address_string(self, addr: str) -> Tuple[Optional[str], Optional[str]]:
        """Parse a US address string into city and state.

        Handles:
        - "123 St, City, State ZIP" (e.g. "123 Main St, San Francisco, CA 94105")
        - "City, State ZIP" (e.g. "San Francisco, CA 94105")
        - "City, State ZIP+4" (e.g. "New York, NY 10001-1234")
        - "City, ST" (e.g. "Austin, TX")
        - "City, StateName" (e.g. "Austin, Texas")
        """
        if not addr:
            return None, None

        # Strategy: Find a state abbreviation or state name in the string,
        # then look backwards for the preceding city name.

        # Scan for 2-letter state abbreviation
        state_abbr_match = re.search(
            r',\s*([A-Z]{2})\b(?!\s*[a-z])', addr
        )
        if state_abbr_match:
            state = state_abbr_match.group(1)
            if self._is_us_state(state):
                # Everything before the comma is "prefix" — extract city from it
                prefix = addr[:state_abbr_match.start()].rstrip(', ')
                city = self._extract_city_from_prefix(prefix)
                if city:
                    return city, state

        # Scan for full state name after a comma
        # Sort by length descending to match longer names first (e.g. "North Carolina" before "North")
        sorted_states = sorted(
            ((name, abbr) for name, abbr in STATE_TO_ABBREV.items() if abbr),
            key=lambda x: len(x[0]), reverse=True
        )
        for state_name, abbr in sorted_states:
            # Find ALL matches and use the LAST one (closest to end of string)
            pattern = r',\s*\b' + re.escape(state_name) + r'\b'
            matches = list(re.finditer(pattern, addr, re.IGNORECASE))
            if matches:
                last_match = matches[-1]
                prefix = addr[:last_match.start()].rstrip(', ')
                city = self._extract_city_from_prefix(prefix)
                if city:
                    return city, abbr

        return None, None

    def _extract_city_from_prefix(self, prefix: str) -> Optional[str]:
        """Extract city name from the part of address before the state.

        Examples:
        - "123 Main St, San Francisco" → "San Francisco"
        - "San Francisco" → "San Francisco"
        - "100 Tech Blvd, Austin" → "Austin"
        """
        if not prefix:
            return None

        # Split by comma and take the last meaningful segment
        parts = [p.strip().strip('"\'') for p in prefix.split(',')]
        parts = [p for p in parts if p and not p.isdigit() and len(p) > 1]

        if not parts:
            return None

        # Last segment is usually the city
        city_candidate = parts[-1]

        # Remove leading street numbers if present (e.g. "100 Tech Blvd, Austin" → take "Austin")
        # But if the last segment starts with a digit, try the previous segment
        if city_candidate and city_candidate[0].isdigit():
            if len(parts) >= 2:
                city_candidate = parts[-2]

        # Clean up
        city_candidate = city_candidate.strip().strip(',.;')

        # Validate: should be at least 2 chars, not a number
        if len(city_candidate) >= 2 and not city_candidate.isdigit():
            return city_candidate

        return None