"""
Named Entity Recognition (NER) — rule-based + pattern enhanced.
Lebih presisi dari regex sederhana: gunakan gazetteer + context patterns.
"""
from typing import List, Dict
import re

# Gazetteer Indonesia (expandable)
PERSON_TITLES = [
    "presiden", "menteri", "gubernur", "bupati", "walikota", "calon",
    "ketua", "wakil", "direktur", "komisaris", "jenderal", "kolonel",
    "mayor", "kapten", "prof", "dr", "ir", "haji", "ustaz", "kyai",
]

ORG_KEYWORDS = [
    "kementerian", "badan", "dewan", "komisi", "partai", "pt", "tbk",
    "universitas", "institut", "polri", "tni", "bpk", "kpk", "ojk",
    "bi", "bps", "bmkg", "bnpb", "baznas", "mui", "nu", "muhammadiyah",
    "perserikatan", "organisasi", "perusahaan", "bank", "asosiasi",
]

LOCATION_KEYWORDS = [
    "jakarta", "surabaya", "bandung", "medan", "semarang", "makassar",
    "yogyakarta", "denpasar", "palembang", "manado", "padang", "solo",
    "indonesia", "jawa", "sumatera", "kalimantan", "sulawesi", "papua",
    "bali", "ntt", "ntb", "aceh", "riau", "lampung", "maluku",
    "provinsi", "kabupaten", "kota", "desa", "kecamatan",
]


def _is_capitalized_phrase(text: str, start: int, end: int) -> bool:
    """Cek apakah span memiliki kata yang diawali huruf besar."""
    span = text[start:end]
    words = span.split()
    return any(w[0].isupper() for w in words if w)


def extract_entities(text: str) -> Dict[str, List[str]]:
    """Extract persons, organizations, locations dari text."""
    persons = set()
    organizations = set()
    locations = set()

    text_lower = text.lower()
    words = text.split()

    # Pattern: Title + Capitalized Name (person)
    for i, word in enumerate(words):
        word_lower = word.lower().strip(".,;:!?\"'()")
        if word_lower in PERSON_TITLES and i + 1 < len(words):
            # Ambil 1-3 kata setelah title sebagai nama
            name_parts = []
            for j in range(i + 1, min(i + 4, len(words))):
                w = words[j].strip(".,;:!?\"'()")
                if w and w[0].isupper():
                    name_parts.append(w)
                else:
                    break
            if name_parts:
                persons.add(" ".join(name_parts))

    # Pattern: Capitalized consecutive words (potential names/orgs)
    cap_pattern = re.finditer(r'\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)+)\b', text)
    for match in cap_pattern:
        phrase = match.group(1)
        phrase_lower = phrase.lower()
        # Classify based on context
        if any(kw in phrase_lower for kw in ORG_KEYWORDS):
            organizations.add(phrase)
        elif any(kw in phrase_lower for kw in LOCATION_KEYWORDS):
            locations.add(phrase)
        elif len(phrase.split()) <= 3:
            persons.add(phrase)

    # Direct keyword matching for organizations
    for kw in ORG_KEYWORDS:
        pattern = re.finditer(rf'\b{re.escape(kw)}\s+([A-Z][a-zA-Z\s]{{2,30}})', text, re.IGNORECASE)
        for m in pattern:
            organizations.add(m.group(0).strip())

    # Location extraction
    for kw in LOCATION_KEYWORDS:
        if kw in text_lower:
            locations.add(kw.title())

    return {
        "persons": list(persons)[:20],
        "organizations": list(organizations)[:15],
        "locations": list(locations)[:15],
    }


def extract_batch(items: List) -> List[Dict]:
    results = []
    for item in items:
        entities = extract_entities(item.text)
        results.append({"id": item.id, "entities": entities})
    return results