"""
KEGG API utilities using BioPython
"""
"""
KEGG API utilities with proper endpoints
"""
import logging
import time
from typing import Dict, List, Optional, Union
from pathlib import Path
import urllib.request
import urllib.error

from .polars_utils import BIOPYTHON_AVAILABLE, POLARS_AVAILABLE

LOG = logging.getLogger(__name__)

# KEGG REST API base URL
KEGG_REST_URL = "https://rest.kegg.jp"


def kegg_api_get(
    endpoint: str,
    database: Optional[str] = None,
    identifier: Optional[str] = None,
    max_retries: int = 3,
    timeout: int = 30
) -> str:
    """
    Get data from KEGG REST API with proper endpoint construction.
    
    Args:
        endpoint: API endpoint (list, get, find, conv, link)
        database: KEGG database name
        identifier: Entry identifier
        max_retries: Maximum retry attempts
        timeout: Request timeout
    
    Returns:
        Raw response text
    
    Raises:
        urllib.error.HTTPError: If request fails
    """
    # Construct URL based on endpoint
    if endpoint == "list":
        if database:
            url = f"{KEGG_REST_URL}/list/{database}"
        else:
            raise ValueError("Database name required for list endpoint")
    
    elif endpoint == "get":
        if identifier:
            url = f"{KEGG_REST_URL}/get/{identifier}"
        else:
            raise ValueError("Identifier required for get endpoint")
    
    elif endpoint == "find":
        if database and identifier:
            url = f"{KEGG_REST_URL}/find/{database}/{identifier}"
        else:
            raise ValueError("Database and identifier required for find endpoint")
    
    elif endpoint == "conv":
        if database and identifier:
            url = f"{KEGG_REST_URL}/conv/{database}/{identifier}"
        else:
            raise ValueError("Target and source databases required for conv endpoint")
    
    elif endpoint == "link":
        if database and identifier:
            url = f"{KEGG_REST_URL}/link/{database}/{identifier}"
        else:
            raise ValueError("Target and source databases required for link endpoint")
    
    else:
        raise ValueError(f"Unknown endpoint: {endpoint}")
    
    LOG.debug(f"KEGG API request: {url}")
    
    # Try BioPython first
    if BIOPYTHON_AVAILABLE and endpoint == "get":
        try:
            from Bio.KEGG import REST
            LOG.debug(f"Using Bio.KEGG.REST for {identifier}")
            
            for attempt in range(max_retries):
                try:
                    response = REST.kegg_get(identifier)
                    return response.read().decode('utf-8')
                except Exception as e:
                    if attempt < max_retries - 1:
                        wait = 2 ** attempt
                        LOG.warning(f"Bio.KEGG.REST attempt {attempt + 1} failed: {e}. Retrying in {wait}s...")
                        time.sleep(wait)
                    else:
                        LOG.warning(f"Bio.KEGG.REST failed for {identifier}: {e}. Falling back to HTTP.")
                        raise
            
        except Exception as e:
            LOG.debug(f"Bio.KEGG.REST failed, using direct HTTP: {e}")
    
    # Direct HTTP request with retries
    last_error = None
    
    for attempt in range(max_retries):
        try:
            # Add delay between retries
            if attempt > 0:
                wait_time = 2 ** attempt
                time.sleep(wait_time)
            
            headers = {
                'User-Agent': 'gpathex /1.0.0 (+https://github.com/yourusername/gpathex )'
            }
            
            req = urllib.request.Request(url, headers=headers)
            
            with urllib.request.urlopen(req, timeout=timeout) as response:
                return response.read().decode('utf-8')
                
        except urllib.error.HTTPError as e:
            last_error = e
            if e.code == 400:
                LOG.warning(f"HTTP 400 Bad Request for {url}")
                # Don't retry 400 errors (bad request)
                break
            elif attempt < max_retries - 1:
                wait = 2 ** attempt
                LOG.warning(f"HTTP attempt {attempt + 1} failed ({e.code}): {e.reason}. Retrying in {wait}s...")
            else:
                LOG.error(f"All {max_retries} HTTP attempts failed for {url}")
        
        except (urllib.error.URLError, TimeoutError) as e:
            last_error = e
            if attempt < max_retries - 1:
                wait = 2 ** attempt
                LOG.warning(f"Network attempt {attempt + 1} failed: {e}. Retrying in {wait}s...")
            else:
                LOG.error(f"All {max_retries} network attempts failed for {url}")
        
        except Exception as e:
            last_error = e
            if attempt < max_retries - 1:
                wait = 2 ** attempt
                LOG.warning(f"Attempt {attempt + 1} failed: {e}. Retrying in {wait}s...")
            else:
                LOG.error(f"All {max_retries} attempts failed for {url}")
    
    if last_error:
        raise last_error
    else:
        raise Exception(f"Failed to fetch {url} after {max_retries} attempts")


