"""
NCBI Taxonomy utilities for gpathex
"""
import gzip
import logging
import tarfile
import tempfile
from pathlib import Path
from typing import Dict, List, Optional, Set
import urllib.request

from .polars_utils import POLARS_AVAILABLE

LOG = logging.getLogger(__name__)

TAXDUMP_URL = "https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/taxdump.tar.gz"


class TaxonomyProcessor:
    """Process NCBI taxonomy dumps."""
    
    def __init__(self, data_dir: Optional[Path] = None):
        self.data_dir = data_dir or Path.cwd() / "taxonomy_data"
        self.data_dir.mkdir(parents=True, exist_ok=True)
        self.nodes = {}  # tax_id -> {'parent': parent_id, 'rank': rank, 'name': name}
        self.names = {}  # tax_id -> scientific_name
        self.lineages = {}  # tax_id -> [rank1, rank2, ...]
    
    def download_taxdump(self, force: bool = False) -> Path:
        """Download NCBI taxdump."""
        tar_path = self.data_dir / "taxdump.tar.gz"
        
        if tar_path.exists() and not force:
            LOG.info(f"Taxdump already exists: {tar_path}")
            return tar_path
        
        LOG.info(f"Downloading NCBI taxonomy from {TAXDUMP_URL}")
        LOG.info("This may take a few minutes (approx. 100MB download)...")
        
        try:
            urllib.request.urlretrieve(
                TAXDUMP_URL,
                tar_path,
                reporthook=self._download_progress
            )
            LOG.info(f"\nDownload complete: {tar_path}")
            return tar_path
            
        except Exception as e:
            LOG.error(f"Failed to download taxonomy: {e}")
            raise
    
    def _download_progress(self, block_num, block_size, total_size):
        """Show download progress."""
        downloaded = block_num * block_size
        percent = downloaded * 100 / total_size if total_size > 0 else 0
        
        if block_num % 100 == 0:  # Update every 100 blocks
            mb_downloaded = downloaded / (1024 * 1024)
            mb_total = total_size / (1024 * 1024) if total_size > 0 else 0
            LOG.info(f"  Downloaded {mb_downloaded:.1f}/{mb_total:.1f} MB ({percent:.1f}%)")
    
    def extract_taxdump(self, tar_path: Optional[Path] = None) -> None:
        """Extract taxdump files."""
        if tar_path is None:
            tar_path = self.data_dir / "taxdump.tar.gz"
        
        if not tar_path.exists():
            raise FileNotFoundError(f"Taxdump not found: {tar_path}")
        
        LOG.info(f"Extracting taxonomy files from {tar_path}")
        
        with tarfile.open(tar_path, 'r:gz') as tar:
            # Extract only needed files
            needed_files = ['nodes.dmp', 'names.dmp', 'merged.dmp']
            
            for member in tar.getmembers():
                if member.name in needed_files:
                    tar.extract(member, self.data_dir)
                    LOG.debug(f"Extracted {member.name}")
        
        LOG.info("Extraction complete")
    
    def parse_nodes(self) -> None:
        """Parse nodes.dmp file."""
        nodes_file = self.data_dir / "nodes.dmp"
        
        if not nodes_file.exists():
            raise FileNotFoundError(f"nodes.dmp not found: {nodes_file}")
        
        LOG.info(f"Parsing taxonomy nodes from {nodes_file}")
        
        with open(nodes_file, 'r', encoding='utf-8') as f:
            for line in f:
                parts = [p.strip() for p in line.strip().split('|')]
                if len(parts) >= 13:
                    tax_id = parts[0]
                    parent_id = parts[1]
                    rank = parts[2]
                    
                    self.nodes[tax_id] = {
                        'parent': parent_id,
                        'rank': rank,
                        'name': ''  # Will be filled from names.dmp
                    }
        
        LOG.info(f"Parsed {len(self.nodes)} taxonomy nodes")
    
    def parse_names(self) -> None:
        """Parse names.dmp file for scientific names."""
        names_file = self.data_dir / "names.dmp"
        
        if not names_file.exists():
            raise FileNotFoundError(f"names.dmp not found: {names_file}")
        
        LOG.info(f"Parsing scientific names from {names_file}")
        
        scientific_names = {}
        
        with open(names_file, 'r', encoding='utf-8') as f:
            for line in f:
                parts = [p.strip() for p in line.strip().split('|')]
                if len(parts) >= 4:
                    tax_id = parts[0]
                    name = parts[1]
                    name_class = parts[3]
                    
                    if name_class == 'scientific name':
                        scientific_names[tax_id] = name
        
        # Update nodes with scientific names
        for tax_id, name in scientific_names.items():
            if tax_id in self.nodes:
                self.nodes[tax_id]['name'] = name
        
        LOG.info(f"Found {len(scientific_names)} scientific names")
    
    def build_lineages(self) -> None:
        """Build complete lineages for all taxa."""
        LOG.info("Building taxonomic lineages...")
        
        for tax_id in self.nodes:
            self._get_lineage(tax_id)
        
        LOG.info(f"Built lineages for {len(self.lineages)} taxa")
    
    def _get_lineage(self, tax_id: str) -> List[str]:
        """Get lineage for a taxon ID (with caching)."""
        if tax_id in self.lineages:
            return self.lineages[tax_id]
        
        lineage = []
        current_id = tax_id
        
        while current_id and current_id in self.nodes:
            node = self.nodes[current_id]
            
            # Add this taxon to lineage if it has a name and rank
            if node['name'] and node['rank']:
                # Map NCBI ranks to standard ranks
                rank = self._normalize_rank(node['rank'])
                if rank:
                    lineage.append(node['name'])
            
            # Move to parent
            if node['parent'] != current_id:  # Avoid infinite loop
                current_id = node['parent']
            else:
                break
        
        # Reverse lineage (root to species)
        lineage.reverse()
        self.lineages[tax_id] = lineage
        return lineage
    
    def _normalize_rank(self, rank: str) -> str:
        """Normalize NCBI ranks to standard taxonomy ranks."""
        rank_map = {
            'superkingdom': 'kingdom',
            'phylum': 'phylum',
            'class': 'class',
            'order': 'order',
            'family': 'family',
            'genus': 'genus',
            'species': 'species',
            'subspecies': 'subspecies',
            'varietas': 'variety',
            'forma': 'form'
        }
        return rank_map.get(rank.lower(), '')
    
    def create_taxonomy_tsv(self, output_file: Path, selected_taxa: Optional[Set[str]] = None) -> int:
        """
        Create taxonomy TSV file in the format:
        tax_id kingdom phylum class order family genus species
        
        Args:
            output_file: Output TSV file path
            selected_taxa: Set of taxon IDs to include (None for all)
        
        Returns:
            Number of taxa written
        """
        LOG.info(f"Creating taxonomy TSV: {output_file}")
        
        # Define standard rank order
        standard_ranks = ['kingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species']
        
        # Count stats
        written = 0
        incomplete = 0
        
        with open(output_file, 'w', encoding='utf-8') as f:
            # Write header
            f.write('\t'.join(['tax_id'] + standard_ranks) + '\n')
            
            # Process taxa
            tax_ids = selected_taxa if selected_taxa else self.nodes.keys()
            
            for tax_id in sorted(tax_ids, key=int):
                if tax_id not in self.nodes:
                    continue
                
                # Get lineage
                lineage = self._get_lineage(tax_id)
                
                # Skip if lineage is too short or empty
                if not lineage or len(lineage) < 2:
                    incomplete += 1
                    continue
                
                # Map lineage to standard ranks
                # For now, just write the lineage as-is
                # (advanced: map to specific ranks based on rank information)
                
                # Write tax_id followed by lineage entries
                # Pad with empty strings if lineage is shorter than standard ranks
                output_line = [tax_id] + lineage
                
                # Ensure we have exactly 8 columns (tax_id + 7 ranks)
                if len(output_line) < 8:
                    output_line += [''] * (8 - len(output_line))
                elif len(output_line) > 8:
                    output_line = output_line[:8]  # Truncate if too long
                
                f.write('\t'.join(output_line) + '\n')
                written += 1
        
        LOG.info(f"Wrote {written} taxa to {output_file}")
        if incomplete > 0:
            LOG.warning(f"Skipped {incomplete} taxa with incomplete lineage")
        
        return written
    
    def get_taxon_lineage(self, tax_id: str) -> Dict[str, str]:
        """Get lineage for a specific taxon ID as a dictionary of rank: name."""
        lineage = self._get_lineage(tax_id)
        
        # Map to standard ranks if possible
        standard_ranks = ['kingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species']
        
        result = {'taxon_id': tax_id}
        
        # If we have enough lineage items, map them to ranks
        if len(lineage) >= len(standard_ranks):
            for i, rank in enumerate(standard_ranks):
                if i < len(lineage):
                    result[rank] = lineage[i]
        else:
            # Just number them
            for i, name in enumerate(lineage):
                result[f'rank_{i}'] = name
        
        return result


