"""API client for the A-Alpha Bio Data Portal.

The `aablocks` package provides programmatic access to the A-Alpha Bio Atlas
for browsing and downloading AlphaSeq data blocks.

## Installation

```bash
pip install aablocks
```

For DataFrame support:

```bash
pip install aablocks pandas   # For pandas
pip install aablocks polars   # For polars
```

## Quick Start

```python
import aablocks as aa

# Login (opens browser)
aa.login()

# List data blocks
blocks = aa.list_blocks()

# Get data as pandas DataFrame
df = aa.get_data("ab1001")
```

## Complete Example

```python
import aablocks as aa

# Authenticate
aa.login()

# Browse data blocks
blocks = aa.list_blocks(format="list")
print(f"Found {len(blocks)} data blocks")

for d in blocks[:3]:
    print(f"{d.id}: {d.name} (v{d.version})")

# Get details
details = aa.get_details("ab1001")
print(f"Modes: {details.modes}")

# Download data
df = aa.get_data("ab1001")
print(df.head())

# ML version
df_ml = aa.get_data("ab1001", mode="ml")

# Read datacard
datacard = aa.get_datacard("ab1001")
print(datacard)
```
"""

from __future__ import annotations

import csv
import gzip
import io
import sys
import time
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import PurePosixPath
from typing import TYPE_CHECKING, Any
from urllib.parse import urlparse

import requests
from tqdm import tqdm

from .auth import get_token_header
from .config import get_api_format, get_base_url, get_progress
from .logging import log_request
from .notices import show_startup_notices

# Import for type hints only - actual imports are done lazily in functions
# to avoid slow import times and allow use without these dependencies
if TYPE_CHECKING:
    import pandas as pd
    import polars as pl


@dataclass
class DataBlock:
    """Data block metadata."""

    id: str
    name: str
    experiment: str
    details: str
    modes: list[str]
    version: str
    locked: bool
    url: str | None
    structure_count: int | None
    a_size: int | None
    alpha_size: int | None
    total_ppi_count: int | None
    unique_ppi_count: int | None
    # Product the block belongs to.
    product: str | None = None
    product_display_name: str | None = None
    product_kind: str | None = None
    # Named product release this block belongs to (e.g. "VHH Q1 2026").
    source: str | None = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> DataBlock:
        """Create a DataBlock from API response dict."""
        return cls(
            id=data["id"],
            name=data["name"],
            experiment=data["experiment"],
            details=data["details"],
            modes=[m["name"] if isinstance(m, dict) else m for m in data["modes"]],
            version=data["version"],
            locked=data.get("locked", False),
            url=data.get("url"),
            structure_count=data.get("structure_count"),
            a_size=data.get("a_size"),
            alpha_size=data.get("alpha_size"),
            total_ppi_count=data.get("total_ppi_count"),
            unique_ppi_count=data.get("unique_ppi_count"),
            product=data.get("product"),
            product_display_name=data.get("product_display_name"),
            product_kind=data.get("product_kind"),
            source=data.get("source"),
        )


