"""FHIR REST client with simple multiprocessing helpers."""

from __future__ import annotations

import logging
import multiprocessing as mp
import os
import time
from collections import OrderedDict
from dataclasses import dataclass
from typing import Any, Callable, Dict, Generator, Iterable, List, Optional, Sequence, Tuple, Union

import pandas as pd
import requests
from requests.auth import HTTPBasicAuth, _basic_auth_str
from requests.adapters import HTTPAdapter
from urllib.parse import urlparse
from urllib3.util.retry import Retry
from tqdm import tqdm

from .auth import Auth, TokenAuth
from .types import FHIRObj

ProcessFunc = Callable[[Dict[str, Any]], Dict[str, List[Dict[str, Any]]]]
logger = logging.getLogger(__name__)
_FHIRPATH_COMPILED: Optional[List[Tuple[str, Any]]] = None


class FHIRRequestError(RuntimeError):
    """Raised when a FHIR request fails or returns invalid JSON."""


@dataclass
class QueryJob:
    resource_type: str
    params: Dict[str, Any]
    max_pages: int


def _join_url(base: str, path: str) -> str:
    """Join a base URL and a resource path with a single slash."""
    if base.endswith("/"):
        base = base[:-1]
    if path.startswith("/"):
        path = path[1:]
    return f"{base}/{path}"


def _resolve_next_url(base_url: str, next_link: str) -> str:
    """Resolve a next-link URL against the server origin to avoid duplicate paths."""
    if next_link.startswith("http://") or next_link.startswith("https://"):
        return next_link
    parsed = urlparse(base_url)
    origin = f"{parsed.scheme}://{parsed.netloc}"
    if next_link.startswith("/"):
        return origin + next_link
    return origin + "/" + next_link


def _flatten_record(resource: Any, out: Dict[str, Any], prefix: str, sep: str) -> None:
    """Recursively flatten a resource into a flat dictionary."""
    if isinstance(resource, dict):
        for key, value in resource.items():
            _flatten_record(value, out, f"{prefix}{sep}{key}", sep)
    elif isinstance(resource, list):
        for idx, value in enumerate(resource):
            _flatten_record(value, out, f"{prefix}{sep}{idx}", sep)
    else:
        out[prefix.lstrip(sep)] = resource


def _freeze_params(params: Dict[str, Any]) -> Tuple[Tuple[str, Any], ...]:
    """Freeze params into a hashable, stable structure."""
    frozen: List[Tuple[str, Any]] = []
    for key, value in sorted(params.items()):
        if isinstance(value, dict):
            frozen.append((key, tuple(sorted(value.items()))))
        elif isinstance(value, (list, tuple, set)):
            frozen.append((key, tuple(value)))
        else:
            frozen.append((key, value))
    return tuple(frozen)


class ResponseCache:
    """Small in-memory TTL cache for GET responses."""

    def __init__(self, ttl_seconds: int = 300, max_entries: int = 1000) -> None:
        self.ttl_seconds = max(0, int(ttl_seconds))
        self.max_entries = max(0, int(max_entries))
        self._store: "OrderedDict[Tuple[Any, ...], Tuple[float, Tuple[Dict[str, Any], str, int, float]]]" = (
            OrderedDict()
        )

    def _prune(self) -> None:
        while self.max_entries and len(self._store) > self.max_entries:
            self._store.popitem(last=False)

    def get(self, key: Tuple[Any, ...], now: float) -> Optional[Tuple[Dict[str, Any], str, int, float]]:
        if self.ttl_seconds <= 0 or self.max_entries <= 0:
            return None
        entry = self._store.get(key)
        if entry is None:
            return None
        inserted_at, payload = entry
        if (now - inserted_at) > self.ttl_seconds:
            self._store.pop(key, None)
            return None
        self._store.move_to_end(key)
        return payload

    def set(self, key: Tuple[Any, ...], payload: Tuple[Dict[str, Any], str, int, float], now: float) -> None:
        if self.ttl_seconds <= 0 or self.max_entries <= 0:
            return
        self._store[key] = (now, payload)
        self._store.move_to_end(key)
        self._prune()