def kegg_list(
    database: str,
    org: Optional[str] = None
) -> List[List[str]]:
    """
    List entries in a KEGG database.
    
    Args:
        database: KEGG database name (e.g., 'organism', 'ko', 'pathway')
        org: Organism code for organism-specific listings
    
    Returns:
        List of [entry, description] pairs
    """
    try:
        # For organism list, use specific endpoint
        if database == "organism":
            response = kegg_api_get("list", database="organism")
        elif org:
            # Organism-specific listing (e.g., hsa for human genes)
            response = kegg_api_get("list", database=org)
        else:
            # Generic database listing
            response = kegg_api_get("list", database=database)
        
        lines = response.strip().split('\n')
        
        result = []
        for line in lines:
            if line.strip():
                parts = line.split('\t', 1)
                if len(parts) == 2:
                    result.append([parts[0].strip(), parts[1].strip()])
                elif len(parts) == 1:
                    result.append([parts[0].strip(), ''])
        
        return result
        
    except Exception as e:
        LOG.error(f"Failed to list {database}: {e}")
        return []


def kegg_find(
    database: str,
    query: str,
    option: Optional[str] = None
) -> List[List[str]]:
    """
    Find entries in KEGG database.
    
    Args:
        database: KEGG database name
        query: Search query
        option: Search option
    
    Returns:
        List of [entry, description] pairs
    """
    try:
        if option:
            response = kegg_api_get("find", database=f"{database}/{option}", identifier=query)
        else:
            response = kegg_api_get("find", database=database, identifier=query)
        
        lines = response.strip().split('\n')
        
        result = []
        for line in lines:
            if line.strip():
                parts = line.split('\t', 1)
                if len(parts) == 2:
                    result.append([parts[0].strip(), parts[1].strip()])
        
        return result
        
    except Exception as e:
        LOG.error(f"Failed to find '{query}' in {database}: {e}")
        return []


def get_organism_list() -> List[List[str]]:
    """
    Get complete KEGG organism list.
    
    Returns:
        List of [org_code, description] pairs
    """
    try:
        LOG.info("Fetching KEGG organism list via REST API")
        
        # Method 1: Try KEGG REST API
        response = kegg_api_get("list", database="organism")
        
        organisms = []
        for line in response.strip().split('\n'):
            if line.strip():
                parts = line.split('\t')
                if len(parts) >= 2:
                    # Format: T01001  hsa     Homo sapiens (human)  Eukaryotes;Animals;Vertebrates;Mammals
                    org_id = parts[0].strip()  # T-number
                    org_code = parts[1].strip()  # 3-4 letter code
                    org_name = parts[2].strip() if len(parts) > 2 else ""
                    
                    organisms.append([org_code, org_name, org_id])
        
        LOG.info(f"Retrieved {len(organisms)} organisms via KEGG REST API")
        return organisms
        
    except Exception as e:
        LOG.error(f"Failed to get organism list from KEGG API: {e}")
        
        # Fallback: Use HTML parsing
        LOG.info("Falling back to HTML parsing for organism list")
        return get_organism_list_html()