def list_blocks(
    all_versions: bool = False,
    format: str | None = None,
) -> list[DataBlock] | str | pd.DataFrame | pl.DataFrame | None:
    """List all data blocks accessible to the current user.

    Args:
        all_versions: If True, return all versions of each data block.
                     If False (default), return only the latest version.
        format: Output format - "csv", "list" (DataBlock objects), "pandas" (DataFrame), or "polars" (DataFrame).
                If None, uses the default from ~/.aab/config.json (defaults to "pandas").

    Returns:
        CSV string, list of DataBlock objects, pandas DataFrame, or polars DataFrame depending on format.
        Returns None and prints an error to stderr if format="pandas"/"polars" is requested but the library is not installed.

    Raises:
        ValueError: If format is not one of "csv", "list", "pandas", or "polars".
        requests.HTTPError: If the API request fails.

    Examples:
        >>> blocks = list_blocks()
        >>> for d in blocks:
        ...     print(f"{d.id}: {d.name}")

        >>> df = list_blocks(format="pandas")
        >>> df[["id", "name", "version"]]

        >>> df = list_blocks(all_versions=True, format="polars")
    """
    format_explicit = format is not None
    if format is None:
        format = get_api_format()

    if format not in ("csv", "list", "pandas", "polars"):
        raise ValueError(
            f"format must be 'csv', 'list', 'pandas', or 'polars', got '{format}'"
        )

    params = {"all_versions": "true"} if all_versions else {}
    resp = _api_get("/datasets", params=params or None, timeout=30)

    data = resp.json()
    objects = data.get("objects", [])

    if format == "csv":
        output = io.StringIO()
        if objects:
            writer = csv.DictWriter(output, fieldnames=objects[0].keys())
            writer.writeheader()
            writer.writerows(objects)
        return output.getvalue()

    list_blocks_alts = (
        "or call list_blocks(format='list') for DataBlock objects or "
        "list_blocks(format='csv') for a CSV string."
    )
    if format == "pandas":
        try:
            import pandas as pd
        except ImportError:
            _print_missing_package_error(
                "list_blocks", "pandas", list_blocks_alts, format_explicit
            )
            return None
        return pd.DataFrame(objects)

    if format == "polars":
        try:
            import polars
        except ImportError:
            _print_missing_package_error(
                "list_blocks", "polars", list_blocks_alts, format_explicit
            )
            return None
        return polars.DataFrame(objects)

    return [DataBlock.from_dict(obj) for obj in objects]


def get_details(block_id: str, version: str | None = None) -> DataBlock:
    """Get metadata for a specific data block.

    Args:
        block_id: The data block identifier (e.g., "ab1001").
        version: Optional version number. If not specified, returns latest.

    Returns:
        DataBlock object with metadata.

    Raises:
        ValueError: If the data block is not found.
        requests.HTTPError: If the API request fails for another reason.

    Examples:
        >>> block = get_details("ab1001")
        >>> print(block.name)
        >>> print(block.modes)  # ['source', 'ml']
        >>> print(block.product_display_name)

        >>> block = get_details("ab1001", version="1")
    """
    params: dict[str, str] = {}
    if version:
        params["version"] = version

    resp = _api_get(
        f"/datasets/{block_id}",
        params=params or None,
        not_found_msg=_block_not_found_msg(block_id),
        timeout=30,
    )
    data = resp.json()
    return DataBlock.from_dict(data["dataset"])


def get_datacard(block_id: str, version: str | None = None) -> dict[str, Any]:
    """Get the datacard content for a data block.

    Args:
        block_id: The data block identifier (e.g., "ab1001").
        version: Optional version number. If not specified, returns latest.

    Returns:
        Datacard content as a dict matching the OpenAPI `DatacardResponse`
        schema. Top-level keys group fields by detail-page tab: `exec`
        (scientific_value, use_cases, differentiations, stat_subs), `ml`
        (schema_columns, noise_signals, recommended_splits, data_views), and
        `bio` (biological_system, figures, findings, considerations).

    Raises:
        ValueError: If the data block is not found.
        requests.HTTPError: If the API request fails for another reason.

    Examples:
        >>> datacard = get_datacard("ab1001")
        >>> print(datacard["exec"]["scientific_value"])

        >>> datacard = get_datacard("ab1001", version="1")
    """
    params: dict[str, str] = {}
    if version:
        params["version"] = version

    resp = _api_get(
        f"/datasets/{block_id}/datacard",
        params=params or None,
        not_found_msg=_block_not_found_msg(block_id),
        timeout=30,
    )
    try:
        return resp.json()
    except ValueError as e:
        raise RuntimeError(
            f"Datacard response was not valid JSON for {block_id}: {e}"
        ) from e


