app.py

import re
import requests
from bs4 import BeautifulSoup
import csv
import json
import yaml
from typing import List, Set
from dataclasses import dataclass, asdict
import logging
from tqdm import tqdm
import hashlib
import os
logging.basicConfig(
	level=logging.INFO,
	format='[%(asctime)s] %(levelname)s: %(message)s',
	datefmt='%H:%M:%S'
)


EMBASSY_URL = "https://travel.state.gov/content/travel/en/us-visas/visa-information-resources/list-of-posts.html"
CSV_FILE = "us_embassies_consulates.csv"
JSON_FILE = "us_embassies_consulates.json"
YAML_FILE = "us_embassies_consulates.yml"


@dataclass(frozen=True)
class Embassy:
	country: str
	city: str
	code: str
	url: str
	continent: str = ''
	full_name: str = ''
	address: str = ''
	telephone: str = ''
	fax: str = ''
	email: str = ''
	website: str = ''
	cancel_reschedule: str = ''
	google_map: str = ''

def get_continent_by_country(country: str) -> str:
	"""
	Return continent name for a given country (in English). If country is empty, use city name.
	"""
	mapping = {
		# Africa
		'algeria': 'Africa', 'egypt': 'Africa', 'nigeria': 'Africa', 'south africa': 'Africa', 'kenya': 'Africa',
		'ghana': 'Africa', 'morocco': 'Africa', 'ethiopia': 'Africa', 'senegal': 'Africa', 'tunisia': 'Africa',
		# Asia
		'china': 'Asia', 'india': 'Asia', 'japan': 'Asia', 'south korea': 'Asia', 'singapore': 'Asia',
		'mongolia': 'Asia', 'bangladesh': 'Asia', 'pakistan': 'Asia', 'thailand': 'Asia', 'vietnam': 'Asia',
		'indonesia': 'Asia', 'malaysia': 'Asia', 'philippines': 'Asia', 'nepal': 'Asia', 'sri lanka': 'Asia',
		# Europe
		'france': 'Europe', 'germany': 'Europe', 'italy': 'Europe', 'spain': 'Europe', 'united kingdom': 'Europe',
		'russia': 'Europe', 'poland': 'Europe', 'austria': 'Europe', 'belgium': 'Europe', 'netherlands': 'Europe',
		'sweden': 'Europe', 'switzerland': 'Europe', 'czech republic': 'Europe', 'hungary': 'Europe',
		# North America
		'united states': 'North America', 'canada': 'North America', 'mexico': 'North America', 'jamaica': 'North America',
		'bahamas': 'North America', 'barbados': 'North America', 'trinidad and tobago': 'North America',
		# South America
		'argentina': 'South America', 'brazil': 'South America', 'chile': 'South America', 'colombia': 'South America',
		'peru': 'South America', 'uruguay': 'South America', 'paraguay': 'South America',
		# Oceania
		'australia': 'Oceania', 'new zealand': 'Oceania', 'fiji': 'Oceania', 'papua new guinea': 'Oceania',
		# Middle East
		'turkey': 'Asia', 'israel': 'Asia', 'saudi arabia': 'Asia', 'united arab emirates': 'Asia', 'qatar': 'Asia',
	}
	country_clean = country.strip().lower()
	return mapping.get(country_clean, '')

def get_continent_by_country_or_city(country: str, city: str) -> str:
	"""
	Return continent by country, or by city if country is empty.
	"""
	cont = get_continent_by_country(country)
	if cont:
		return cont
	return get_continent_by_country(city)

def parse_embassy_name(name: str) -> tuple[str, str, str]:
	"""
	Parse the embassy name string into country, city, and code.
	Example: 'Albania - Tirana (TIA)' -> ('Albania', 'Tirana', 'TIA')
	"""
	match = re.match(r"^(.*?)\s*-\s*(.*?)\s*\(([^)]+)\)$", name)
	if match:
		country, city, code = match.groups()
		return country.strip(), city.strip(), code.strip()
	match = re.match(r"^(.*?)\s*\(([^)]+)\)$", name)
	if match:
		city_or_country, code = match.groups()
		return '', city_or_country.strip(), code.strip()
	return '', name.strip(), ''