def get_organism_list_html() -> List[List[str]]:
    """
    Fallback method: Parse KEGG organism list from HTML.
    """
    import urllib.request
    from html.parser import HTMLParser
    
    class KEGGOrgHTMLParser(HTMLParser):
        def __init__(self):
            super().__init__()
            self.organisms = []
            self.current_org = {}
            self.in_org_row = False
            self.in_org_link = False
            self.in_name_cell = False
            self.collected_name = False
        
        def handle_starttag(self, tag, attrs):
            attrs_dict = dict(attrs)
            
            if tag == 'tr':
                # Check if this is an organism row
                for attr_name, attr_value in attrs:
                    if attr_name == 'onclick' and 'show_organism' in attr_value:
                        self.in_org_row = True
                        self.current_org = {}
            
            elif tag == 'a' and self.in_org_row:
                href = attrs_dict.get('href', '')
                if 'show_organism?org=' in href:
                    self.in_org_link = True
                    # Extract organism code from href
                    org_code = href.split('org=')[1].split("'")[0]
                    self.current_org['code'] = org_code
            
            elif tag == 'td' and self.in_org_row and not self.collected_name:
                self.in_name_cell = True
        
        def handle_data(self, data):
            if self.in_name_cell and data.strip():
                if 'name' not in self.current_org:
                    self.current_org['name'] = data.strip()
                    self.collected_name = True
        
        def handle_endtag(self, tag):
            if tag == 'tr' and self.in_org_row:
                if self.current_org:
                    self.organisms.append([
                        self.current_org.get('code', ''),
                        self.current_org.get('name', ''),
                        ''  # No T-number from HTML
                    ])
                self.in_org_row = False
                self.in_org_link = False
                self.in_name_cell = False
                self.collected_name = False
                self.current_org = {}
    
    try:
        url = "https://www.kegg.jp/kegg/catalog/org_list.html"
        
        req = urllib.request.Request(
            url,
            headers={
                'User-Agent': 'Mozilla/5.0 (gpathex /1.0.0)'
            }
        )
        
        with urllib.request.urlopen(req, timeout=30) as response:
            html_content = response.read().decode('utf-8', errors='ignore')
        
        parser = KEGGOrgHTMLParser()
        parser.feed(html_content)
        
        LOG.info(f"Parsed {len(parser.organisms)} organisms from HTML")
        return parser.organisms
        
    except Exception as e:
        LOG.error(f"Failed to parse organism list from HTML: {e}")
        return []


