#!/usr/bin/env python3
"""
Get KEGG organism classification ranks
"""
import sys
import logging
import string
from pathlib import Path
from typing import Dict, List, Optional, Set
from collections import defaultdict

from ..utils import polars_utils, parsers, taxonomy
from ..config import config

LOG = logging.getLogger(__name__)


# OLD 
# def parse_kegg_organisms_file(keg_file: Path) -> Dict[str, str]:
    # """
    # Parse KEGG organisms file (br08610.keg).
    
    # Returns: {org_code: taxon_id}
    # """
    # if not keg_file.exists():
        # LOG.error(f"KEGG organisms file not found: {keg_file}")
        # return {}
    
    # result = {}
    # current_taxon = ""
    # current_level = "A"
    
    # # Map letters to levels (A=top, B=sub, etc.)
    # level_order = {char: idx for idx, char in enumerate(string.ascii_uppercase)}
    
    # LOG.info(f"Parsing KEGG organisms file: {keg_file}")
    
    # with open(keg_file, 'r', encoding='utf-8') as f:
        # for line_num, line in enumerate(f, 1):
            # line = line.strip()
            # if not line:
                # continue
            
            # tag = line[0]
            
            # # Skip non-letter tags
            # if tag not in level_order:
                # continue
            
            # # Taxon line (contains TAX:)
            # if "TAX:" in line:
                # # Extract taxon ID
                # start = line.find("TAX:") + 4
                # end = line.find("]", start)
                # if start > 4 and end > start:
                    # current_taxon = line[start:end]
                    # current_level = tag
                # continue
            
            # # Organism line (at one level deeper)
            # if (tag in level_order and current_level in level_order and
                # level_order[tag] - level_order[current_level] == 1):
                
                # # Extract organism code
                # parts = line.split()
                # if len(parts) >= 2:
                    # org_code = parts[1]
                    # if not org_code.isdigit():  # Skip numeric IDs
                        # result[org_code] = current_taxon
    
    # LOG.info(f"Found {len(result)} organisms with taxon IDs")
    # return result

