# ABOUTME: Functions for extracting ATP concentration from ChEMBL assay metadata
# ABOUTME: Supports both structured assay_parameters and free-text assay_description parsing

import re
from typing import Optional

import pandas as pd

from .api.downloader import check_and_download_chembl_db


# Unit conversion factors to µM
UNIT_CONVERSION_TO_UM = {
    "µM": 1.0,
    "uM": 1.0,
    "μM": 1.0,
    "mM": 1000.0,
    "nM": 0.001,
    "M": 1_000_000.0,
}


def convert_atp_units_to_um(value: float, unit: str) -> Optional[float]:
    """Convert ATP concentration from various units to µM.

    Args:
        value: The concentration value
        unit: The unit string (e.g., "mM", "nM", "µM")

    Returns:
        Concentration in µM, or None if unit is not recognized
    """
    if unit in UNIT_CONVERSION_TO_UM:
        return value * UNIT_CONVERSION_TO_UM[unit]
    return None


# Regex patterns for ATP concentration extraction from assay descriptions
# Pattern 1: "X µM/uM/mM/nM ATP" or "X µM/uM/mM/nM atp"
# Pattern 2: "ATP (X µM/uM/mM/nM)"
# Pattern 3: "at X µM/uM/mM/nM ATP"
# Pattern 4: "ATP = X µM/uM/mM/nM" or "ATP: X µM/uM/mM/nM"
# Pattern 5: "[ATP] = X µM/uM/mM/nM"
# Pattern 6: "[gamma-32P]ATP at X µM/uM/mM/nM"

# Unit pattern to match various representations of concentration units
UNIT_PATTERN = r"(µM|uM|μM|mM|nM|M)\b"

# ATP patterns
ATP_PATTERNS = [
    # "10 µM ATP", "10 uM ATP", "10 mM ATP"
    re.compile(
        rf"(\d+(?:\.\d+)?)\s*{UNIT_PATTERN}\s+(?:of\s+)?(?:\[?(?:gamma-32P|γ-32P)\]?)?ATP",
        re.IGNORECASE,
    ),
    # "ATP (2 mM)", "ATP (10 µM)"
    re.compile(
        rf"ATP\s*\(\s*(\d+(?:\.\d+)?)\s*{UNIT_PATTERN}\s*\)",
        re.IGNORECASE,
    ),
    # "ATP = 100 uM", "ATP: 100 uM"
    re.compile(
        rf"ATP\s*[=:]\s*(\d+(?:\.\d+)?)\s*{UNIT_PATTERN}",
        re.IGNORECASE,
    ),
    # "[ATP] = 10 µM"
    re.compile(
        rf"\[ATP\]\s*=\s*(\d+(?:\.\d+)?)\s*{UNIT_PATTERN}",
        re.IGNORECASE,
    ),
    # "[gamma-32P]ATP at 10 µM" or "gamma-32P-ATP at 10 µM"
    re.compile(
        rf"\[?(?:gamma-32P|γ-32P)\]?[- ]?ATP\s+(?:at\s+)?(\d+(?:\.\d+)?)\s*{UNIT_PATTERN}",
        re.IGNORECASE,
    ),
]


def parse_atp_from_description(description: Optional[str]) -> Optional[float]:
    """Extract ATP concentration from assay description text.

    Searches for common patterns like:
    - "10 µM ATP", "10 uM ATP", "10 mM ATP"
    - "ATP (2 mM)"
    - "at 2 uM ATP"
    - "ATP = 100 uM"
    - "[gamma-32P]ATP at 10 µM"

    Args:
        description: The assay description text

    Returns:
        ATP concentration in µM, or None if not found
    """
    if not description:
        return None

    for pattern in ATP_PATTERNS:
        match = pattern.search(description)
        if match:
            value = float(match.group(1))
            unit = match.group(2)
            return convert_atp_units_to_um(value, unit)

    return None