def fetch_embassy_details(url: str) -> dict:
	"""
	Fetch and parse embassy detail page for extra information.
	Returns a dict with keys: full_name, address, telephone, fax, email, website, cancel_reschedule, google_map.
	"""
	try:
		html = fetch_html(url)
	except Exception:
		return {}
	soup = BeautifulSoup(html, 'html.parser')
	return {
		'full_name': extract_full_name(soup),
		'address': extract_address(soup),
		'telephone': extract_telephone(soup),
		'fax': extract_fax(soup),
		'email': extract_email(soup),
		'website': extract_website(soup),
		'cancel_reschedule': extract_cancel_reschedule(soup),
		'google_map': extract_google_map(soup),
	}

def extract_full_name(soup: BeautifulSoup) -> str:
	contact_title = soup.find('div', class_='contact-title')
	if contact_title:
		return contact_title.get_text(strip=True)
	h1 = soup.find('h1')
	if h1:
		return h1.get_text(strip=True)
	if soup.title:
		return soup.title.get_text(strip=True)
	return ''

def extract_address(soup: BeautifulSoup) -> str:
	contact = soup.find('div', class_='contactInformation')
	if not contact:
		return ''
	addr = contact.find('div', class_='info_address')
	if addr:
		return addr.get_text(separator=' ', strip=True)
	return ''

def extract_telephone(soup: BeautifulSoup) -> str:
	contact = soup.find('div', class_='contactInformation')
	if not contact:
		return ''
	tel = contact.find('div', string=lambda s: s and 'Telephone' in s)
	if tel and tel.find_next_sibling('div'):
		return tel.find_next_sibling('div').get_text(separator=' ', strip=True)
	for tel_box in contact.find_all('div', class_='contact-box'):
		desc = tel_box.find('div', class_='contact-desc')
		data = tel_box.find('div', class_='contact-data')
		if desc and data and 'telephone' in desc.get_text(strip=True).lower():
			return data.get_text(separator=' ', strip=True)
	return ''

def extract_fax(soup: BeautifulSoup) -> str:
	contact = soup.find('div', class_='contactInformation')
	if not contact:
		return ''
	fax = contact.find('div', string=lambda s: s and 'Fax' in s)
	if fax and fax.find_next_sibling('div'):
		return fax.find_next_sibling('div').get_text(strip=True)
	return ''

def extract_email(soup: BeautifulSoup) -> str:
	import re
	contact = soup.find('div', class_='contactInformation')
	if not contact:
		return ''
	email = contact.find('a', href=re.compile(r'^mailto:'))
	if email:
		return email.get_text(strip=True)
	return ''

def extract_website(soup: BeautifulSoup) -> str:
	import re
	contact = soup.find('div', class_='contactInformation')
	if not contact:
		return ''
	web = contact.find('a', href=re.compile(r'^https?://'))
	if web:
		return web.get('href')
	return ''

def extract_cancel_reschedule(soup: BeautifulSoup) -> str:
	import re
	contact = soup.find('div', class_='contactInformation')
	if not contact:
		return ''
	cancel = contact.find('div', string=lambda s: s and 'Cancel and Reschedule' in s)
	if cancel and cancel.find_next_sibling('div'):
		val = cancel.find_next_sibling('div')
		a = val.find('a', href=re.compile(r'^mailto:'))
		return a.get_text(strip=True) if a else val.get_text(strip=True)
	for box in contact.find_all('div', class_='contact-box'):
		desc = box.find('div', class_='contact-desc')
		data = box.find('div', class_='contact-data')
		if desc and data and 'cancel and reschedule' in desc.get_text(strip=True).lower():
			a = data.find('a', href=re.compile(r'^mailto:'))
			return a.get_text(strip=True) if a else data.get_text(strip=True)
	return ''

