"""
Polars utilities with graceful fallbacks
"""
import logging
from typing import Any, Dict, List, Optional, Union
from pathlib import Path

LOG = logging.getLogger(__name__)

# Try to import Polars
try:
    import polars as pl
    POLARS_AVAILABLE = True
    LOG.debug("Polars successfully imported")
except ImportError:
    POLARS_AVAILABLE = False
    pl = None
    LOG.debug("Polars not available, using fallbacks")

# Try to import BioPython
try:
    from Bio import SeqIO
    from Bio.KEGG import REST
    BIOPYTHON_AVAILABLE = True
    LOG.debug("BioPython successfully imported")
except ImportError:
    BIOPYTHON_AVAILABLE = False
    SeqIO = None
    REST = None
    LOG.debug("BioPython not available")

# Try to import polars-bio
try:
    from polars_bio import read_fasta as pl_read_fasta
    POLARS_BIO_AVAILABLE = True
    LOG.debug("polars-bio successfully imported")
except ImportError:
    POLARS_BIO_AVAILABLE = False
    pl_read_fasta = None
    LOG.debug("polars-bio not available")
except Exception as e:
    # Handle other import errors
    LOG.warning(f"Error importing polars-bio: {e}")
    POLARS_BIO_AVAILABLE = False
    pl_read_fasta = None


def check_polars_available() -> bool:
    """Check if Polars is available."""
    return POLARS_AVAILABLE


def check_biopython_available() -> bool:
    """Check if BioPython is available."""
    return BIOPYTHON_AVAILABLE


def check_polars_bio_available() -> bool:
    """Check if polars-bio is available."""
    return POLARS_BIO_AVAILABLE


def create_dataframe(
    data: List[Dict[str, Any]], 
    schema: Optional[Dict[str, Any]] = None
) -> Union['pl.DataFrame', List[Dict[str, Any]]]:
    """
    Create DataFrame from list of dicts using Polars if available.
    
    Args:
        data: List of dictionaries
        schema: Optional schema for Polars
    
    Returns:
        Polars DataFrame or original list
    """
    if not data:
        return [] if not POLARS_AVAILABLE else pl.DataFrame()
    
    if POLARS_AVAILABLE:
        try:
            if schema:
                # Convert schema to Polars types
                polars_schema = {}
                for col, dtype in schema.items():
                    if dtype == str:
                        polars_schema[col] = pl.Utf8
                    elif dtype == int:
                        polars_schema[col] = pl.Int64
                    elif dtype == float:
                        polars_schema[col] = pl.Float64
                    elif dtype == bool:
                        polars_schema[col] = pl.Boolean
                    else:
                        polars_schema[col] = pl.Utf8
                
                return pl.DataFrame(data, schema=polars_schema)
            else:
                return pl.DataFrame(data)
        except Exception as e:
            LOG.warning(f"Failed to create Polars DataFrame: {e}")
    
    return data


def save_dataframe(
    df: Union['pl.DataFrame', List], 
    filepath: Union[str, Path], 
    **kwargs
) -> None:
    """
    Save DataFrame to file using Polars if available.
    
    Args:
        df: DataFrame or list to save
        filepath: Output file path
        **kwargs: Additional arguments for write methods
    """
    filepath = Path(filepath)
    filepath.parent.mkdir(parents=True, exist_ok=True)
    
    if POLARS_AVAILABLE and isinstance(df, pl.DataFrame):
        ext = filepath.suffix.lower()
        
        try:
            if ext == '.parquet':
                df.write_parquet(filepath, **kwargs)
            elif ext == '.csv':
                df.write_csv(filepath, **kwargs)
            elif ext in ['.tsv', '.txt']:
                df.write_csv(filepath, separator='\t', **kwargs)
            elif ext == '.json':
                df.write_json(filepath, **kwargs)
            elif ext == '.jsonl':
                df.write_ndjson(filepath, **kwargs)
            elif ext in ['.feather', '.arrow']:
                df.write_ipc(filepath, **kwargs)
            else:
                # Default to CSV
                df.write_csv(filepath, **kwargs)
            
            LOG.info(f"Saved DataFrame to {filepath} using Polars")
            return
            
        except Exception as e:
            LOG.warning(f"Failed to save with Polars: {e}. Falling back to manual.")
    
    # Fallback to manual writing
    with open(filepath, 'w', encoding='utf-8') as f:
        if isinstance(df, list) and df and isinstance(df[0], dict):
            # Write as TSV
            import csv
            fieldnames = df[0].keys()
            writer = csv.DictWriter(f, fieldnames=fieldnames, delimiter='\t')
            writer.writeheader()
            writer.writerows(df)
        elif isinstance(df, list):
            # Write each item on a new line
            for row in df:
                if isinstance(row, (list, tuple)):
                    f.write('\t'.join(str(x) for x in row) + '\n')
                else:
                    f.write(str(row) + '\n')
        else:
            f.write(str(df))
    
    LOG.info(f"Saved data to {filepath} (manual)")