def get_atp_from_assay_parameters(
    assay_chembl_ids: list[str],
    prefix: Optional[list[str]] = None,
    version: Optional[str] = None,
) -> pd.DataFrame:
    """Query assay_parameters table for ATP concentration data.

    Searches for ATP-related entries in assay_parameters.text_value and
    extracts concentration values.

    Args:
        assay_chembl_ids: List of ChEMBL assay IDs to query
        prefix: Optional prefix for ChEMBL database location
        version: Optional ChEMBL version

    Returns:
        DataFrame with columns: assay_chembl_id, atp_concentration_um
    """
    from chembl_downloader import query

    configs = check_and_download_chembl_db(prefix=prefix, version=version)

    if not assay_chembl_ids:
        return pd.DataFrame(columns=["assay_chembl_id", "atp_concentration_um"])

    placeholders = ", ".join([f"'{aid}'" for aid in assay_chembl_ids])

    query_str = f"""
    SELECT
        a.chembl_id AS assay_chembl_id,
        ap.text_value,
        ap.standard_value,
        ap.standard_units
    FROM assays a
    JOIN assay_parameters ap ON a.assay_id = ap.assay_id
    WHERE a.chembl_id IN ({placeholders})
      AND (ap.text_value LIKE '%ATP%'
           OR ap.type LIKE '%ATP%'
           OR ap.comments LIKE '%ATP%')
    """

    result = query(
        query_str,
        version=configs["version"],
        prefix=configs["prefix"],
    )

    if result.empty:
        return pd.DataFrame(columns=["assay_chembl_id", "atp_concentration_um"])

    # Extract ATP concentration from text_value using regex
    atp_data = []
    for _, row in result.iterrows():
        text = row.get("text_value", "")
        conc = parse_atp_from_description(text) if text else None
        if conc is not None:
            atp_data.append(
                {"assay_chembl_id": row["assay_chembl_id"], "atp_concentration_um": conc}
            )

    if not atp_data:
        return pd.DataFrame(columns=["assay_chembl_id", "atp_concentration_um"])

    return pd.DataFrame(atp_data).drop_duplicates()


def get_atp_from_assay_descriptions(
    assay_chembl_ids: list[str],
    prefix: Optional[list[str]] = None,
    version: Optional[str] = None,
) -> pd.DataFrame:
    """Query assay table and extract ATP concentration from descriptions.

    Args:
        assay_chembl_ids: List of ChEMBL assay IDs to query
        prefix: Optional prefix for ChEMBL database location
        version: Optional ChEMBL version

    Returns:
        DataFrame with columns: assay_chembl_id, atp_concentration_um
    """
    from chembl_downloader import query

    configs = check_and_download_chembl_db(prefix=prefix, version=version)

    if not assay_chembl_ids:
        return pd.DataFrame(columns=["assay_chembl_id", "atp_concentration_um"])

    placeholders = ", ".join([f"'{aid}'" for aid in assay_chembl_ids])

    query_str = f"""
    SELECT
        chembl_id AS assay_chembl_id,
        description AS assay_description
    FROM assays
    WHERE chembl_id IN ({placeholders})
      AND description LIKE '%ATP%'
    """

    result = query(
        query_str,
        version=configs["version"],
        prefix=configs["prefix"],
    )

    if result.empty:
        return pd.DataFrame(columns=["assay_chembl_id", "atp_concentration_um"])

    # Extract ATP concentration from descriptions
    result["atp_concentration_um"] = result["assay_description"].apply(
        parse_atp_from_description
    )

    # Filter to rows with valid ATP concentration and return
    return (
        result[result["atp_concentration_um"].notna()][
            ["assay_chembl_id", "atp_concentration_um"]
        ]
        .drop_duplicates()
        .reset_index(drop=True)
    )


def get_atp_concentration_for_assays(
    assay_chembl_ids: list[str],
    prefix: Optional[list[str]] = None,
    version: Optional[str] = None,
) -> pd.DataFrame:
    """Get ATP concentration for assays, combining multiple sources.

    Priority order:
    1. assay_parameters table (structured data)
    2. assay_description field (regex extraction)

    Args:
        assay_chembl_ids: List of ChEMBL assay IDs
        prefix: Optional prefix for ChEMBL database location
        version: Optional ChEMBL version

    Returns:
        DataFrame with columns: assay_chembl_id, atp_concentration_um
    """
    # Get ATP from assay_parameters (structured data)
    atp_params = get_atp_from_assay_parameters(
        assay_chembl_ids, prefix=prefix, version=version
    )

    # Get ATP from assay descriptions (regex extraction)
    atp_desc = get_atp_from_assay_descriptions(
        assay_chembl_ids, prefix=prefix, version=version
    )

    if atp_params.empty and atp_desc.empty:
        return pd.DataFrame(columns=["assay_chembl_id", "atp_concentration_um"])

    # Combine: prefer assay_parameters data when available
    if atp_params.empty:
        return atp_desc
    if atp_desc.empty:
        return atp_params

    # Merge, keeping assay_parameters values where both exist
    combined = pd.merge(
        atp_desc,
        atp_params,
        on="assay_chembl_id",
        how="outer",
        suffixes=("_desc", "_param"),
    )

    # Use assay_parameters value if available, otherwise description value
    combined["atp_concentration_um"] = combined["atp_concentration_um_param"].fillna(
        combined["atp_concentration_um_desc"]
    )

    return combined[["assay_chembl_id", "atp_concentration_um"]].reset_index(drop=True)
