"""
One-time script to extract company legal suffix data from cleanco and iso-20275.

Requires: pip install cleanco iso-20275

Outputs company_suffixes.json with suffix → {type, jurisdictions} mappings.
"""

import json
import logging
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(message)s")
log = logging.getLogger(__name__)


def load_cleanco_suffixes() -> dict[str, dict]:
    """Extract suffix data from cleanco's terms_by_type and terms_by_country."""
    from cleanco.clean import terms_by_type, terms_by_country

    suffixes: dict[str, dict] = {}

    # terms_by_type: {"Corporation": ["inc", "inc.", "incorporated", ...], ...}
    for entity_type, terms in terms_by_type.items():
        for term in terms:
            key = term.lower().replace(".", "").strip()
            if not key:
                continue
            if key not in suffixes:
                suffixes[key] = {"type": entity_type, "jurisdictions": []}
            # Don't overwrite type if already set — first wins

    # terms_by_country: {"United States": ["inc", "llc", ...], ...}
    for country, terms in terms_by_country.items():
        for term in terms:
            key = term.lower().replace(".", "").strip()
            if not key:
                continue
            if key in suffixes:
                if country not in suffixes[key]["jurisdictions"]:
                    suffixes[key]["jurisdictions"].append(country)
            else:
                suffixes[key] = {"type": "", "jurisdictions": [country]}

    log.info(f"  cleanco: {len(suffixes)} suffixes")
    return suffixes


def load_gleif_abbreviations() -> dict[str, dict]:
    """Extract entity legal form abbreviations from iso-20275."""
    import iso20275

    suffixes: dict[str, dict] = {}
    paths = iso20275.get_csv_paths()
    entries = iso20275.read_entries_from_csvpath(paths[-1])  # latest version

    for elf_list in entries.values():
        for elf in elf_list:
            country = elf.country or ""
            # Collect abbreviations from both local and transliterated
            raw_abbrevs = []
            if elf.local_abbreviations:
                raw_abbrevs.extend(elf.local_abbreviations.split(";"))
            if elf.transliterated_abbreviations:
                raw_abbrevs.extend(elf.transliterated_abbreviations.split(";"))

            for abbrev in raw_abbrevs:
                key = abbrev.lower().replace(".", "").strip()
                if not key or len(key) < 2:
                    continue
                if key not in suffixes:
                    suffixes[key] = {"type": "", "jurisdictions": []}
                if country and country not in suffixes[key]["jurisdictions"]:
                    suffixes[key]["jurisdictions"].append(country)

    log.info(f"  GLEIF ELF: {len(suffixes)} abbreviations")
    return suffixes


def merge_suffixes(cleanco: dict, gleif: dict) -> dict:
    """Merge cleanco (primary) with GLEIF (supplementary)."""
    merged = dict(cleanco)

    for key, gleif_entry in gleif.items():
        if key in merged:
            # Add any new jurisdictions from GLEIF
            for j in gleif_entry["jurisdictions"]:
                if j not in merged[key]["jurisdictions"]:
                    merged[key]["jurisdictions"].append(j)
        else:
            merged[key] = gleif_entry

    return merged


# Multi-word suffixes that should be matched as a unit
MULTI_WORD_SUFFIXES = {
    "pty ltd": {"type": "Limited", "jurisdictions": ["Australia", "South Africa"]},
    "pvt ltd": {"type": "Limited", "jurisdictions": ["India", "Pakistan", "Sri Lanka"]},
    "co ltd": {"type": "Limited", "jurisdictions": ["Japan", "Thailand", "Taiwan"]},
    "co inc": {"type": "Corporation", "jurisdictions": ["United States"]},
    "and co": {"type": "Partnership", "jurisdictions": []},
    "& co": {"type": "Partnership", "jurisdictions": []},
}


def main() -> None:
    log.info("Loading cleanco suffixes...")
    cleanco = load_cleanco_suffixes()

    log.info("Loading GLEIF ELF abbreviations...")
    try:
        gleif = load_gleif_abbreviations()
    except Exception as e:
        log.warning(f"  Could not load GLEIF data: {e}")
        gleif = {}

    log.info("Merging...")
    merged = merge_suffixes(cleanco, gleif)

    # Add/override multi-word suffixes with correct types
    for key, entry in MULTI_WORD_SUFFIXES.items():
        if key not in merged:
            merged[key] = entry
        else:
            # Fill in type if missing
            if not merged[key]["type"] and entry["type"]:
                merged[key]["type"] = entry["type"]
            # Add jurisdictions
            for j in entry["jurisdictions"]:
                if j not in merged[key]["jurisdictions"]:
                    merged[key]["jurisdictions"].append(j)

    # Sort by key
    result = dict(sorted(merged.items()))

    log.info(f"Total: {len(result)} suffixes")

    # Spot-check
    for key in ("inc", "ltd", "gmbh", "sa", "plc", "ag", "llc", "nv", "pty ltd"):
        entry = result.get(key, {})
        log.info(f"  {key}: type={entry.get('type', '?')}, jurisdictions={entry.get('jurisdictions', [])[:3]}")

    # Write output
    out_path = Path(__file__).parent.parent / "data" / "company_suffixes.json"
    out_path.write_text(json.dumps(result, indent=2, ensure_ascii=False) + "\n")
    log.info(f"Written to {out_path}")


if __name__ == "__main__":
    main()
