"""core lookup, search and annotation logic for icdlookup.

all public functions in this module are re-exported from ``icdlookup.__init__``.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

from icdlookup.loader import (
    build_index,
    load_data,
    normalise_code,
    _add_dot,
)


def lookup(
    code: str,
    *,
    detail: bool = False,
    reference: str | Path | None = None,
) -> str | dict[str, str] | None:
    """look up a single icd code and return its description.

    parameters
    ----------
    code : str
        an icd code such as ``"A00.1"`` or ``"A001"``.  lookup is
        case-insensitive and dot-insensitive.
    detail : bool, default False
        when *True* return the full hierarchy dict instead of just the
        description string.
    reference : str | Path | None
        optional path to a custom reference csv.

    returns
    -------
    str | dict[str, str] | None
        the description string, the full row dict if *detail=True*,
        or *None* if the code is not found.
    """
    index = build_index(reference)
    norm = normalise_code(code)
    row = index.get(norm)

    if row is None:
        return None

    if detail:
        return dict(row)  # return a copy

    return row.get("description")


def search(
    keyword: str,
    *,
    limit: int = 20,
    reference: str | Path | None = None,
) -> list[tuple[str, str]]:
    """search for icd codes whose description contains *keyword*.

    parameters
    ----------
    keyword : str
        case-insensitive search term.
    limit : int, default 20
        maximum number of results to return.
    reference : str | Path | None
        optional path to a custom reference csv.

    returns
    -------
    list[tuple[str, str]]
        a list of ``(icd10_code, description)`` tuples.
    """
    rows = load_data(reference)
    kw = keyword.strip().lower()
    results: list[tuple[str, str]] = []

    for row in rows:
        if kw in row.get("description", "").lower():
            results.append((row["icd10_code"], row["description"]))
            if len(results) >= limit:
                break

    return results


def annotate(
    df: Any,
    code_column: str,
    *,
    description_column: str = "icd_description",
    reference: str | Path | None = None,
) -> Any:
    """add an icd description column to a pandas dataframe.

    parameters
    ----------
    df : pandas.DataFrame
        the input dataframe.
    code_column : str
        name of the column containing icd codes.
    description_column : str, default "icd_description"
        name of the new column to create.
    reference : str | Path | None
        optional path to a custom reference csv.

    returns
    -------
    pandas.DataFrame
        the original dataframe with the new description column appended.

    raises
    ------
    ImportError
        if pandas is not installed.
    """
    try:
        import pandas as pd  # noqa: F811
    except ImportError:
        raise ImportError(
            "pandas is required for annotate(). "
            "install it with: pip install icdlookup[pandas]"
        ) from None

    index = build_index(reference)

    df = df.copy()
    df[description_column] = df[code_column].map(
        lambda c: _lookup_description(c, index)
    )
    return df


def _lookup_description(code: Any, index: dict[str, dict[str, str]]) -> str | None:
    """helper used by annotate() to map a single code."""
    if not isinstance(code, str):
        return None
    norm = normalise_code(code)
    row = index.get(norm)
    return row["description"] if row else None