def extract_google_map(soup: BeautifulSoup) -> str:
	import re
	iframe = soup.find('iframe', src=re.compile(r'google.com/maps'))
	if iframe:
		return iframe.get('src')
	return ''

def fetch_html(url: str) -> str:
	"""
	Fetch HTML content from a URL, using .cache if available.
	"""
	cache_dir = os.path.join(os.path.dirname(__file__), '.cache')
	os.makedirs(cache_dir, exist_ok=True)
	url_hash = hashlib.md5(url.encode('utf-8')).hexdigest()
	cache_file = os.path.join(cache_dir, f'{url_hash}.html')
	if os.path.exists(cache_file) and os.path.getsize(cache_file) > 0:
		with open(cache_file, 'r', encoding='utf-8') as f:
			return f.read()
	response = requests.get(url, timeout=30)
	response.raise_for_status()
	html = response.text
	with open(cache_file, 'w', encoding='utf-8') as f:
		f.write(html)
	return html


def get_section_parent(soup: BeautifulSoup, anchor: str) -> BeautifulSoup | None:
	"""
	Return the <p> tag containing embassy links for a given anchor, or None.
	"""
	h2 = soup.find('a', {'id': anchor})
	if not h2:
		return None
	section = h2.find_parent('h2')
	if not section:
		return None
	return section.find_next_sibling('p')



def extract_embassy_links_from_section(p: BeautifulSoup | None) -> List[Embassy]:
	"""
	Extract all embassy links from a <p> section.
	"""
	results: List[Embassy] = []
	if not p:
		return results
	results.extend(extract_embassy_links_from_divs(p))
	results.extend(extract_embassy_links_from_direct_a(p))
	return results



def create_embassy_from_link(name: str, href: str) -> Embassy:
	country, city, code = parse_embassy_name(name)
	if not country and city.strip() != '':
		country = city
	url = normalize_url(href)
	continent = get_continent_by_country_or_city(country, city)
	return Embassy(country, city, code, url, continent=continent)

def extract_embassy_links_from_divs(p: BeautifulSoup) -> List[Embassy]:
	"""Extract embassy links from <div> blocks inside <p>."""
	links: List[Embassy] = []
	for div in p.find_all('div', recursive=False):
		for item in div.find_all('a', recursive=True):
			name = item.get_text(strip=True)
			href = item.get('href')
			links.append(create_embassy_from_link(name, href))
	return links



def extract_embassy_links_from_direct_a(p: BeautifulSoup) -> List[Embassy]:
	"""Extract embassy links from direct <a> children of <p>."""
	links: List[Embassy] = []
	for item in p.find_all('a', recursive=False):
		name = item.get_text(strip=True)
		href = item.get('href')
		links.append(create_embassy_from_link(name, href))
	return links

def normalize_url(href: str | None) -> str:
	"""
	Add prefix to relative URLs. If already absolute, return as is.
	"""
	if not href:
		return ''
	if href.startswith('http://') or href.startswith('https://'):
		return href
	base = 'https://travel.state.gov'
	if href.startswith('/'):
		return base + href
	return base + '/' + href

def parse_embassies(html: str) -> List[Embassy]:
	"""
	Parse embassy/consulate data from HTML.
	Returns a list of Embassy dataclass instances.
	"""
	soup = BeautifulSoup(html, 'html.parser')
	anchors = ["A_H", "I_R", "S_Z"]
	results: List[Embassy] = []
	logging.info("Extracting embassy links from main list page...")
	for anchor in anchors:
		p = get_section_parent(soup, anchor)
		results.extend(extract_embassy_links_from_section(p))
	embassies = deduplicate_embassies(results)
	logging.info(f"Found {len(embassies)} embassies/consulates. Fetching details for each...")
	enriched: List[Embassy] = []
	for embassy in tqdm(embassies, desc="Crawling embassy details", unit="embassy"):
		details = fetch_embassy_details(embassy.url)
		enriched.append(Embassy(
			country=embassy.country,
			city=embassy.city,
			code=embassy.code,
			url=embassy.url,
			continent=embassy.continent,
			full_name=details.get('full_name', ''),
			address=details.get('address', ''),
			telephone=details.get('telephone', ''),
			fax=details.get('fax', ''),
			email=details.get('email', ''),
			website=details.get('website', ''),
			cancel_reschedule=details.get('cancel_reschedule', ''),
			google_map=details.get('google_map', ''),
		))
	return enriched