def download_and_create_taxonomy(output_file: str, selected_taxa: Optional[List[str]] = None) -> int:
    """
    Convenience function: download and create taxonomy TSV in one step.
    
    Args:
        output_file: Output TSV file path
        selected_taxa: List of taxon IDs to include (None for all)
    
    Returns:
        Number of taxa written
    """
    processor = TaxonomyProcessor()
    
    try:
        # Download
        tar_path = processor.download_taxdump()
        
        # Extract
        processor.extract_taxdump(tar_path)
        
        # Parse
        processor.parse_nodes()
        processor.parse_names()
        processor.build_lineages()
        
        # Create TSV
        taxa_set = set(selected_taxa) if selected_taxa else None
        return processor.create_taxonomy_tsv(Path(output_file), taxa_set)
        
    except Exception as e:
        LOG.error(f"Failed to create taxonomy file: {e}")
        raise


__all__ = [
    'TaxonomyProcessor',
    'download_and_create_taxonomy',
    'TAXDUMP_URL'
]

def create_comprehensive_taxonomy_tsv(self, output_file: Path, target_taxon_ids: Set[str]) -> int:
    """
    Create comprehensive taxonomy TSV including all target taxon IDs.
    Builds full taxonomy tree for missing IDs.
    
    Args:
        output_file: Output TSV file path
        target_taxon_ids: Set of taxon IDs to include
    
    Returns:
        Number of taxa written
    """
    LOG.info(f"Creating comprehensive taxonomy TSV for {len(target_taxon_ids)} target taxa")
    
    # First, ensure we have all target taxon IDs in our database
    missing_ids = [tid for tid in target_taxon_ids if tid not in self.nodes]
    
    if missing_ids:
        LOG.info(f"{len(missing_ids)} target taxon IDs not in local database")
        LOG.info("Attempting to fetch missing taxa from NCBI...")
        
        # Try to fetch missing taxa via NCBI API
        fetched_count = self._fetch_missing_taxa(missing_ids[:1000])  # Limit to 1000
        
        if fetched_count > 0:
            LOG.info(f"Fetched {fetched_count} missing taxa from NCBI")
    
    # Define standard rank order
    standard_ranks = ['kingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species']
    
    written = 0
    with open(output_file, 'w', encoding='utf-8') as f:
        # Write header
        f.write('\t'.join(['tax_id'] + standard_ranks) + '\n')
        
        # Process target taxon IDs
        for tax_id in sorted(target_taxon_ids, key=lambda x: int(x) if x.isdigit() else 0):
            if tax_id not in self.nodes:
                LOG.debug(f"Taxon ID {tax_id} still not in database")
                continue
            
            # Get lineage
            lineage = self._get_lineage(tax_id)
            
            # Write tax_id followed by lineage entries
            output_line = [tax_id] + lineage
            
            # Ensure we have exactly 8 columns (tax_id + 7 ranks)
            if len(output_line) < 8:
                output_line += [''] * (8 - len(output_line))
            elif len(output_line) > 8:
                output_line = output_line[:8]  # Truncate if too long
            
            f.write('\t'.join(output_line) + '\n')
            written += 1
    
    LOG.info(f"Wrote {written} taxa to {output_file}")
    return written