#NEW FORMAT
def parse_kegg_organisms_file(keg_file: Path) -> Dict[str, str]:
    """
    Parse KEGG organisms file (br08610.keg).
    
    Returns: {org_code: taxon_id}
    """
    if not keg_file.exists():
        LOG.error(f"KEGG organisms file not found: {keg_file}")
        return {}
    
    result = {}
    current_taxon = ""
    current_level = "A"
    
    # Map letters to levels (A=top, B=sub, etc.)
    level_order = {char: idx for idx, char in enumerate(string.ascii_uppercase)}
    
    LOG.info(f"Parsing KEGG organisms file: {keg_file}")
    
    with open(keg_file, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            
            tag = line[0]
            
            # Skip non-letter tags
            if tag not in level_order:
                continue
            
            # Taxon line (contains TAX:)
            if "TAX:" in line:
                # Extract taxon ID
                start = line.find("TAX:") + 4
                end = line.find("]", start)
                if start > 4 and end > start:
                    taxon_id = line[start:end].strip()
                    if taxon_id and taxon_id != "":  # Skip empty taxon IDs
                        current_taxon = taxon_id
                        current_level = tag
                    else:
                        LOG.debug(f"Empty taxon ID at line {line_num}: {line}")
                        current_taxon = ""
                continue
            
            # Organism line (at one level deeper)
            if (tag in level_order and current_level in level_order and
                level_order[tag] - level_order[current_level] == 1):
                
                # Extract organism code
                parts = line.split()
                if len(parts) >= 2:
                    org_code = parts[1]
                    if not org_code.isdigit():  # Skip numeric IDs
                        if current_taxon:  # Only add if we have a taxon ID
                            result[org_code] = current_taxon
                        else:
                            LOG.warning(f"Skipping {org_code}: No taxon ID")
    
    LOG.info(f"Found {len(result)} organisms with valid taxon IDs")
    
    # Log some statistics
    if LOG.level <= logging.DEBUG:
        LOG.debug(f"Sample of organism-taxon mappings:")
        count = 0
        for org, taxon in list(result.items())[:5]:
            LOG.debug(f"  {org} -> {taxon}")
            count += 1
    
    return result

def parse_ncbi_taxon_file(taxon_file: Path) -> Dict[str, Dict[str, str]]:
    """
    Parse NCBI taxonomy file.
    
    Expected format: tax_id\tkingdom\tphylum\tclass\torder\tfamily\tgenus\tspecies
    
    Returns: {taxon_id: {kingdom: name, phylum: name, ...}}
    """
    if not taxon_file.exists():
        LOG.error(f"Taxonomy file not found: {taxon_file}")
        return {}
    
    result = {}
    
    LOG.info(f"Parsing NCBI taxonomy file: {taxon_file}")
    
    standard_ranks = ['kingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species']
    
    with open(taxon_file, 'r', encoding='utf-8') as f:
        # Skip header
        header = f.readline()
        
        for line_num, line in enumerate(f, 2):
            line = line.strip()
            if not line:
                continue
            
            parts = line.split('\t')
            if len(parts) >= 1:
                taxon_id = parts[0]
                
                taxon_info = {'taxon_id': taxon_id}
                
                # Map parts to standard ranks
                for i, rank in enumerate(standard_ranks, 1):
                    if i < len(parts) and parts[i]:
                        taxon_info[rank] = parts[i]
                    else:
                        taxon_info[rank] = ''
                
                # Try to extract species name
                if not taxon_info.get('species') and len(parts) > 1:
                    # Last part might be species
                    for i in range(len(parts)-1, 0, -1):
                        if parts[i] and ' ' in parts[i]:  # Likely species name
                            taxon_info['species'] = parts[i]
                            break
                
                result[taxon_id] = taxon_info
    
    LOG.info(f"Loaded {len(result)} taxon entries")
    return result


def download_taxonomy_if_needed(taxon_file: Path, taxon_ids: Optional[Set[str]] = None) -> bool:
    """
    Download NCBI taxonomy if file doesn't exist.
    
    Args:
        taxon_file: Target taxonomy file
        taxon_ids: Set of taxon IDs to include (None for all)
    
    Returns:
        True if successful, False otherwise
    """
    if taxon_file.exists():
        LOG.info(f"Taxonomy file already exists: {taxon_file}")
        return True
    
    LOG.info(f"Taxonomy file not found: {taxon_file}")
    LOG.info("Downloading NCBI taxonomy database...")
    
    try:
        # Create parent directory
        taxon_file.parent.mkdir(parents=True, exist_ok=True)
        
        # Download and create taxonomy
        count = taxonomy.download_and_create_taxonomy(str(taxon_file), list(taxon_ids) if taxon_ids else None)
        
        LOG.info(f"✅ Successfully created taxonomy file with {count} taxa")
        return True
        
    except Exception as e:
        LOG.error(f"❌ Failed to download/create taxonomy: {e}")
        return False


def get_organism_ranks(
    keg_file: str,
    taxon_file: str,
    output_file: str,
    download_tax: bool = False
) -> int:
    """
    Get KEGG organism classification ranks.
    """
    keg_path = Path(keg_file)
    taxon_path = Path(taxon_file)
    output_path = Path(output_file)
    
    # Validate inputs
    if not keg_path.exists():
        LOG.error(f"KEGG organisms file not found: {keg_path}")
        return 1
    
    # Download taxonomy if needed
    if download_tax or not taxon_path.exists():
        LOG.info("Need taxonomy database...")
        
        # Get taxon IDs from KEGG file
        org_taxon_map = parse_kegg_organisms_file(keg_path)
        taxon_ids = set(org_taxon_map.values())
        
        if not download_taxonomy_if_needed(taxon_path, taxon_ids):
            LOG.error("Taxonomy download failed")
            return 1
    
    LOG.info(f"🍹 Getting organism classification ranks")
    LOG.info(f"📁 KEGG organisms: {keg_path}")
    LOG.info(f"📁 NCBI taxonomy: {taxon_path}")
    LOG.info(f"📁 Output: {output_path}")
    
    # Parse files
    org_taxon_map = parse_kegg_organisms_file(keg_path)
    taxonomy_data = parse_ncbi_taxon_file(taxon_path)
    
    if not org_taxon_map:
        LOG.error("No organisms found in KEGG file")
        return 1
    
    if not taxonomy_data:
        LOG.error("No taxonomy data loaded")
        return 1
    
    # Get ranks for each organism
    results = []
    found_count = 0
    not_found_count = 0
    
    for org_code, taxon_id in org_taxon_map.items():
        if taxon_id in taxonomy_data:
            taxon_info = taxonomy_data[taxon_id]
            
            # Prepare result
            result = {
                'org_code': org_code,
                'taxon_id': taxon_id,
                'kingdom': taxon_info.get('kingdom', ''),
                'phylum': taxon_info.get('phylum', ''),
                'class': taxon_info.get('class', ''),
                'order': taxon_info.get('order', ''),
                'family': taxon_info.get('family', ''),
                'genus': taxon_info.get('genus', ''),
                'species': taxon_info.get('species', '')
            }
            
            results.append(result)
            found_count += 1
        else:
            LOG.warning(f"Taxon ID {taxon_id} not found in taxonomy for {org_code}")
            results.append({
                'org_code': org_code,
                'taxon_id': taxon_id,
                'kingdom': '',
                'phylum': '',
                'class': '',
                'order': '',
                'family': '',
                'genus': '',
                'species': ''
            })
            not_found_count += 1
    
    # Save results
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    if polars_utils.POLARS_AVAILABLE:
        import polars as pl
        
        # Create DataFrame
        df = pl.DataFrame(results)
        
        # Determine output format
        if output_path.suffix == '.parquet':
            df.write_parquet(output_path)
        elif output_path.suffix == '.csv':
            df.write_csv(output_path)
        else:
            # Default to TSV
            df.write_csv(output_path, separator='\t')
        
        LOG.info(f"💾 Saved {len(results)} records to {output_path}")
    else:
        # Manual TSV writing
        with open(output_path, 'w', encoding='utf-8') as f:
            # Write header
            if results:
                headers = results[0].keys()
                f.write('\t'.join(headers) + '\n')
                
                # Write data
                for result in results:
                    row = [str(result.get(h, '')) for h in headers]
                    f.write('\t'.join(row) + '\n')
        
        LOG.info(f"💾 Saved {len(results)} records to {output_path} (manual)")
    
    # Create summary
    summary_path = output_path.with_suffix('.summary.txt')
    with open(summary_path, 'w') as f:
        f.write("Organism Classification Summary\n")
        f.write("=" * 50 + "\n")
        f.write(f"KEGG organisms file: {keg_path}\n")
        f.write(f"NCBI taxonomy file: {taxon_path}\n")
        f.write(f"Output file: {output_path}\n")
        f.write(f"Download taxonomy: {download_tax}\n\n")
        f.write(f"Total organisms: {len(org_taxon_map)}\n")
        f.write(f"Organisms with taxonomy: {found_count}\n")
        f.write(f"Organisms without taxonomy: {not_found_count}\n")
        
        if not_found_count > 0:
            f.write("\nOrganisms without taxonomy:\n")
            f.write("-" * 40 + "\n")
            for result in results:
                if not result['kingdom']:  # No taxonomy data
                    f.write(f"{result['org_code']}: taxon {result['taxon_id']}\n")
    
    LOG.info(f"📝 Summary saved to: {summary_path}")
    
    # Print final summary
    print("\n" + "="*60)
    print("🍹 ORGANISM CLASSIFICATION COMPLETE")
    print("="*60)
    print(f"🧬 Total organisms: {len(org_taxon_map)}")
    print(f"✅ With taxonomy: {found_count}")
    print(f"❌ Without taxonomy: {not_found_count}")
    print(f"📁 Output file: {output_path}")
    print("="*60)
    
    if not_found_count > 0:
        LOG.warning(f"⚠️  {not_found_count} organisms not found in taxonomy")
    
    return 0


def main():
    """Command-line entry point."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Get KEGG organism classification ranks"
    )
    parser.add_argument(
        "--keg", "-k",
        required=True,
        help="KEGG Organisms htex file"
    )
    parser.add_argument(
        "--taxon", "-t",
        default="ncbi_taxonomy.tsv",
        help="NCBI Taxonomy file (default: ncbi_taxonomy.tsv)"
    )
    parser.add_argument(
        "--out", "-o",
        default=config.get("DEFAULT_RANKS_FILE"),
        help=f"Output file (default: {config.get('DEFAULT_RANKS_FILE')})"
    )
    parser.add_argument(
        "--download-tax",
        action="store_true",
        help="Download NCBI taxonomy if needed"
    )
    parser.add_argument(
        "--force-download",
        action="store_true",
        help="Force download of taxonomy even if file exists"
    )
    parser.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Verbose output"
    )
    
    args = parser.parse_args()
    
    # Set up logging
    log_level = "DEBUG" if args.verbose else config.get("LOG_LEVEL")
    logging.basicConfig(
        level=getattr(logging, log_level.upper()),
        format=config.get("LOG_FORMAT")
    )
    
    # Get organism ranks
    return get_organism_ranks(
        keg_file=args.keg,
        taxon_file=args.taxon,
        output_file=args.out,
        download_tax=args.download_tax or args.force_download
    )


if __name__ == "__main__":
    sys.exit(main())
