"""Universal data parser for flexible data formats."""

from __future__ import annotations

import base64
import io
import json
from typing import Any, Dict, List, Optional, Union

import pandas as pd
from pydantic import BaseModel


async def parse_flexible_data(
    data: Union[List[Dict[str, Any]], str, bytes],
    format: str = "json",
    column_types: Optional[Dict[str, str]] = None,
) -> pd.DataFrame:
    """
    Parse data from various formats into a pandas DataFrame.

    Args:
        data: The input data (JSON list, CSV string, or base64-encoded binary)
        format: The format of the data ("json", "csv", "parquet")
        column_types: Optional column type hints

    Returns:
        Parsed pandas DataFrame
    """
    if format == "json":
        if isinstance(data, list):
            df = pd.DataFrame(data)
        elif isinstance(data, str):
            df = pd.DataFrame(json.loads(data))
        else:
            raise ValueError("JSON data must be a list or JSON string")

    elif format == "csv":
        if isinstance(data, str):
            df = pd.read_csv(io.StringIO(data))
        elif isinstance(data, bytes):
            df = pd.read_csv(io.BytesIO(data))
        else:
            raise ValueError("CSV data must be string or bytes")

    elif format == "parquet":
        if isinstance(data, str):
            # Assume base64 encoded
            binary_data = base64.b64decode(data)
            df = pd.read_parquet(io.BytesIO(binary_data))
        elif isinstance(data, bytes):
            df = pd.read_parquet(io.BytesIO(data))
        else:
            raise ValueError("Parquet data must be base64 string or bytes")

    else:
        raise ValueError(f"Unsupported format: {format}")

    # Apply column types if provided
    if column_types:
        for col, dtype in column_types.items():
            if col in df.columns:
                try:
                    df[col] = df[col].astype(dtype)
                except Exception:
                    pass  # Skip if conversion fails

    return df


async def parse_data_source(
    source: BaseModel, max_rows: Optional[int] = None
) -> pd.DataFrame:
    """
    Universal parser that handles any data format.
    This is what makes your API work with ANY integration.

    Args:
        source: Data source configuration
        max_rows: Maximum rows to read (for memory safety)
    """
    # Check data size if it's a string
    if isinstance(source.data, str):
        data_size_mb = len(source.data.encode()) / (1024 * 1024)
        if data_size_mb > 100:  # 100MB limit for direct string data
            raise ValueError(
                f"Data too large ({data_size_mb:.1f} MB). Please use file upload for large datasets."
            )

    if source.format == "json":
        if isinstance(source.data, str):
            data = json.loads(source.data)
        else:
            data = source.data
        df = pd.DataFrame(data)

    elif source.format == "csv":
        if isinstance(source.data, bytes):
            df = pd.read_csv(io.BytesIO(source.data))
        else:
            df = pd.read_csv(io.StringIO(source.data))

    elif source.format == "parquet":
        if isinstance(source.data, bytes):
            df = pd.read_parquet(io.BytesIO(source.data))
        else:
            # Base64 encoded parquet
            df = pd.read_parquet(io.BytesIO(base64.b64decode(source.data)))

    # Apply column type hints if provided
    if source.column_types:
        for col, dtype in source.column_types.items():
            if col in df.columns:
                df[col] = df[col].astype(dtype)

    # Apply ID column hint if provided
    if source.id_column and source.id_column in df.columns:
        # Store this hint for auto-detection
        df.attrs["suggested_id_column"] = source.id_column

    # Apply row limit if specified
    if max_rows and len(df) > max_rows:
        import logging

        logging.warning(
            f"DataFrame has {len(df)} rows. Limiting to {max_rows} rows for memory safety."
        )
        df = df.head(max_rows)

    return df