def flatten_data(
    bundle: Dict[str, Any] | FHIRObj, col_sep: str = "_"
) -> Dict[str, List[Dict[str, Any]]]:
    """Flatten bundle entries into row dictionaries by resourceType."""
    if isinstance(bundle, FHIRObj):
        bundle = bundle.to_dict()
    records: Dict[str, List[Dict[str, Any]]] = {}
    for entry in bundle.get("entry", []) or []:
        resource = entry.get("resource") or {}
        resource_type = resource.get("resourceType", "Resource")
        records.setdefault(resource_type, [])
        base: Dict[str, Any] = {}
        _flatten_record(resource, base, "", col_sep)
        records[resource_type].append({k: v for k, v in base.items() if v is not None})
    return records


def _bundle_entry_count(bundle: Dict[str, Any] | FHIRObj) -> int:
    entries = bundle.entry if isinstance(bundle, FHIRObj) else bundle.get("entry", [])
    return len(entries or [])


def _processed_row_count(processed: Iterable[Dict[str, List[Dict[str, Any]]]]) -> int:
    total = 0
    for bundle_output in processed:
        for rows in bundle_output.values():
            total += len(rows)
    return total


def _format_fancy_summary(
    bundles: int, entries: int, rows: int, mode: str, color: bool = True
) -> str:
    label = "FHIR Summary"
    if color:
        label = "\033[1;36mFHIR Summary\033[0m"
    line = "—" * 58
    if color:
        line = f"\033[2m{line}\033[0m"
    return (
        f"\n{line}\n"
        f"  {label}\n"
        f"  bundles: {bundles:<8}  entries: {entries:<8}  rows: {rows:<8}\n"
        f"  mode: {mode}\n"
        f"{line}"
    )


def _compile_fhir_paths(fhir_paths: Sequence[Union[str, Tuple[str, str]]]) -> List[Tuple[str, Any]]:
    """Compile FHIRPath expressions into callable functions."""
    try:
        import fhirpathpy  # type: ignore
    except Exception as exc:  # pragma: no cover - optional dependency
        raise ImportError(
            "FHIRPath support requires `fhirpathpy`. Install it to use fhir_paths."
        ) from exc
    compiled: List[Tuple[str, Any]] = []
    for path in fhir_paths:
        name, expr = (path, path) if isinstance(path, str) else path
        compiled.append((name, fhirpathpy.compile(path=expr)))
    return compiled


def _parse_fhir_path_worker(bundle: Dict[str, Any] | FHIRObj) -> Dict[str, List[Dict[str, Any]]]:
    """Worker-friendly wrapper for parsing FHIRPath using global compiled paths."""
    if _FHIRPATH_COMPILED is None:
        raise RuntimeError("FHIRPath worker not initialized.")
    return parse_fhir_path(bundle, _FHIRPATH_COMPILED)


def _init_fhirpath_worker(compiled: List[Tuple[str, Any]]) -> None:
    """Initializer to set compiled FHIRPath expressions in worker processes."""
    global _FHIRPATH_COMPILED
    _FHIRPATH_COMPILED = compiled


def parse_fhir_path(
    bundle: Dict[str, Any] | FHIRObj,
    compiled_fhir_paths: List[Tuple[str, Any]],
) -> Dict[str, List[Dict[str, Any]]]:
    """Extract fields from bundle entries using compiled FHIRPath expressions."""
    if isinstance(bundle, FHIRObj):
        bundle = bundle.to_dict()
    records: Dict[str, List[Dict[str, Any]]] = {}
    for entry in bundle.get("entry", []) or []:
        resource = entry.get("resource") or {}
        resource_type = resource.get("resourceType", "Resource")
        records.setdefault(resource_type, [])
        base: Dict[str, Any] = {}
        for name, compiled in compiled_fhir_paths:
            result = compiled(resource=resource)
            if isinstance(result, list):
                if len(result) == 0:
                    value = None
                elif len(result) == 1:
                    value = result[0]
                else:
                    value = result
            else:
                value = result
            if name not in base or base[name] is None:
                base[name] = value
        records[resource_type].append({k: v for k, v in base.items() if v is not None})
    return records