def get_schema(
    block_id: str,
    version: str | None = None,
    mode: str | None = None,
) -> dict[str, Any]:
    """Get the CSV schema for a data block.

    The returned dict can be unpacked directly into pd.read_csv():
        schema = get_schema("ab1001")
        df = pd.read_csv("data.csv", **schema)

    Args:
        block_id: The data block identifier (e.g., "ab1001").
        version: Optional version number. If not specified, returns latest.
        mode: Data mode (e.g., "source", "ml"). If not specified, uses "source".

    Returns:
        Dict with 'dtype' (column name -> pandas dtype) and 'parse_dates' (list of column names).

    Raises:
        ValueError: If the data block is not found.
        requests.HTTPError: If the API request fails for another reason.

    Examples:
        >>> schema = get_schema("ab1001")
        >>> schema
        {'dtype': {'id': 'Int64', 'name': 'string'}, 'parse_dates': ['created_at']}

        >>> schema = get_schema("ab1001", mode="ml")
        >>> df = pd.read_csv("data.csv", **schema)
    """
    params: dict[str, str] = {}
    if version:
        params["version"] = version
    if mode:
        params["mode"] = mode

    resp = _api_get(
        f"/datasets/{block_id}/schema",
        params=params or None,
        not_found_msg=_block_not_found_msg(block_id),
        timeout=30,
    )
    return resp.json()