def deduplicate_embassies(embassies: List[Embassy]) -> List[Embassy]:
	"""
	Remove duplicates and filter out navigation links.
	"""
	seen: Set[Embassy] = set()
	clean_results: List[Embassy] = []
	for entry in embassies:
		if (
			not entry.url
			or entry.url.startswith('#')
			or entry.url.endswith('#top')
			or 'back to top' in entry.city.lower()
			or 'back to top' in entry.country.lower()
			or 'back to top' in entry.code.lower()
			or entry in seen
		):
			continue
		seen.add(entry)
		clean_results.append(entry)
	return clean_results

def save_csv(data: List[Embassy], filename: str) -> None:
	"""Save embassy/consulate data to a CSV file."""
	try:
		write_to_file(filename, serialize_csv, data)
	except Exception as e:
		print(f"Error writing CSV: {e}")

def save_json(data: List[Embassy], filename: str) -> None:
	"""Save embassy/consulate data to a JSON file."""
	try:
		write_to_file(filename, serialize_json, data)
	except Exception as e:
		print(f"Error writing JSON: {e}")

def save_yaml(data: List[Embassy], filename: str) -> None:
	"""Save embassy/consulate data to a YAML file."""
	try:
		write_to_file(filename, serialize_yaml, data)
	except Exception as e:
		print(f"Error writing YAML: {e}")


def serialize_csv(f, data: List[Embassy]) -> None:
	fieldnames = ['country', 'city', 'code', 'continent', 'url', 'full_name', 'address', 'telephone', 'fax', 'email', 'website', 'cancel_reschedule', 'google_map']
	writer = csv.DictWriter(f, fieldnames=fieldnames)
	writer.writeheader()
	for embassy in data:
		writer.writerow(asdict(embassy))

def serialize_json(f, data: List[Embassy]) -> None:
	json.dump([asdict(e) for e in data], f, ensure_ascii=False, indent=2)

def serialize_yaml(f, data: List[Embassy]) -> None:
	yaml.dump([asdict(e) for e in data], f, allow_unicode=True, sort_keys=False)

def write_to_file(filename: str, serializer, data: List[Embassy]) -> None:
	mode = 'w'
	newline = '' if serializer == serialize_csv else None
	encoding = 'utf-8'
	with open(filename, mode, newline=newline, encoding=encoding) as f:
		serializer(f, data)

def export_all_formats(data: List[Embassy]) -> None:
	"""Export embassy/consulate data to CSV, JSON, and YAML files."""
	save_csv(data, CSV_FILE)
	save_json(data, JSON_FILE)
	save_yaml(data, YAML_FILE)

def main() -> None:
	"""Main execution function."""
	try:
		logging.info("Fetching main embassy list page...")
		html = fetch_html(EMBASSY_URL)
		logging.info("Parsing embassies and crawling details...")
		embassies = parse_embassies(html)
		logging.info("Exporting data to CSV, JSON, and YAML...")
		export_all_formats(embassies)
		logging.info(f"Extracted {len(embassies)} embassies/consulates.")
		logging.info(f"Files generated: {CSV_FILE}, {JSON_FILE}, {YAML_FILE}")
	except Exception as e:
		logging.error(f"Error: {e}")

if __name__ == "__main__":
	main()