def get_organism_info(org_code: str) -> Dict[str, str]:
    """
    Get detailed organism information from KEGG organism page.
    
    Args:
        org_code: KEGG organism code (e.g., "agw")
    
    Returns:
        Dictionary with organism information
    """
    import urllib.request
    import re
    
    url = f"https://www.kegg.jp/kegg-bin/show_organism?org={org_code}"
    
    try:
        req = urllib.request.Request(
            url,
            headers={'User-Agent': 'gpathex/1.0.0'}
        )
        
        with urllib.request.urlopen(req, timeout=30) as response:
            html_content = response.read().decode('utf-8', errors='ignore')
            
            info = {
                'org_code': org_code,
                'taxon_id': '',
                'name': '',
                'genbank_accession': '',
                'assembly_accession': '',
                'bioproject': '',
                'chromosome': '',
                'protein_genes': '',
                'rna_genes': ''
            }
            
            # Extract T-number (Entry)
            entry_match = re.search(r'Entry\s*</td><td[^>]*>([^<]+)', html_content)
            if entry_match:
                info['taxon_id'] = entry_match.group(1).strip()
            
            # Extract Name
            name_match = re.search(r'Name\s*</td><td[^>]*>([^<]+)', html_content)
            if name_match:
                info['name'] = name_match.group(1).strip()
            
            # Extract Data source line
            data_source_match = re.search(r'Data source\s*</td><td[^>]*>([^<]+)', html_content)
            if data_source_match:
                data_source = data_source_match.group(1).strip()
                
                # Extract GenBank Assembly
                gca_match = re.search(r'GCA_[\d\.]+', data_source)
                if gca_match:
                    info['assembly_accession'] = gca_match.group(0)
                
                # Extract BioProject
                bioproject_match = re.search(r'BioProject:\s*(\d+)', data_source)
                if bioproject_match:
                    info['bioproject'] = bioproject_match.group(1)
            
            # Extract Chromosome/GenBank accession
            chromosome_match = re.search(r'Chromosome\s*</td><td[^>]*>([^<]+)', html_content)
            if chromosome_match:
                chromosome = chromosome_match.group(1).strip()
                # Look for GB: CP... or GenBank: ...
                gb_match = re.search(r'GB:\s*([A-Z]+\d+)', chromosome)
                if gb_match:
                    info['genbank_accession'] = gb_match.group(1)
                else:
                    # Try other patterns
                    acc_match = re.search(r'([A-Z]+\d+(\.\d+)?)', chromosome)
                    if acc_match:
                        info['genbank_accession'] = acc_match.group(1)
            
            # Extract statistics
            stats_match = re.search(r'Statistics\s*</td><td[^>]*>([^<]+)', html_content)
            if stats_match:
                stats = stats_match.group(1).strip()
                # Number of protein genes
                protein_match = re.search(r'Number of protein genes:\s*(\d+)', stats)
                if protein_match:
                    info['protein_genes'] = protein_match.group(1)
                
                # Number of RNA genes
                rna_match = re.search(r'Number of RNA genes:\s*(\d+)', stats)
                if rna_match:
                    info['rna_genes'] = rna_match.group(1)
            
            return info
            
    except Exception as e:
        LOG.error(f"Failed to get organism info for {org_code}: {e}")
        return {
            'org_code': org_code,
            'taxon_id': '',
            'name': '',
            'genbank_accession': '',
            'assembly_accession': '',
            'bioproject': '',
            'chromosome': '',
            'protein_genes': '',
            'rna_genes': ''
        }

def construct_ncbi_urls(organism_info: Dict[str, str]) -> Dict[str, str]:
    """
    Construct NCBI URLs from organism information.
    
    Args:
        organism_info: Dictionary from get_organism_info()
    
    Returns:
        Dictionary with various NCBI URLs
    """
    urls = {
        'genbank_nucleotide_url': '',
        'genome_assembly_url': '',
        'bioproject_url': '',
        'ftp_url': ''
    }
    
    org_code = organism_info.get('org_code', '')
    genbank_acc = organism_info.get('genbank_accession', '')
    assembly_acc = organism_info.get('assembly_accession', '')
    bioproject = organism_info.get('bioproject', '')
    
    # Construct GenBank nucleotide URL
    if genbank_acc:
        urls['genbank_nucleotide_url'] = f"https://www.ncbi.nlm.nih.gov/nuccore/{genbank_acc}"
    
    # Construct Genome assembly URL
    if assembly_acc:
        if assembly_acc.startswith('GCA_'):
            urls['genome_assembly_url'] = f"https://www.ncbi.nlm.nih.gov/datasets/genome/{assembly_acc}/"
        elif assembly_acc.startswith('GCF_'):
            urls['genome_assembly_url'] = f"https://www.ncbi.nlm.nih.gov/datasets/genome/{assembly_acc}/"
    
    # Construct BioProject URL
    if bioproject:
        urls['bioproject_url'] = f"https://www.ncbi.nlm.nih.gov/bioproject/{bioproject}"
    
    # Construct FTP URL from assembly accession
    if assembly_acc:
        if assembly_acc.startswith('GCA_'):
            # GenBank assembly
            acc_num = assembly_acc.replace('GCA_', '')
            num_part = acc_num.split('.')[0]
            if len(num_part) >= 9:
                dir1 = num_part[:3]
                dir2 = num_part[3:6]
                dir3 = num_part[6:9]
                urls['ftp_url'] = f"ftp://ftp.ncbi.nlm.nih.gov/genomes/all/GCA/{dir1}/{dir2}/{dir3}/{assembly_acc}"
        elif assembly_acc.startswith('GCF_'):
            # RefSeq assembly
            acc_num = assembly_acc.replace('GCF_', '')
            num_part = acc_num.split('.')[0]
            if len(num_part) >= 9:
                dir1 = num_part[:3]
                dir2 = num_part[3:6]
                dir3 = num_part[6:9]
                urls['ftp_url'] = f"ftp://ftp.ncbi.nlm.nih.gov/genomes/all/GCF/{dir1}/{dir2}/{dir3}/{assembly_acc}"
    
    return urls