def get_data(
    block_id: str,
    version: str | None = None,
    mode: str | None = None,
    format: str | None = None,
    output_path: str | None = None,
    output_compressed: bool = False,
    progress: bool | None = None,
    schema: bool = True,
    load_negatives: bool = True,
) -> str | pd.DataFrame | pl.DataFrame | None:
    """Get CSV data for a data block.

    Args:
        block_id: The data block identifier (e.g., "ab1001").
        version: Optional version number. If not specified, returns latest.
        mode: Data mode (e.g., "source", "ml"). If not specified, uses "source".
        format: Output format - "csv" (string), "pandas" (DataFrame), or "polars" (DataFrame).
                If None, uses the default from ~/.aab/config.json (defaults to "pandas").
        output_path: If specified, writes to this path and returns None.
        output_compressed: If True, write gzipped data; if False, decompress first.
        progress: If True, show download progress bar.
                  If None, uses the default from ~/.aab/config.json (defaults to True).
        schema: If True (default), fetch and apply the data block's column schema
                for recommended dtype handling. Set to False to skip schema enforcement
        load_negatives: If True (default), Retain negative control rows where
                        mata_description starts with "ANeg" or matalpha_description
                        starts with "AlphaNeg". If False, filters out controls.

    Returns:
        CSV data as string, pandas DataFrame, or polars DataFrame depending on format.
        Returns None if output_path is specified, or if format="pandas"/"polars" is
        requested but the library is not installed (an error is printed to stderr).

    Raises:
        ValueError: If format is not one of "csv", "pandas", or "polars",
            or if the data block is not found.
        requests.HTTPError: If the API request fails for another reason.

    Examples:
        >>> df = get_data("ab1001", format="pandas")
        >>> df.head()

        >>> df = get_data("ab1001", mode="ml", format="polars")

        >>> csv = get_data("ab1001")
        >>> print(csv[:100])

        >>> get_data("ab1001", output_path="data.csv.gz")  # Download to file
    """
    format_explicit = format is not None
    if format is None:
        format = get_api_format()

    if format not in ("csv", "pandas", "polars"):
        raise ValueError(f"format must be 'csv', 'pandas', or 'polars', got '{format}'")

    if progress is None:
        progress = get_progress()

    params: dict[str, str] = {"redirect": "false"}
    if version:
        params["version"] = version
    if mode:
        params["mode"] = mode

    resp = _api_get(
        f"/datasets/{block_id}/data",
        params=params,
        not_found_msg=_block_not_found_msg(block_id),
        timeout=30,
    )

    body = resp.json()
    s3_url: str = body["url"]
    schema_data: dict[str, Any] = body.get("schema", {}) if schema else {}

    log_request(
        "GET",
        f"/datasets/{block_id}/data",
        {
            "id": block_id,
            "version": version or "latest",
            "url": s3_url,
            "schema": str(schema_data) if schema_data else "none",
        },
    )

    def _download_bytes(url: str) -> bytes:
        """Download raw bytes from a URL with an optional progress bar."""
        with requests.get(url, stream=True, timeout=300) as r:
            r.raise_for_status()
            total = int(r.headers.get("Content-Length", 0))
            chunks = []
            if progress and total:
                try:
                    with tqdm(
                        total=total, unit="B", unit_scale=True, desc="Downloading"
                    ) as pbar:
                        for chunk in r.iter_content(chunk_size=8192):
                            chunks.append(chunk)
                            pbar.update(len(chunk))
                except ImportError:
                    for chunk in r.iter_content(chunk_size=8192):
                        chunks.append(chunk)
            else:
                for chunk in r.iter_content(chunk_size=8192):
                    chunks.append(chunk)
            return b"".join(chunks)

    negative_message = "Filtering out negative control rows (ANeg/AlphaNeg)..."
    loading_dataframe_message = "Loading DataFrame..."

    # Write file to disk if output_path specified
    if output_path:
        data = _download_bytes(s3_url)
        if not load_negatives:
            csv_text = gzip.decompress(data).decode("utf-8")
            csv_text = _filter_negatives_csv(csv_text, progress=progress)
            data = csv_text.encode("utf-8")
            if output_compressed:
                data = gzip.compress(data)
        else:
            if not output_compressed:
                data = gzip.decompress(data)
        with open(output_path, "wb") as f:
            f.write(data)
        return None

    get_data_alts = (
        "call get_data(format='csv') to get the CSV as a string, or pass "
        "output_path=... to write the file directly to disk."
    )
    if format == "pandas":
        try:
            import pandas as pd
        except ImportError:
            _print_missing_package_error(
                "get_data", "pandas", get_data_alts, format_explicit
            )
            return None
        # low_memory=False: read entire file at once to avoid spurious DtypeWarning
        # from chunked inference. Safe because the file is already fully downloaded.
        read_kwargs = {"compression": "gzip", "low_memory": False, **schema_data}
        if progress:
            data = _download_bytes(s3_url)
            tqdm.write(loading_dataframe_message)
            df = pd.read_csv(io.BytesIO(data), **read_kwargs)
        else:
            df = pd.read_csv(s3_url, **read_kwargs)
        if not load_negatives:
            if progress:
                tqdm.write(negative_message)
            # Filter out negative control rows (ANeg/AlphaNeg prefixes) using pandas boolean mask
            _check_description_columns(df.columns)
            mask = df["mata_description"].str.startswith("ANeg", na=False) | df[
                "matalpha_description"
            ].str.startswith("AlphaNeg", na=False)
            df = df[~mask].reset_index(drop=True)
        return df

    if format == "polars":
        try:
            import polars
        except ImportError:
            _print_missing_package_error(
                "get_data", "polars", get_data_alts, format_explicit
            )
            return None
        polars_kwargs = pandas_schema_to_polars(schema_data)
        data = _download_bytes(s3_url)
        if progress:
            tqdm.write(loading_dataframe_message)
        df = polars.read_csv(io.BytesIO(data), **polars_kwargs)
        if not load_negatives:
            if progress:
                tqdm.write(negative_message)
            # Filter out negative control rows (ANeg/AlphaNeg prefixes) using polars expressions
            _check_description_columns(df.columns)
            mask = polars.col("mata_description").str.starts_with("ANeg") | polars.col(
                "matalpha_description"
            ).str.starts_with("AlphaNeg")
            df = df.filter(~mask)
        return df

    # For csv format, decompress
    data = _download_bytes(s3_url)
    csv_text = gzip.decompress(data).decode("utf-8")
    if not load_negatives:
        csv_text = _filter_negatives_csv(csv_text, progress=progress)
    return csv_text