def _request_json(
    session: requests.Session,
    url: str,
    params: Dict[str, Any],
    timeout: int,
    log: Optional[logging.Logger],
    cache: Optional[ResponseCache] = None,
) -> Tuple[Dict[str, Any], str, int, float]:
    """GET a URL, validate response, parse JSON, and return payload plus metadata."""
    start = time.perf_counter()
    if cache is not None:
        cache_key = (
            url,
            _freeze_params(params),
            id(session.auth) if hasattr(session, "auth") else None,
            session.headers.get("Authorization"),
            session.headers.get("Accept"),
        )
        cached = cache.get(cache_key, start)
        if cached is not None:
            if log is not None:
                log.debug("FHIR GET cache hit url=%s params=%s", url, params)
            data, response_url, status_code, _ = cached
            elapsed_ms = (time.perf_counter() - start) * 1000.0
            return data, response_url, status_code, elapsed_ms
    try:
        response = session.get(url, params=params, timeout=timeout)
        response.raise_for_status()
    except requests.RequestException as exc:
        if log is not None:
            log.exception("FHIR request failed url=%s params=%s", url, params)
        raise FHIRRequestError(f"Request failed: {url}") from exc

    elapsed_ms = (time.perf_counter() - start) * 1000.0
    if log is not None:
        log.debug(
            "FHIR GET %s params=%s status=%s elapsed_ms=%.1f",
            url,
            params,
            response.status_code,
            elapsed_ms,
        )
    try:
        data = response.json()
    except ValueError as exc:
        if log is not None:
            log.exception("FHIR response JSON decode failed url=%s", response.url)
        raise FHIRRequestError(f"Invalid JSON response from {response.url}") from exc
    if cache is not None:
        cache.set(cache_key, (data, response.url, response.status_code, elapsed_ms), time.perf_counter())
    return data, response.url, response.status_code, elapsed_ms


def _worker_fetch_bundles(
    args: Tuple[
        str,
        QueryJob,
        Dict[str, str],
        int,
        Dict[str, Any],
        bool,
        str,
        Optional[str],
        Optional[str],
        Optional[str],
        int,
        int,
    ]
) -> List[Dict[str, Any]]:
    """Worker helper for parallel bundle fetching."""
    (
        base_url,
        job,
        headers,
        timeout,
        request_params,
        log_requests,
        auth_mode,
        auth_url,
        refresh_url,
        token_auth_method,
        cache_ttl_seconds,
        cache_max_entries,
    ) = args
    log = logger if log_requests else None
    session = requests.Session()
    cache = ResponseCache(ttl_seconds=cache_ttl_seconds, max_entries=cache_max_entries)
    if headers:
        session.headers.update(headers)
    if auth_mode == "basic":
        username = os.environ.get("FHIR_USER")
        password = os.environ.get("FHIR_PASSWORD")
        if username and password:
            session.auth = HTTPBasicAuth(username, password)
    elif auth_mode == "token_env":
        username = os.environ.get("FHIR_USER")
        password = os.environ.get("FHIR_PASSWORD")
        if username and password and auth_url:
            session.auth = TokenAuth(
                username=username,
                password=password,
                auth_url=auth_url,
                refresh_url=refresh_url,
                session=session,
                auth_method=token_auth_method or "basic",
            )
    bundles: List[Dict[str, Any]] = []
    request_url = _join_url(base_url, job.resource_type)
    params = {**request_params, **job.params}
    pages = 0
    while True:
        bundle, response_url, _, _ = _request_json(
            session=session,
            url=request_url,
            params=params,
            timeout=timeout,
            log=log,
            cache=cache,
        )
        if bundle.get("resourceType") != "Bundle":
            bundle = {
                "resourceType": "Bundle",
                "type": "read",
                "total": 1,
                "entry": [{"fullUrl": response_url, "resource": bundle}],
            }
        bundles.append(bundle)
        pages += 1
        if 0 < job.max_pages <= pages:
            break
        next_link = None
        for link in bundle.get("link", []) or []:
            if link.get("relation") == "next":
                next_link = link.get("url")
                break
        if not next_link:
            break
        request_url = _resolve_next_url(base_url, next_link)
        params = {}
    return bundles