def _fetch_missing_taxa(self, taxon_ids: List[str]) -> int:
    """Fetch missing taxa from NCBI API."""
    import requests
    import time
    
    fetched = 0
    
    for tax_id in taxon_ids:
        try:
            url = f"https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
            params = {
                'db': 'taxonomy',
                'id': tax_id,
                'retmode': 'xml'
            }
            
            response = requests.get(url, params=params, timeout=30)
            if response.status_code == 200:
                # Parse XML response (simplified)
                # This would need actual XML parsing
                # For now, just mark as fetched
                self.nodes[tax_id] = {
                    'parent': '1',  # Default to root
                    'rank': 'no rank',
                    'name': f"Taxon_{tax_id}"
                }
                fetched += 1
            
            # Be nice to NCBI API
            time.sleep(0.1)
            
        except Exception as e:
            LOG.debug(f"Failed to fetch taxon {tax_id}: {e}")
    
    return fetched

def download_and_create_taxonomy_for_kegg(kegg_file: str, output_file: str, force: bool = False) -> int:
    """
    Create taxonomy TSV specifically for KEGG organisms.
    
    Args:
        kegg_file: Path to KEGG organisms .keg file
        output_file: Output TSV file path
        force: Force download even if file exists
    
    Returns:
        Number of taxa written
    """
    from ..commands.get_ranks import parse_kegg_organisms_file
    
    kegg_path = Path(kegg_file)
    output_path = Path(output_file)
    
    if output_path.exists() and not force:
        LOG.info(f"Taxonomy file already exists: {output_path}")
        return 1
    
    # Parse KEGG file to get required taxon IDs
    LOG.info(f"Parsing KEGG file to get required taxon IDs: {kegg_path}")
    org_taxon_map = parse_kegg_organisms_file(kegg_path)
    
    if not org_taxon_map:
        LOG.error("No organisms found in KEGG file")
        return 0
    
    taxon_ids = set(org_taxon_map.values())
    LOG.info(f"Found {len(taxon_ids)} unique taxon IDs in KEGG file")
    
    # Create processor
    processor = TaxonomyProcessor()
    
    try:
        # Download taxdump
        tar_path = processor.download_taxdump(force)
        
        # Extract
        processor.extract_taxdump(tar_path)
        
        # Parse
        processor.parse_nodes()
        processor.parse_names()
        processor.build_lineages()
        
        # Create comprehensive TSV
        return processor.create_comprehensive_taxonomy_tsv(output_path, taxon_ids)
        
    except Exception as e:
        LOG.error(f"Failed to create taxonomy file: {e}")
        raise
        