# TODO: FOR NOW I NEED SOMETHING FOR THIS TOO COMPLICATED THING!
# def get_ncbi_ftp_url(taxon_id: str, organism_name: str = "", org_code: str = "") -> str:
    # """
    # Get NCBI FTP URL by scraping KEGG organism page.
    
    # Args:
        # taxon_id: KEGG T-number (optional)
        # organism_name: Organism name (optional)
        # org_code: KEGG organism code (required)
    
    # Returns:
        # NCBI FTP URL or empty string
    # """
    # if not org_code:
        # LOG.error("org_code is required for get_ncbi_ftp_url")
        # return ""
    
    # try:
        # # Get detailed organism info
        # org_info = get_organism_info(org_code)
        
        # # Construct URLs
        # urls = construct_ncbi_urls(org_info)
        
        # # Return FTP URL if available
        # if urls['ftp_url']:
            # LOG.debug(f"Found FTP URL for {org_code}: {urls['ftp_url']}")
            # return urls['ftp_url']
        
        # # Fallback: Try to construct from assembly accession
        # assembly_acc = org_info.get('assembly_accession', '')
        # if assembly_acc:
            # LOG.debug(f"Using assembly accession for {org_code}: {assembly_acc}")
            # return construct_ncbi_urls(org_info)['ftp_url']
        
        # LOG.warning(f"No FTP URL found for {org_code}")
        # return ""
        
    # except Exception as e:
        # LOG.error(f"Error getting FTP URL for {org_code}: {e}")
        # return ""

def get_ncbi_ftp_url(taxon_id: str) -> str:
    # """
    # Simple get_ncbi_ftp_url - returns empty string.
    # Users can fill this later if needed.
    
    # Args:
        # taxon_id: KEGG T-number (e.g., "T01001")
    
    # Returns:
        # Empty string (placeholder for future enhancement)
    # """
    # # Return empty string - users can fill this manually or via separate tool
    # return ""
    raise NotImplementedError("NAH! Not yet!")

def kegg_conv(
    target_db: str, 
    source_db: str, 
    source_id: Optional[str] = None
) -> Dict[str, List[str]]:
    """
    Convert IDs between KEGG databases.
    
    Args:
        target_db: Target database
        source_db: Source database
        source_id: Source ID (optional, if None converts all)
    
    Returns:
        Dictionary of {source_id: [target_ids]}
    """
    if source_id:
        query = f"{target_db}/{source_db}:{source_id}"
    else:
        query = f"{target_db}/{source_db}"
    
    try:
        response = kegg_api_get(query)
        lines = response.strip().split('\n')
        
        result = {}
        for line in lines:
            if line.strip():
                parts = line.split('\t')
                if len(parts) == 2:
                    source = parts[0].strip()
                    target = parts[1].strip()
                    
                    if source not in result:
                        result[source] = []
                    result[source].append(target)
        
        return result
        
    except Exception as e:
        LOG.error(f"Failed to convert {source_db} to {target_db}: {e}")
        return {}