def download_structures(
    block_id: str,
    version: str | None = None,
    output_path: str | None = None,
    progress: bool | None = None,
) -> str:
    """Download all structure files for a data block as a zip archive.

    Args:
        block_id: The data block identifier (e.g., "ab1001").
        version: Optional version number. If not specified, returns latest.
        output_path: Output file path. Defaults to the filename from the S3 URL.
        progress: If True, show download progress bar.
                  If None, uses the default from ~/.aab/config.json.

    Returns:
        The output file path.

    Raises:
        ValueError: If the data block is not found.
        requests.HTTPError: If the API request fails for another reason.

    Examples:
        >>> download_structures("ab1614")
        >>> download_structures("ab1614", output_path="my_structures.zip")
    """
    if progress is None:
        progress = get_progress()

    params: dict[str, str] = {}
    if version:
        params["version"] = version

    resp = _api_get(
        f"/datasets/{block_id}/structures/download",
        params=params or None,
        not_found_msg=_block_not_found_msg(block_id),
        timeout=30,
    )

    s3_url: str = resp.json()["url"]

    if output_path is None:
        output_path = PurePosixPath(urlparse(s3_url).path).name

    # Stream download from S3 to file
    start_time = time.time()
    error_msg = None
    status_code = None
    try:
        with requests.get(s3_url, stream=True, timeout=600) as r:
            status_code = r.status_code
            r.raise_for_status()
            total = int(r.headers.get("Content-Length", 0))
            with open(output_path, "wb") as f:
                if progress and total:
                    try:
                        with tqdm(
                            total=total, unit="B", unit_scale=True, desc="Downloading"
                        ) as pbar:
                            for chunk in r.iter_content(chunk_size=8192):
                                f.write(chunk)
                                pbar.update(len(chunk))
                    except ImportError:
                        for chunk in r.iter_content(chunk_size=8192):
                            f.write(chunk)
                else:
                    for chunk in r.iter_content(chunk_size=8192):
                        f.write(chunk)
    except requests.RequestException as e:
        error_msg = str(e)
        raise
    finally:
        latency_ms = int((time.time() - start_time) * 1000)
        log_request(
            "GET",
            f"/datasets/{block_id}/structures/download (S3)",
            {"output_path": output_path},
            status_code,
            error_msg,
            latency_ms,
        )

    return output_path


def download_structure(
    block_id: str,
    filename: str,
    version: str | None = None,
    output_path: str | None = None,
) -> str:
    """Download a single structure file for a dataset.

    Args:
        block_id: The data block identifier (e.g., "ab1614").
        filename: The structure filename to download (e.g., "structure.cif").
        version: Optional version number. If not specified, returns latest.
        output_path: If specified, writes to this path and returns the path.
                     If None, returns the file content as a string.

    Returns:
        The file content as a string (no output_path), or the output file path.

    Raises:
        ValueError: If the structure file or data block is not found.
        requests.HTTPError: If the API request fails for another reason.

    Examples:
        >>> content = download_structure("ab1614", "structure.cif")
        >>> download_structure("ab1614", "structure.cif", output_path="my_file.cif")
    """
    params: dict[str, str] = {}
    if version:
        params["version"] = version

    resp = _api_get(
        f"/datasets/{block_id}/structures/{filename}",
        params=params or None,
        not_found_msg=(
            f"Structure file '{filename}' not found for data block "
            f"'{block_id}'. Check the block ID and filename, and that you "
            f"have access to it."
        ),
        timeout=300,
    )

    if output_path is None:
        return resp.text

    with open(output_path, "w") as f:
        f.write(resp.text)

    return output_path


def pandas_schema_to_polars(schema: dict[str, Any]) -> dict[str, Any]:
    """Convert a pandas CSV schema dict to polars-compatible kwargs.

    Takes the schema dict from get_schema() (designed for pd.read_csv) and
    converts it to kwargs suitable for polars.read_csv().

    Args:
        schema: Pandas schema dict with 'dtype' and optionally 'parse_dates'.

    Returns:
        Dict with 'schema_overrides' and 'try_parse_dates' for polars.read_csv().

    Examples:
        >>> schema = get_schema("ab1001")
        >>> polars_kwargs = pandas_schema_to_polars(schema)
        >>> df = polars.read_csv("data.csv", **polars_kwargs)
    """
    import polars as pl

    PANDAS_TO_POLARS = {
        "int64": pl.Int64,
        "Int64": pl.Int64,
        "int32": pl.Int32,
        "Int32": pl.Int32,
        "int16": pl.Int16,
        "Int16": pl.Int16,
        "int8": pl.Int8,
        "Int8": pl.Int8,
        "uint64": pl.UInt64,
        "UInt64": pl.UInt64,
        "uint32": pl.UInt32,
        "UInt32": pl.UInt32,
        "uint16": pl.UInt16,
        "UInt16": pl.UInt16,
        "uint8": pl.UInt8,
        "UInt8": pl.UInt8,
        "float64": pl.Float64,
        "Float64": pl.Float64,
        "float32": pl.Float32,
        "Float32": pl.Float32,
        "bool": pl.Boolean,
        "boolean": pl.Boolean,
        "string": pl.Utf8,
        "object": pl.Utf8,
        "str": pl.Utf8,
        "category": pl.Categorical,
    }

    overrides: dict[str, Any] = {}
    for col, pd_dtype in schema.get("dtype", {}).items():
        overrides[col] = PANDAS_TO_POLARS.get(pd_dtype, pl.Utf8)

    result: dict[str, Any] = {}
    if overrides:
        result["schema_overrides"] = overrides
    if schema.get("parse_dates"):
        result["try_parse_dates"] = True
    return result


