"""
Common utilities for GPATHEX with modern dependencies
"""
import logging
import sys
import time
from pathlib import Path
from typing import List, Dict, Any, Optional, Union, Iterator
import urllib.request
import urllib.error

# Try to import modern dependencies
try:
    import polars as pl
    POLARS_AVAILABLE = True
except ImportError:
    POLARS_AVAILABLE = False
    pl = None

try:
    from Bio.KEGG import REST
    from Bio.KEGG import Enzyme, Compound, Map
    BIOPYTHON_AVAILABLE = True
except ImportError:
    BIOPYTHON_AVAILABLE = False
    REST = None

try:
    from polars_bio import read_fasta
    POLARS_BIO_AVAILABLE = True
except ImportError:
    POLARS_BIO_AVAILABLE = False

LOG = logging.getLogger(__name__)

__version__ = "1.0.0"
__author__ = ("Maulana Malik Nashrulloh",)
__email__ = "maulana@genbinesia.or.id"

def setup_logging(level: str = "INFO") -> None:
    """Configure logging with rich formatting if available."""
    try:
        from rich.logging import RichHandler
        from rich.console import Console
        
        console = Console()
        handler = RichHandler(
            console=console,
            show_time=True,
            show_level=True,
            show_path=True,
            rich_tracebacks=True
        )
        
        logging.basicConfig(
            level=getattr(logging, level.upper()),
            format="%(message)s",
            datefmt="[%X]",
            handlers=[handler]
        )
        LOG.info("Using rich logging for beautiful output! ✨")
        
    except ImportError:
        # Fallback to basic logging
        logging.basicConfig(
            stream=sys.stderr,
            level=getattr(logging, level.upper()),
            format="[%(levelname)s] %(asctime)s - %(name)s - %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S"
        )

def safe_request(url: str, max_retries: int = 3, timeout: int = 30) -> bytes:
    """Safely request URL with retries using Bio.KEGG if available."""
    if BIOPYTHON_AVAILABLE and "kegg.jp" in url:
        try:
            # Use Bio.KEGG.REST for KEGG API calls
            LOG.debug(f"Using Bio.KEGG.REST for: {url}")
            response = REST.kegg_get(url.split("=")[-1].split("&")[0])
            return response.read().encode('utf-8')
        except Exception as e:
            LOG.warning(f"Bio.KEGG.REST failed, falling back to urllib: {e}")
    
    # Fallback to urllib
    last_error = None
    for attempt in range(max_retries):
        try:
            with urllib.request.urlopen(url, timeout=timeout) as response:
                return response.read()
        except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError) as e:
            last_error = e
            if attempt < max_retries - 1:
                wait_time = 2 ** attempt
                LOG.warning(f"Request failed (attempt {attempt + 1}/{max_retries}): {e}. Retrying in {wait_time}s...")
                time.sleep(wait_time)
    
    raise last_error or Exception("Request failed")

def read_org(file_path: Union[str, Path]) -> pl.DataFrame:
    """
    Read .org file using Polars for speed.
    
    Returns:
        DataFrame with columns: [org_abbr, name, url]
    """
    file_path = Path(file_path)
    if not file_path.exists():
        raise FileNotFoundError(f"Organization file not found: {file_path}")
    
    LOG.info(f"Reading organism information from {file_path}")
    
    if POLARS_AVAILABLE:
        try:
            df = pl.read_csv(
                file_path,
                separator="\t",
                has_header=False,
                new_columns=["org_abbr", "name", "url"],
                null_values=["", "NA", "null"],
                comment_prefix="#"
            )
            LOG.info(f"Loaded {len(df)} records using Polars")
            return df
        except Exception as e:
            LOG.warning(f"Polars read failed, falling back to manual: {e}")
    
    # Fallback to manual reading
    result = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if line and not line.startswith('#'):
                parts = line.split("\t")
                if len(parts) >= 3:
                    result.append(parts[:3])
    
    if POLARS_AVAILABLE:
        df = pl.DataFrame(result, schema=["org_abbr", "name", "url"])
    else:
        # Return as list if Polars not available
        return result
    
    LOG.info(f"Found {len(df)} records")
    return df

def read_fasta_file(file_path: Union[str, Path]) -> Union[pl.DataFrame, Iterator]:
    """
    Read FASTA file using the best available method.
    Priority: polars-bio > biopython > manual parser
    
    Returns:
        Polars DataFrame or iterator of sequences
    """
    file_path = Path(file_path)
    
    if POLARS_BIO_AVAILABLE:
        LOG.debug(f"Using polars-bio to read {file_path}")
        return read_fasta(file_path)
    
    elif BIOPYTHON_AVAILABLE:
        from Bio import SeqIO
        LOG.debug(f"Using BioPython SeqIO to read {file_path}")
        return SeqIO.parse(file_path, "fasta")
    
    else:
        # Fallback to manual parser
        LOG.debug(f"Using manual FASTA parser for {file_path}")
        from .FastaReader import open_fasta
        return open_fasta(file_path)