def kegg_link(
    target_db: str, 
    source_db: str, 
    source_id: Optional[str] = None
) -> Dict[str, List[str]]:
    """
    Link entries between KEGG databases.
    
    Args:
        target_db: Target database
        source_db: Source database
        source_id: Source ID (optional)
    
    Returns:
        Dictionary of {source_id: [target_ids]}
    """
    if source_id:
        query = f"{target_db}/{source_db}:{source_id}"
    else:
        query = f"{target_db}/{source_db}"
    
    try:
        response = kegg_api_get(query)
        lines = response.strip().split('\n')
        
        result = {}
        for line in lines:
            if line.strip():
                parts = line.split('\t')
                if len(parts) == 2:
                    source = parts[0].strip()
                    target = parts[1].strip()
                    
                    if source not in result:
                        result[source] = []
                    result[source].append(target)
        
        return result
        
    except Exception as e:
        LOG.error(f"Failed to link {source_db} to {target_db}: {e}")
        return {}

def parse_kegg_response(
    response: str, 
    entry_type: str = "generic"
) -> Dict:
    """
    Parse KEGG API response.
    
    Args:
        response: Raw response text
        entry_type: Type of entry (ko, compound, pathway, etc.)
    
    Returns:
        Parsed data as dictionary
    """
    lines = response.strip().split('\n')
    data = {}
    current_field = None
    current_values = []
    
    for line in lines:
        line = line.rstrip()
        
        if not line:
            continue
        
        # Field indicators (uppercase at start of line)
        if line[0].isupper() and line[:12].isupper():
            # Save previous field
            if current_field and current_values:
                if current_field in data:
                    # Append if field already exists
                    if isinstance(data[current_field], list):
                        data[current_field].extend(current_values)
                    else:
                        data[current_field] = [data[current_field]] + current_values
                else:
                    if len(current_values) == 1:
                        data[current_field] = current_values[0]
                    else:
                        data[current_field] = current_values
            
            # Start new field
            parts = line.split(None, 1)
            current_field = parts[0]
            current_values = [parts[1].strip()] if len(parts) > 1 else []
        else:
            # Continuation line
            if current_field is not None:
                current_values.append(line.strip())
    
    # Don't forget the last field
    if current_field and current_values:
        if current_field in data:
            if isinstance(data[current_field], list):
                data[current_field].extend(current_values)
            else:
                data[current_field] = [data[current_field]] + current_values
        else:
            if len(current_values) == 1:
                data[current_field] = current_values[0]
            else:
                data[current_field] = current_values
    
    return data


def download_kegg_file(
    kegg_id: str, 
    output_dir: Union[str, Path],
    file_type: str = "keg"
) -> Path:
    """
    Download file from KEGG.
    
    Args:
        kegg_id: KEGG identifier
        output_dir: Output directory
        file_type: File type (keg, kgml, etc.)
    
    Returns:
        Path to downloaded file
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    if file_type == "keg":
        filename = f"{kegg_id}.keg"
        url = f"https://www.kegg.jp/kegg-bin/download_htext?htext={kegg_id}&format=htext&filedir="
    elif file_type == "kgml":
        filename = f"{kegg_id}.kgml"
        url = f"https://www.kegg.jp/kegg-bin/download?format=kgml&entry={kegg_id}"
    else:
        raise ValueError(f"Unsupported file type: {file_type}")
    
    output_path = output_dir / filename
    
    if output_path.exists():
        LOG.info(f"File already exists: {output_path}")
        return output_path
    
    try:
        LOG.info(f"Downloading {kegg_id} from KEGG")
        response = urllib.request.urlopen(url)
        
        with open(output_path, 'wb') as f:
            f.write(response.read())
        
        LOG.info(f"Downloaded to {output_path}")
        return output_path
        
    except Exception as e:
        LOG.error(f"Failed to download {kegg_id}: {e}")
        raise
		