class FHIRRestClient:
    """Simple FHIR REST client focused on bundle retrieval and DataFrame conversion.

    Parameters
    ----------
    base_url:
        Base FHIR server URL (e.g., https://example.org/fhir).
    auth_url:
        Token endpoint for password-based auth. If omitted, uses FHIR_AUTH_URL env.
    auth:
        Optional Auth or requests.Session instance to use.
    num_processes:
        Process count for multiprocessing (>=1).
    request_timeout:
        Per-request timeout in seconds.
    request_params:
        Default query params merged into every request.
    bearer_token:
        Token used for parallel requests (or set FHIR_TOKEN env).
    log_requests:
        If True, logs request timing and summaries.
    token_auth_method:
        Token auth method for password-based token endpoints ("basic" or "post").
    cache_ttl_seconds:
        In-memory cache TTL for GET responses.
    cache_max_entries:
        Maximum number of cached GET responses.
    """

    def __init__(
        self,
        base_url: Optional[str] = None,
        auth_url: Optional[str] = None,
        auth: Optional[Union[Auth, requests.Session]] = None,
        num_processes: int = 4,
        request_timeout: int = 30,
        request_params: Optional[Dict[str, Any]] = None,
        bearer_token: Optional[str] = None,
        log_requests: bool = False,
        extra_headers: Optional[Dict[str, str]] = None,
        retries: Optional[Retry] = None,
        return_fhir_obj: bool = False,
        token_auth_method: Optional[str] = None,
        cache_ttl_seconds: int = 300,
        cache_max_entries: int = 1000,
    ) -> None:
        """Initialize a FHIR REST client."""
        self.base_url = base_url or os.environ.get("FHIR_BASE_URL")
        if not self.base_url:
            raise ValueError("base_url is required (or set FHIR_BASE_URL).")
        if "://" not in self.base_url:
            # Default to https when scheme is missing (common config mistake).
            self.base_url = f"https://{self.base_url}"
        self.auth_url = auth_url or os.environ.get("FHIR_AUTH_URL")
        self.request_timeout = request_timeout
        self.request_params = request_params or {}
        self.num_processes = max(1, num_processes)
        self._close_session_on_exit = False
        self.log_requests = log_requests
        self.extra_headers = extra_headers or {}
        self.return_fhir_obj = return_fhir_obj
        self._token_auth_method = (token_auth_method or "basic").lower()
        self._cache_ttl_seconds = cache_ttl_seconds
        self._cache_max_entries = cache_max_entries
        self._response_cache = ResponseCache(
            ttl_seconds=self._cache_ttl_seconds,
            max_entries=self._cache_max_entries,
        )

        if isinstance(auth, Auth):
            self.session = auth.session
        elif isinstance(auth, requests.Session):
            self.session = auth
        else:
            self._close_session_on_exit = True
            if self.auth_url:
                self.session = Auth(
                    auth_url=self.auth_url,
                    refresh_url=os.environ.get("FHIR_REFRESH_URL"),
                    auth_method="env",
                    token_auth_method=self._token_auth_method,
                ).session
            else:
                self.session = requests.Session()

        if self.extra_headers:
            self.session.headers.update(self.extra_headers)

        if retries is None:
            retries = Retry(
                total=3,
                backoff_factor=0.5,
                status_forcelist=[502, 503, 504],
                allowed_methods=["GET"],
                raise_on_status=False,
            )
        adapter = HTTPAdapter(max_retries=retries)
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)

        self._bearer_token = bearer_token or os.environ.get("FHIR_TOKEN")
        logger.info(
            "Initialized FHIRRestClient base_url=%s processes=%s auth_url=%s",
            self.base_url,
            self.num_processes,
            "set" if self.auth_url else "none",
        )

    def __enter__(self) -> "FHIRRestClient":
        """Return self for context manager usage."""
        return self

    def close(self) -> None:
        """Close the underlying session if owned."""
        if self._close_session_on_exit:
            self.session.close()

    def __exit__(self, exc_type, exc, exc_tb) -> None:
        """Close the session on context manager exit."""
        self.close()

    def fetch_bundles(
        self,
        resource_type: str,
        params: Optional[Dict[str, Any]] = None,
        max_pages: int = -1,
        show_progress: bool = True,
    ) -> Generator[Dict[str, Any], None, int]:
        """Fetch bundles for a resource with pagination.

        Parameters
        ----------
        resource_type:
            FHIR resource type (e.g., Patient, Observation).
        params:
            Query parameters for the search.
        max_pages:
            Maximum number of pages to fetch (-1 for all).
        show_progress:
            Whether to show a tqdm progress bar.

        Yields
        ------
        dict
            FHIR Bundle JSON objects.
        """
        if max_pages == 0:
            return 0
        request_url = _join_url(self.base_url, resource_type)
        merged_params = {**self.request_params, **(params or {})}
        pages = 0
        total_requests = 0
        progress = tqdm(
            desc=f"Query ({resource_type})",
            disable=not show_progress,
            total=None if max_pages == -1 else max_pages,
        )
        try:
            while True:
                bundle, response_url, _, _ = _request_json(
                    session=self.session,
                    url=request_url,
                    params=merged_params,
                    timeout=self.request_timeout,
                    log=logger if self.log_requests else None,
                    cache=self._response_cache,
                )
                total_requests += 1
                if bundle.get("resourceType") != "Bundle":
                    bundle = {
                        "resourceType": "Bundle",
                        "type": "read",
                        "total": 1,
                        "entry": [{"fullUrl": response_url, "resource": bundle}],
                    }
                if self.return_fhir_obj:
                    bundle = FHIRObj(**bundle)
                yield bundle
                pages += 1
                progress.update(1)
                if max_pages != -1 and pages >= max_pages:
                    break
                next_link = None
                for link in bundle.get("link", []) or []:
                    if link.get("relation") == "next":
                        next_link = link.get("url")
                        break
                if not next_link:
                    break
                request_url = _resolve_next_url(self.base_url, next_link)
                merged_params = {}
        finally:
            progress.close()
        if self.log_requests:
            logger.info(
                "Completed fetch_bundles resource=%s pages=%s requests=%s",
                resource_type,
                pages,
                total_requests,
            )
        return pages

    def fetch_bundles_parallel(
        self,
        resource_type: str,
        param_list: Sequence[Dict[str, Any]],
        max_pages: int = -1,
        show_progress: bool = True,
    ) -> Generator[Dict[str, Any], None, int]:
        """Fetch bundles in parallel for multiple parameter sets.

        Parameters
        ----------
        resource_type:
            FHIR resource type.
        param_list:
            List of query param dicts; each runs as a separate job.
        max_pages:
            Maximum pages per job (-1 for all).
        show_progress:
            Whether to show a tqdm progress bar.
        """
        headers: Dict[str, str] = dict(self.extra_headers)
        auth_mode = "none"
        token: Optional[str] = None
        if hasattr(self.session, "auth"):
            auth = self.session.auth
            if isinstance(auth, HTTPBasicAuth):
                auth_mode = "basic"
            elif isinstance(auth, TokenAuth):
                # Prefer the live token from the auth object.
                token = getattr(auth, "token", None)
                if token:
                    auth_mode = "bearer"
                    headers["Authorization"] = f"Bearer {token}"
                else:
                    # Use token auth in worker (re-auth from env) instead of reusing parent token.
                    auth_mode = "token_env"
            else:
                # Use token auth in worker (re-auth from env) instead of reusing parent token.
                auth_mode = "token_env"
        # Fall back to explicit bearer token (param/env) only if we didn't use auth.
        if auth_mode == "none":
            token = self._bearer_token
            if token:
                auth_mode = "bearer"
                headers["Authorization"] = f"Bearer {token}"
        if auth_mode == "none" and not headers:
            raise ValueError(
                "Parallel fetch requires token or basic auth. "
                "Configure Auth(auth_type='basic'|'token') or pass bearer_token."
            )
        jobs = [QueryJob(resource_type, params, max_pages) for params in param_list]
        if self.log_requests:
            logger.info(
                "Starting parallel fetch resource=%s jobs=%s processes=%s",
                resource_type,
                len(jobs),
                self.num_processes,
            )
        progress = tqdm(
            total=len(jobs),
            desc=f"Query Parallel ({resource_type})",
            disable=not show_progress,
        )
        bundles_count = 0
        with mp.Pool(processes=self.num_processes) as pool:
            for bundles in pool.imap_unordered(
                _worker_fetch_bundles,
                [
                    (
                        self.base_url,
                        job,
                        headers,
                        self.request_timeout,
                        self.request_params,
                        self.log_requests,
                        auth_mode,
                        self.auth_url,
                        os.environ.get("FHIR_REFRESH_URL"),
                        self._token_auth_method,
                        self._cache_ttl_seconds,
                        self._cache_max_entries,
                    )
                    for job in jobs
                ],
            ):
                for bundle in bundles:
                    yield bundle
                    bundles_count += 1
                progress.update(1)
        progress.close()
        if self.log_requests:
            logger.info(
                "Completed parallel fetch resource=%s bundles=%s",
                resource_type,
                bundles_count,
            )
        return bundles_count

    def bundles_to_dataframe(
        self,
        bundles: Iterable[Dict[str, Any] | FHIRObj],
        process_function: Optional[ProcessFunc] = None,
        fhir_paths: Optional[Sequence[Union[str, Tuple[str, str]]]] = None,
        mode: str = "flatten",
        use_multiprocessing: bool = True,
    ) -> Union[pd.DataFrame, Dict[str, pd.DataFrame]]:
        """Convert bundles to DataFrame(s).

        Parameters
        ----------
        bundles:
            Iterable of FHIR Bundle JSON objects.
        process_function:
            Custom processing function. Required when mode="custom".
        fhir_paths:
            FHIRPath expressions to extract fields (requires fhirpathpy).
        mode:
            "flatten" (default) or "custom".
        use_multiprocessing:
            Whether to process bundles with multiprocessing.
        """
        bundles_list = list(bundles)
        if self.log_requests:
            logger.info(
                "Building dataframe mode=%s multiprocessing=%s fhir_paths=%s",
                mode,
                use_multiprocessing,
                bool(fhir_paths),
            )
        if fhir_paths is not None:
            compiled = _compile_fhir_paths(fhir_paths)
            # Store compiled paths globally for multiprocessing worker access.
            global _FHIRPATH_COMPILED
            _FHIRPATH_COMPILED = compiled
            process_function = _parse_fhir_path_worker
            mode = "custom"
        if mode == "flatten":
            process_function = flatten_data
        elif mode == "custom":
            if process_function is None:
                raise ValueError("custom mode requires a process_function.")
        else:
            raise ValueError("mode must be 'flatten' or 'custom'.")
        if use_multiprocessing and self.num_processes > 1:
            if process_function is _parse_fhir_path_worker and fhir_paths is not None:
                with mp.Pool(
                    processes=self.num_processes,
                    initializer=_init_fhirpath_worker,
                    initargs=(compiled,),
                ) as pool:
                    processed = list(pool.imap(process_function, bundles_list))
            else:
                with mp.Pool(processes=self.num_processes) as pool:
                    processed = list(pool.imap(process_function, bundles_list))
        else:
            processed = [process_function(bundle) for bundle in bundles_list]

        bundle_count = len(bundles_list)
        entry_count = sum(_bundle_entry_count(bundle) for bundle in bundles_list)
        row_count = _processed_row_count(processed)
        logger.info(
            _format_fancy_summary(
                bundles=bundle_count,
                entries=entry_count,
                rows=row_count,
                mode=mode,
            )
        )

        return _records_to_dataframe(processed)

    def query_to_dataframe(
        self,
        resource_type: str,
        params: Optional[Dict[str, Any]] = None,
        max_pages: int = -1,
        process_function: Optional[ProcessFunc] = None,
        fhir_paths: Optional[Sequence[Union[str, Tuple[str, str]]]] = None,
        mode: str = "flatten",
        use_multiprocessing: bool = True,
        parallel_params: Optional[Sequence[Dict[str, Any]]] = None,
    ) -> Union[pd.DataFrame, Dict[str, pd.DataFrame]]:
        """Run a query and return DataFrame(s).

        Parameters
        ----------
        resource_type:
            FHIR resource type.
        params:
            Query parameters for a single request flow.
        max_pages:
            Maximum pages to fetch (-1 for all).
        process_function:
            Custom processing function.
        fhir_paths:
            FHIRPath expressions to extract fields.
        mode:
            "flatten" (default) or "custom".
        use_multiprocessing:
            Whether to process bundles with multiprocessing.
        parallel_params:
            If provided, runs parallel jobs for each param dict.
        """
        if self.log_requests:
            logger.info(
                "Query to dataframe resource=%s mode=%s parallel=%s",
                resource_type,
                mode,
                parallel_params is not None,
            )
        if parallel_params is not None:
            bundles = list(
                self.fetch_bundles_parallel(
                    resource_type=resource_type,
                    param_list=parallel_params,
                    max_pages=max_pages,
                )
            )
        else:
            bundles = list(
                self.fetch_bundles(
                    resource_type=resource_type,
                    params=params,
                    max_pages=max_pages,
                )
            )
        return self.bundles_to_dataframe(
            bundles=bundles,
            process_function=process_function,
            fhir_paths=fhir_paths,
            mode=mode,
            use_multiprocessing=use_multiprocessing,
        )

    def query_from_dataframe(
        self,
        df: pd.DataFrame,
        resource_type: str,
        column_map: Dict[str, str],
        params: Optional[Dict[str, Any]] = None,
        max_pages: int = -1,
        process_function: Optional[ProcessFunc] = None,
        fhir_paths: Optional[Sequence[Union[str, Tuple[str, str]]]] = None,
        mode: str = "flatten",
        use_multiprocessing: bool = True,
        parallel_fetch: bool = False,
        param_prefix: Optional[Dict[str, str]] = None,
    ) -> Union[pd.DataFrame, Dict[str, pd.DataFrame]]:
        """Run one query per DataFrame row and return DataFrame(s).

        Parameters
        ----------
        df:
            Input DataFrame containing query values.
        resource_type:
            FHIR resource type.
        column_map:
            Mapping of FHIR param -> DataFrame column.
        params:
            Base query parameters added to each request.
        max_pages:
            Maximum pages per row (-1 for all).
        process_function:
            Custom processing function.
        fhir_paths:
            FHIRPath expressions to extract fields.
        mode:
            "flatten" (default) or "custom".
        use_multiprocessing:
            Whether to process bundles with multiprocessing.
        parallel_fetch:
            Whether to fetch bundles in parallel (uses active Auth token or basic auth).
        param_prefix:
            Optional mapping of FHIR param -> prefix to prepend to each value
            (e.g., {"subject": "Patient/"}).
        """
        if self.log_requests:
            logger.info(
                "Query from dataframe resource=%s rows=%s mode=%s",
                resource_type,
                len(df),
                mode,
            )
        if df.empty:
            raise ValueError("Input DataFrame is empty.")
        for _, col in column_map.items():
            if col not in df.columns:
                raise ValueError(f"Column '{col}' not found in DataFrame.")
            if df[col].isnull().any():
                raise ValueError(f"Column '{col}' contains null values.")

        base_params = params or {}
        prefixes = param_prefix or {}
        param_list = []
        for row in df.itertuples(index=False):
            row_params: Dict[str, Any] = {}
            for fhir_param, col in column_map.items():
                value = row[df.columns.get_loc(col)]
                prefix = prefixes.get(fhir_param, "")
                row_params[fhir_param] = f"{prefix}{value}" if prefix else value
            param_list.append({**base_params, **row_params})

        bundles: List[Dict[str, Any]] = []
        if parallel_fetch and self.num_processes > 1:
            bundles = list(
                self.fetch_bundles_parallel(
                    resource_type=resource_type,
                    param_list=param_list,
                    max_pages=max_pages,
                )
            )
        else:
            for row_params in param_list:
                bundles.extend(
                    list(
                        self.fetch_bundles(
                            resource_type=resource_type,
                            params=row_params,
                            max_pages=max_pages,
                            show_progress=False,
                        )
                    )
                )

        return self.bundles_to_dataframe(
            bundles=bundles,
            process_function=process_function,
            fhir_paths=fhir_paths,
            mode=mode,
            use_multiprocessing=use_multiprocessing,
        )

    def query_df(
        self,
        resource_type: str,
        params: Optional[Dict[str, Any]] = None,
        max_pages: int = -1,
        process_function: Optional[ProcessFunc] = None,
        fhir_paths: Optional[Sequence[Union[str, Tuple[str, str]]]] = None,
        mode: str = "flatten",
        use_multiprocessing: bool = True,
        parallel_params: Optional[Sequence[Dict[str, Any]]] = None,
    ) -> Union[pd.DataFrame, Dict[str, pd.DataFrame]]:
        """Short alias for query_to_dataframe."""
        return self.query_to_dataframe(
            resource_type=resource_type,
            params=params,
            max_pages=max_pages,
            process_function=process_function,
            fhir_paths=fhir_paths,
            mode=mode,
            use_multiprocessing=use_multiprocessing,
            parallel_params=parallel_params,
        )

    def query_df_from(
        self,
        df: pd.DataFrame,
        resource_type: str,
        column_map: Dict[str, str],
        params: Optional[Dict[str, Any]] = None,
        max_pages: int = -1,
        process_function: Optional[ProcessFunc] = None,
        fhir_paths: Optional[Sequence[Union[str, Tuple[str, str]]]] = None,
        mode: str = "flatten",
        use_multiprocessing: bool = True,
        parallel_fetch: bool = False,
        param_prefix: Optional[Dict[str, str]] = None,
    ) -> Union[pd.DataFrame, Dict[str, pd.DataFrame]]:
        """Short alias for query_from_dataframe."""
        return self.query_from_dataframe(
            df=df,
            resource_type=resource_type,
            column_map=column_map,
            params=params,
            max_pages=max_pages,
            process_function=process_function,
            fhir_paths=fhir_paths,
            mode=mode,
            use_multiprocessing=use_multiprocessing,
            parallel_fetch=parallel_fetch,
            param_prefix=param_prefix,
        )

    def bundles_to_df(
        self,
        bundles: Iterable[Dict[str, Any] | FHIRObj],
        process_function: Optional[ProcessFunc] = None,
        fhir_paths: Optional[Sequence[Union[str, Tuple[str, str]]]] = None,
        mode: str = "flatten",
        use_multiprocessing: bool = True,
    ) -> Union[pd.DataFrame, Dict[str, pd.DataFrame]]:
        """Short alias for bundles_to_dataframe."""
        return self.bundles_to_dataframe(
            bundles=bundles,
            process_function=process_function,
            fhir_paths=fhir_paths,
            mode=mode,
            use_multiprocessing=use_multiprocessing,
        )


def _records_to_dataframe(
    processed: Iterable[Dict[str, List[Dict[str, Any]]]]
) -> Union[pd.DataFrame, Dict[str, pd.DataFrame]]:
    """Convert processed bundle records into DataFrame(s)."""
    records: Dict[str, List[Dict[str, Any]]] = {}
    for bundle_output in processed:
        for resource_type, rows in bundle_output.items():
            records.setdefault(resource_type, [])
            records[resource_type].extend(rows)
    if not records:
        return pd.DataFrame()
    dfs = {
        resource_type: pd.DataFrame(rows).dropna(axis=1, how="all")
        for resource_type, rows in records.items()
    }
    return next(iter(dfs.values())) if len(dfs) == 1 else dfs