def create_taxonomy_from_api(taxon_ids: Set[str], output_file: Path) -> int:
    """
    Create taxonomy file directly from NCBI API.
    This is slower but more comprehensive.
    """
    import requests
    import xml.etree.ElementTree as ET
    import time
    
    LOG.info(f"Creating taxonomy from NCBI API for {len(taxon_ids)} taxon IDs")
    
    results = {}
    
    # Process in batches
    batch_size = 100
    taxon_list = list(taxon_ids)
    
    for i in range(0, len(taxon_list), batch_size):
        batch = taxon_list[i:i+batch_size]
        batch_str = ",".join(batch)
        
        try:
            url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
            params = {
                'db': 'taxonomy',
                'id': batch_str,
                'retmode': 'xml'
            }
            
            response = requests.get(url, params=params, timeout=60)
            if response.status_code != 200:
                LOG.warning(f"API request failed: {response.status_code}")
                continue
            
            root = ET.fromstring(response.content)
            
            for taxon in root.findall('.//Taxon'):
                tax_id = taxon.find('TaxId').text
                sci_name = taxon.find('ScientificName').text
                rank = taxon.find('Rank').text
                
                # Get lineage
                lineage_map = {}
                for lineage_taxon in taxon.find('LineageEx').findall('Taxon'):
                    l_rank = lineage_taxon.find('Rank').text
                    l_name = lineage_taxon.find('ScientificName').text
                    if l_rank and l_name:
                        lineage_map[l_rank.lower()] = l_name
                
                # Map to standard ranks
                standard_ranks = ['superkingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species']
                lineage = [lineage_map.get(rank, '') for rank in standard_ranks]
                
                # Ensure species is the organism name
                if not lineage[-1] and sci_name:
                    lineage[-1] = sci_name
                
                results[tax_id] = lineage
            
            LOG.info(f"Processed batch {i//batch_size + 1}/{(len(taxon_list)+batch_size-1)//batch_size}")
            
            # Rate limiting
            time.sleep(0.5)
            
        except Exception as e:
            LOG.warning(f"Error processing batch: {e}")
            continue
    
    # Write results
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write('\t'.join(['tax_id'] + standard_ranks) + '\n')
        
        for tax_id in sorted(taxon_ids):
            if tax_id in results:
                lineage = results[tax_id]
                f.write(f"{tax_id}\t" + '\t'.join(lineage) + '\n')
            else:
                # Write empty entry
                f.write(f"{tax_id}\t" + '\t\t\t\t\t\t\t\n')
    
    LOG.info(f"Wrote {len(results)} taxa to {output_file}")
    return len(results)
    