def load_dataframe(
    filepath: Union[str, Path], 
    **kwargs
) -> Union['pl.DataFrame', List]:
    """
    Load DataFrame from file using Polars if available.
    
    Args:
        filepath: Input file path
        **kwargs: Additional arguments for read methods
    
    Returns:
        Polars DataFrame or list
    """
    filepath = Path(filepath)
    
    if POLARS_AVAILABLE and filepath.exists():
        ext = filepath.suffix.lower()
        
        try:
            if ext == '.parquet':
                return pl.read_parquet(filepath, **kwargs)
            elif ext == '.csv':
                return pl.read_csv(filepath, **kwargs)
            elif ext in ['.tsv', '.txt']:
                return pl.read_csv(filepath, separator='\t', **kwargs)
            elif ext == '.json':
                return pl.read_json(filepath, **kwargs)
            elif ext == '.jsonl':
                return pl.read_ndjson(filepath, **kwargs)
            elif ext in ['.feather', '.arrow']:
                return pl.read_ipc(filepath, **kwargs)
        except Exception as e:
            LOG.warning(f"Failed to load with Polars: {e}")
    
    # Fallback to manual reading
    if not filepath.exists():
        raise FileNotFoundError(f"File not found: {filepath}")
    
    data = []
    with open(filepath, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith('#'):
                data.append(line.split('\t'))
    
    return data


def filter_dataframe(
    df: Union['pl.DataFrame', List], 
    condition: Any
) -> Union['pl.DataFrame', List]:
    """
    Filter DataFrame using Polars if available.
    
    Args:
        df: DataFrame or list to filter
        condition: Filter condition (string expression or callable)
    
    Returns:
        Filtered DataFrame or list
    """
    if POLARS_AVAILABLE and isinstance(df, pl.DataFrame):
        if isinstance(condition, str):
            try:
                return df.filter(pl.col(condition))
            except:
                # If column doesn't exist, try different approach
                return df
        elif callable(condition):
            # Note: This is less efficient with Polars
            return df.filter(pl.col("*").map_batches(lambda s: condition(s)))
    
    # Fallback for lists
    if isinstance(df, list) and callable(condition):
        return [row for row in df if condition(row)]
    
    return df


def merge_dataframes(
    df1: Union['pl.DataFrame', List], 
    df2: Union['pl.DataFrame', List], 
    on: str, 
    how: str = "inner"
) -> Union['pl.DataFrame', List]:
    """
    Merge two DataFrames using Polars if available.
    
    Args:
        df1: First DataFrame
        df2: Second DataFrame
        on: Column to join on
        how: Join type (inner, left, outer, cross)
    
    Returns:
        Merged DataFrame or list
    """
    if (POLARS_AVAILABLE and isinstance(df1, pl.DataFrame) 
        and isinstance(df2, pl.DataFrame)):
        
        how_map = {
            "inner": "inner",
            "left": "left",
            "outer": "outer",
            "cross": "cross"
        }
        
        join_method = how_map.get(how, "inner")
        return df1.join(df2, on=on, how=join_method)
    
    # Fallback for lists of dicts
    if (isinstance(df1, list) and isinstance(df2, list) 
        and df1 and df2 and isinstance(df1[0], dict)):
        
        result = []
        dict2 = {row.get(on): row for row in df2 if on in row}
        
        for row1 in df1:
            if on in row1 and row1[on] in dict2:
                merged = {**row1, **dict2[row1[on]]}
                result.append(merged)
            elif how == "left" and on in row1:
                result.append(row1.copy())
        
        return result
    
    return df1


def group_dataframe(
    df: Union['pl.DataFrame', List], 
    by: str, 
    agg: Optional[Dict[str, str]] = None
) -> Union['pl.DataFrame', List]:
    """
    Group DataFrame using Polars if available.
    
    Args:
        df: DataFrame to group
        by: Column to group by
        agg: Aggregation functions {column: function}
    
    Returns:
        Grouped DataFrame
    """
    if POLARS_AVAILABLE and isinstance(df, pl.DataFrame):
        if agg:
            # Convert string functions to Polars expressions
            agg_exprs = []
            for col, func in agg.items():
                if func == "count":
                    agg_exprs.append(pl.len().alias(f"{col}_count"))
                elif func == "sum":
                    agg_exprs.append(pl.col(col).sum().alias(f"{col}_sum"))
                elif func == "mean":
                    agg_exprs.append(pl.col(col).mean().alias(f"{col}_mean"))
                elif func == "first":
                    agg_exprs.append(pl.col(col).first().alias(f"{col}_first"))
                elif func == "last":
                    agg_exprs.append(pl.col(col).last().alias(f"{col}_last"))
            
            return df.group_by(by).agg(*agg_exprs)
        else:
            return df.group_by(by).agg(pl.all())
    
    # Fallback for lists
    if isinstance(df, list) and df and isinstance(df[0], dict):
        groups = {}
        for row in df:
            key = row.get(by)
            if key not in groups:
                groups[key] = []
            groups[key].append(row)
        
        result = []
        for key, rows in groups.items():
            group_dict = {by: key}
            if agg:
                for agg_col, agg_func in agg.items():
                    if agg_func == "count":
                        group_dict[f"{agg_col}_count"] = len(rows)
                    elif agg_func == "sum":
                        # Sum numeric columns
                        col_values = [r.get(agg_col, 0) for r in rows]
                        numeric_values = [v for v in col_values if isinstance(v, (int, float))]
                        group_dict[f"{agg_col}_sum"] = sum(numeric_values)
                    elif agg_func == "first":
                        group_dict[f"{agg_col}_first"] = rows[0].get(agg_col) if rows else None
                    elif agg_func == "last":
                        group_dict[f"{agg_col}_last"] = rows[-1].get(agg_col) if rows else None
            else:
                group_dict["rows"] = rows
            
            result.append(group_dict)
        
        return result
    
    return df


# Export constants for use in other modules
__all__ = [
    'POLARS_AVAILABLE',
    'BIOPYTHON_AVAILABLE', 
    'POLARS_BIO_AVAILABLE',
    'check_polars_available',
    'check_biopython_available',
    'check_polars_bio_available',
    'create_dataframe',
    'save_dataframe',
    'load_dataframe',
    'filter_dataframe',
    'merge_dataframes',
    'group_dataframe',
    'pl_read_fasta'  # Export for use in fasta.py
]
