#!/usr/bin/env python3
from __future__ import annotations

import re
import unicodedata
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Literal, Optional, Tuple

import joblib

RecordType = Literal["PERSON", "JOINT_PERSON", "ORG", "TRUST", "UNKNOWN"]
PlanLabel = Literal[
    "PERSON_COMMA",
    "PERSON_SPACE",
    "JOINT",
    "ORG",
    "TRUST",
    "UNKNOWN",
]

SUFFIXES = {"JR", "SR", "II", "III", "IV", "V", "MD", "PHD", "ESQ", "CPA"}
PARTICLES = {
    "DE",
    "DEL",
    "DELA",
    "DELLA",
    "DI",
    "LA",
    "LE",
    "VAN",
    "VON",
    "ST",
    "SAINT",
}
JOINT_SPLIT_REGEX = re.compile(r"\s(?:&|AND)\s")


@dataclass
class ParsedPerson:
    full_name: str
    first_name: Optional[str]
    middle_name: Optional[str]
    last_name: Optional[str]
    surname: Optional[str]
    suffix: Optional[str]

    def as_dict(self) -> Dict[str, Optional[str]]:
        return {
            "full_name": self.full_name,
            "first_name": self.first_name,
            "middle_name": self.middle_name,
            "last_name": self.last_name,
            "surname": self.surname,
            "suffix": self.suffix,
        }


def normalize_name(input_name: str) -> str:
    return (
        unicodedata.normalize("NFKC", input_name)
        .replace("“", '"')
        .replace("”", '"')
        .replace("‘", "'")
        .replace("’", "'")
        .strip()
        .upper()
    )


def clean_token(token: str) -> str:
    return re.sub(r"^,+|,+$", "", token.strip())


def split_tokens(text: str) -> List[str]:
    raw = (text or "").upper()
    return [
        tok for tok in (clean_token(t) for t in re.split(r"[^A-Z0-9\\.]+", raw)) if tok
    ]


def suffix_of(token: str) -> Optional[str]:
    norm = token.replace(".", "").replace(",", "").upper().strip()
    if norm in SUFFIXES:
        return token.strip()
    return None


def parse_person_comma(name: str) -> Optional[ParsedPerson]:
    if "," not in name:
        return None
    left, right = name.split(",", 1)
    left_tokens = split_tokens(left)
    right_tokens = split_tokens(right)
    if not left_tokens or not right_tokens:
        return None

    suffix = None
    left_last = left_tokens[-1]
    left_suffix = suffix_of(left_last)
    if left_suffix:
        suffix = left_suffix
        left_tokens = left_tokens[:-1]
    if not left_tokens:
        return None

    right_last = right_tokens[-1]
    right_suffix = suffix_of(right_last)
    if right_suffix:
        suffix = right_suffix
        right_tokens = right_tokens[:-1]
    if not right_tokens:
        return None

    first_name = right_tokens[0]
    middle_name = " ".join(right_tokens[1:]) if len(right_tokens) > 1 else None
    last_name = " ".join(left_tokens)
    full_parts = [*right_tokens, *left_tokens]
    if suffix:
        full_parts.append(suffix)

    return ParsedPerson(
        full_name=" ".join(full_parts),
        first_name=first_name,
        middle_name=middle_name,
        last_name=last_name,
        surname=last_name,
        suffix=suffix,
    )


def parse_person_space(name: str) -> Optional[ParsedPerson]:
    tokens = split_tokens(name)
    if not tokens:
        return None
    suffix = None
    maybe_suffix = suffix_of(tokens[-1])
    if maybe_suffix:
        suffix = maybe_suffix
        tokens = tokens[:-1]
    if not tokens:
        return None

    if len(tokens) == 1:
        token = tokens[0]
        full_parts = [token]
        if suffix:
            full_parts.append(suffix)
        return ParsedPerson(
            full_name=" ".join(full_parts),
            first_name=None,
            middle_name=None,
            last_name=token,
            surname=token,
            suffix=suffix,
        )

    span = 1
    if len(tokens) >= 3 and tokens[1].upper() in PARTICLES:
        span = 2
        if len(tokens) >= 4 and tokens[2].upper() in PARTICLES:
            span = 3

    last_tokens = tokens[:span]
    given_tokens = tokens[span:]
    if not given_tokens and len(tokens) >= 2:
        last_tokens = tokens[:1]
        given_tokens = tokens[1:]
    if not given_tokens:
        return None

    first_name = given_tokens[0]
    middle_name = " ".join(given_tokens[1:]) if len(given_tokens) > 1 else None
    last_name = " ".join(last_tokens)
    full_parts = [*given_tokens, *last_tokens]
    if suffix:
        full_parts.append(suffix)

    return ParsedPerson(
        full_name=" ".join(full_parts),
        first_name=first_name,
        middle_name=middle_name,
        last_name=last_name,
        surname=last_name,
        suffix=suffix,
    )


def parse_joint(name: str) -> List[ParsedPerson]:
    parts = [p.strip() for p in JOINT_SPLIT_REGEX.split(name) if p.strip()]
    if len(parts) <= 1:
        single = parse_person_comma(name) if "," in name else parse_person_space(name)
        return [single] if single else []

    parsed = [
        parse_person_comma(part) if "," in part else parse_person_space(part)
        for part in parts
    ]

    if (
        len(parts) == 2
        and parsed[0] is None
        and parsed[1] is not None
        and len(split_tokens(parts[0])) == 1
        and parsed[1].last_name
    ):
        patched_name = f"{parts[0]} {parsed[1].last_name}"
        patched = (
            parse_person_comma(patched_name)
            if "," in patched_name
            else parse_person_space(patched_name)
        )
        if patched:
            parsed[0] = patched

    return [person for person in parsed if person is not None]


def canonical_from_people(people: List[ParsedPerson]) -> Optional[str]:
    if not people:
        return None
    return " | ".join(person.full_name for person in people)


def plan_to_record_type(plan: PlanLabel) -> RecordType:
    if plan == "PERSON_COMMA" or plan == "PERSON_SPACE":
        return "PERSON"
    if plan == "JOINT":
        return "JOINT_PERSON"
    return plan  # type: ignore[return-value]


def predict_parse(
    plan: PlanLabel, normalized_name: str
) -> Tuple[RecordType, Optional[str], List[ParsedPerson]]:
    if plan == "PERSON_COMMA":
        person = parse_person_comma(normalized_name)
        people = [person] if person else []
        return "PERSON", canonical_from_people(people), people
    if plan == "PERSON_SPACE":
        person = parse_person_space(normalized_name)
        people = [person] if person else []
        return "PERSON", canonical_from_people(people), people
    if plan == "JOINT":
        people = parse_joint(normalized_name)
        return "JOINT_PERSON", canonical_from_people(people), people
    return plan_to_record_type(plan), None, []


def load_model(path: str | Path = "model.joblib") -> Any:
    return joblib.load(path)


def parse_name(model: Any, raw_name: str) -> Dict[str, Any]:
    normalized_name = normalize_name(raw_name)
    pred_plan = model.predict([normalized_name])[0]
    pred_rt, pred_canonical, pred_people = predict_parse(pred_plan, normalized_name)

    return {
        "raw_name": raw_name,
        "normalized_name": normalized_name,
        "pred_plan": pred_plan,
        "record_type": pred_rt,
        "canonical_full_name": pred_canonical,
        "parsed_persons": [p.as_dict() for p in pred_people],
    }