def _api_get(
    endpoint: str,
    params: dict[str, str] | None = None,
    not_found_msg: str | None = None,
    **kwargs: Any,
) -> requests.Response:
    """Make an authenticated GET request with logging.

    If ``not_found_msg`` is given, a 404 response is re-raised as a
    ``ValueError`` with that message instead of the raw ``HTTPError``.
    """
    show_startup_notices()
    headers = get_token_header()
    start_time = time.time()
    error_msg = None
    status_code = None
    try:
        resp = requests.get(
            f"{get_base_url()}{endpoint}",
            headers=headers,
            params=params,
            **kwargs,
        )
        status_code = resp.status_code
        resp.raise_for_status()
        return resp
    except requests.HTTPError as e:
        error_msg = str(e)
        if not_found_msg and status_code == 404:
            raise ValueError(not_found_msg) from None
        raise
    except requests.RequestException as e:
        error_msg = str(e)
        raise
    finally:
        latency_ms = int((time.time() - start_time) * 1000)
        log_request("GET", endpoint, params, status_code, error_msg, latency_ms)


def _block_not_found_msg(block_id: str) -> str:
    """Friendly error message for a missing or inaccessible data block."""
    return (
        f"Data block '{block_id}' not found. Check the block ID "
        f"(e.g. 'ab1001') and that you have access to it."
    )


def _check_description_columns(columns: Iterable[str]) -> None:
    """Raise ValueError if required description columns are missing."""
    col_set = set(columns)
    missing = [
        c for c in ("mata_description", "matalpha_description") if c not in col_set
    ]
    if missing:
        raise ValueError(
            f"Cannot filter negatives: missing column(s) {', '.join(missing)}"
        )


def _filter_negatives_csv(csv_text: str, progress: bool) -> str:
    """Remove negative control rows from CSV text.

    Filters out rows where mata_description starts with 'ANeg'
    or matalpha_description starts with 'AlphaNeg'.
    """
    lines = csv_text.split("\n")
    if not lines or not lines[0]:
        raise ValueError("Cannot filter negatives: CSV data has no header row")

    reader = csv.reader(lines)
    header = next(reader)
    _check_description_columns(header)
    mata_idx = header.index("mata_description")
    matalpha_idx = header.index("matalpha_description")

    data_lines = lines[1:]
    if progress:
        lines_and_rows = tqdm(
            zip(data_lines, reader),
            total=len(data_lines),
            desc="Filtering  ",  # make the same width as Downloading to align progress bars
            unit=" rows",
        )
    else:
        lines_and_rows = zip(data_lines, reader)

    kept = [lines[0]]
    for line, row in lines_and_rows:
        if not row:
            continue
        if row[mata_idx].startswith("ANeg"):
            continue
        if row[matalpha_idx].startswith("AlphaNeg"):
            continue
        kept.append(line)
    kept.append("")  # preserve trailing newline
    return "\n".join(kept)


def _print_missing_package_error(
    func_name: str, package: str, alternatives: str, format_explicit: bool
) -> None:
    """Print a friendly stderr error when an optional dataframe lib is missing."""
    call = f"{func_name}(format='{package}')" if format_explicit else f"{func_name}()"
    print(
        f"Error: {call} requires {package}. Either `pip install {package}`, "
        f"{alternatives}",
        file=sys.stderr,
    )