def parse_kegg_rest_response(response: bytes) -> pl.DataFrame:
    """
    Parse KEGG REST API response using Bio.KEGG if available.
    
    Args:
        response: Bytes from KEGG REST API
    
    Returns:
        Polars DataFrame with parsed data
    """
    text = response.decode('utf-8')
    
    if BIOPYTHON_AVAILABLE:
        # Try to use Bio.KEGG parsers based on content
        if "ENTRY" in text[:100]:
            from Bio.KEGG import parse
            
            records = list(parse(text))
            if records:
                # Convert to list of dicts
                data = []
                for record in records:
                    if hasattr(record, 'entry'):
                        data.append({
                            'entry': getattr(record, 'entry', ''),
                            'name': getattr(record, 'name', ''),
                            'definition': getattr(record, 'definition', ''),
                            'pathway': getattr(record, 'pathway', ''),
                            'dblinks': getattr(record, 'dblinks', ''),
                        })
                
                if POLARS_AVAILABLE and data:
                    return pl.DataFrame(data)
    
    # Fallback to tabular parsing
    lines = text.strip().split('\n')
    data = []
    for line in lines:
        if line and not line.startswith('#'):
            data.append(line.split('\t'))
    
    if POLARS_AVAILABLE and data:
        # Infer schema from first few rows
        return pl.DataFrame(data)
    
    return data if data else []

def kegg_api_lookup(identifier: str, database: str = "ko") -> Optional[pl.DataFrame]:
    """
    Look up information from KEGG API using Bio.KEGG.
    
    Args:
        identifier: KEGG identifier (e.g., "K00001", "hsa00010")
        database: KEGG database name (ko, pathway, compound, etc.)
    
    Returns:
        Polars DataFrame with results or None
    """
    if not BIOPYTHON_AVAILABLE:
        LOG.warning("BioPython not available for KEGG API lookup")
        return None
    
    try:
        LOG.debug(f"Looking up {identifier} in KEGG {database} database")
        
        if database == "ko":
            from Bio.KEGG import REST
            response = REST.kegg_get(f"ko:{identifier}")
        elif database == "pathway":
            response = REST.kegg_get(f"path:{identifier}")
        elif database == "compound":
            response = REST.kegg_get(f"cpd:{identifier}")
        else:
            response = REST.kegg_get(f"{database}:{identifier}")
        
        return parse_kegg_rest_response(response.read())
        
    except Exception as e:
        LOG.error(f"KEGG API lookup failed for {identifier}: {e}")
        return None

def create_dataframe(data: List[Dict], schema: Optional[Dict] = None) -> Union[pl.DataFrame, List[Dict]]:
    """
    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 POLARS_AVAILABLE and data:
        try:
            if schema:
                return pl.DataFrame(data, schema=schema)
            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 Polars 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()
        
        if ext == '.parquet':
            df.write_parquet(filepath, **kwargs)
        elif ext == '.csv':
            df.write_csv(filepath, **kwargs)
        elif ext == '.tsv':
            df.write_csv(filepath, separator='\t', **kwargs)
        elif ext in ['.json', '.jsonl']:
            df.write_json(filepath, **kwargs)
        else:
            # Default to CSV
            df.write_csv(filepath, **kwargs)
        
        LOG.info(f"Saved DataFrame to {filepath} using Polars")
    
    else:
        # 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
                writer = csv.DictWriter(f, fieldnames=df[0].keys(), delimiter='\t')
                writer.writeheader()
                writer.writerows(df)
            else:
                # Write as plain text
                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')
        
        LOG.info(f"Saved data to {filepath} (manual)")

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
    
    Returns:
        Filtered DataFrame or list
    """
    if POLARS_AVAILABLE and isinstance(df, pl.DataFrame):
        if isinstance(condition, str):
            return df.filter(pl.col(condition))
        elif callable(condition):
            # Convert callable to Polars expression if possible
            return df.filter(condition)
    
    # Fallback to manual filtering for lists
    if isinstance(df, list):
        if 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)
    
    Returns:
        Merged DataFrame or list
    """
    if POLARS_AVAILABLE and isinstance(df1, pl.DataFrame) and isinstance(df2, pl.DataFrame):
        return df1.join(df2, on=on, how=how)
    
    # Fallback to manual merge for lists of dicts
    if isinstance(df1, list) and isinstance(df2, list) and df1 and df2 and isinstance(df1[0], dict):
        result = []
        dict2 = {row[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)
        
        return result
    
    return df1  # Return first if can't merge

def group_dataframe(df: Union[pl.DataFrame, List], by: str, agg: Dict = None) -> Union[pl.DataFrame, List]:
    """
    Group DataFrame using Polars if available.
    
    Args:
        df: DataFrame to group
        by: Column to group by
        agg: Aggregation functions
    
    Returns:
        Grouped DataFrame
    """
    if POLARS_AVAILABLE and isinstance(df, pl.DataFrame):
        if agg:
            return df.group_by(by).agg(**agg)
        return df.group_by(by).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[agg_col] = len(rows)
                    elif agg_func == "sum":
                        # Sum numeric columns
                        col_values = [r.get(agg_col, 0) for r in rows]
                        group_dict[agg_col] = sum(v for v in col_values if isinstance(v, (int, float)))
            else:
                group_dict["rows"] = rows
            result.append(group_dict)
        
        return result
    
    return df
    