#!/usr/bin/env python3
"""
Download protein sequences from NCBI
"""
import sys
import logging
import time
from pathlib import Path
from typing import List, Dict
from concurrent.futures import ThreadPoolExecutor, as_completed
import urllib.request
import urllib.error

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

LOG = logging.getLogger(__name__)


def parse_organism_file(org_file: Path) -> List[Dict]:
    """
    Parse organism file to get URLs for protein downloads.
    
    Expected format: org_code, name, ncbi_ftp_url
    """
    organisms = []
    
    with open(org_file, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            
            parts = line.split('\t')
            if len(parts) >= 3:
                org_code = parts[0]
                name = parts[1]
                url = parts[2]
                
                # Only include organisms with NCBI FTP URLs
                if "ftp.ncbi.nlm.nih.gov" in url:
                    organisms.append({
                        'org_code': org_code,
                        'name': name,
                        'url': url.rstrip('/'),  # Remove trailing slash
                        'line_number': line_num
                    })
                else:
                    LOG.debug(f"Skipping {org_code}: Not an NCBI FTP URL")
    
    LOG.info(f"Found {len(organisms)} organisms with NCBI FTP URLs")
    return organisms


def download_protein_file(
    organism: Dict,
    output_dir: Path,
    format: str = "fasta.gz",
    max_retries: int = 3,
    skip_existing: bool = True
) -> Dict:
    """
    Download protein file for a single organism.
    
    Returns:
        Dictionary with download result
    """
    org_code = organism['org_code']
    ftp_url = organism['url']
    
    result = {
        'org_code': org_code,
        'name': organism['name'],
        'success': False,
        'error': None,
        'filepath': None,
        'size_bytes': 0,
        'attempts': 0
    }
    
    # Extract base directory name from FTP URL
    # Format: ftp://ftp.ncbi.nlm.nih.gov/genomes/all/GCF/XXX/XXX/XXX/
    url_parts = ftp_url.split('/')
    if len(url_parts) < 8:
        result['error'] = f"Invalid FTP URL format: {ftp_url}"
        return result
    
    # Get the last part (usually accession number)
    accession = url_parts[-1] if url_parts[-1] else url_parts[-2]
    
    # Construct protein file URL
    if format == "fasta.gz":
        filename = f"{accession}_translated_cds.faa.gz"
    else:
        filename = f"{accession}_translated_cds.faa"
    
    protein_url = f"{ftp_url}/{filename}"
    output_filename = f"{org_code}.pep.{format}"
    output_path = output_dir / output_filename
    
    # Check if file exists
    if skip_existing and output_path.exists():
        result['success'] = True
        result['filepath'] = str(output_path)
        result['size_bytes'] = output_path.stat().st_size
        result['message'] = "File already exists"
        return result
    
    # Try download with retries
    last_error = None
    
    for attempt in range(max_retries):
        result['attempts'] = attempt + 1
        
        try:
            # Add delay between requests
            if attempt > 0:
                delay = 2 ** attempt
                LOG.debug(f"Retry {attempt + 1}/{max_retries} for {org_code}, waiting {delay}s...")
                time.sleep(delay)
            
            LOG.debug(f"Downloading {org_code} proteins (attempt {attempt + 1}/{max_retries})")
            
            # Make request
            req = urllib.request.Request(
                protein_url,
                headers={
                    'User-Agent': 'gpathex /1.0.0'
                }
            )
            
            with urllib.request.urlopen(req, timeout=60) as response:  # 60 second timeout for large files
                content = response.read()
                
                if not content:
                    raise ValueError("Empty response from server")
                
                # Save file
                with open(output_path, 'wb') as f:
                    f.write(content)
                
                result['success'] = True
                result['filepath'] = str(output_path)
                result['size_bytes'] = len(content)
                result['message'] = f"Downloaded {len(content):,} bytes"
                
                LOG.info(f"✅ Downloaded {org_code} proteins ({len(content):,} bytes)")
                return result
        
        except urllib.error.HTTPError as e:
            if e.code == 404:
                result['error'] = f"Protein file not found (404)"
                result['message'] = "Protein file not available for this organism"
                LOG.warning(f"⚠️  Protein file not found for {org_code} (404)")
                break  # Don't retry 404 errors
            elif e.code == 403:
                LOG.error(f"Access forbidden (403) for {url}. Check permissions.")
                break  # Don't retry 403 errors
            elif e.code == 500:
                LOG.error(f"Server error (500) for {url}. Try again later.")    
            else:
                last_error = e
                result['error'] = f"HTTP error {e.code}: {e.reason}"
                LOG.warning(f"HTTP error for {org_code}: {e.code} {e.reason}")
        
        except urllib.error.URLError as e:
            last_error = e
            result['error'] = f"URL error: {e.reason}"
            LOG.warning(f"URL error for {org_code}: {e.reason}")
        
        except TimeoutError as e:
            last_error = e
            result['error'] = f"Timeout after 60 seconds"
            LOG.warning(f"Timeout for {org_code}")
        
        except Exception as e:
            last_error = e
            result['error'] = str(e)
            LOG.warning(f"Error downloading {org_code}: {e}")
    
    # If we get here, all attempts failed
    if last_error:
        result['error'] = str(last_error)
    
    result['message'] = f"Failed after {result['attempts']} attempts"
    LOG.error(f"❌ Failed to download {org_code} proteins: {result['error']}")
    return result


def download_protein_files(
    org_file: str,
    output_dir: str = "proteins",
    concurrent: int = 5,
    format: str = "fasta.gz",
    skip_ncbi_check: bool = False,
    skip_existing: bool = True
) -> int:
    """
    Download protein files for multiple organisms.
    """
    org_path = Path(org_file)
    output_path = Path(output_dir)
    
    if not org_path.exists():
        LOG.error(f"Organism file not found: {org_path}")
        return 1
    
    # Create output directory
    output_path.mkdir(parents=True, exist_ok=True)
    
    LOG.info(f"🍹 Starting protein download for organisms in {org_path}")
    LOG.info(f"📁 Output directory: {output_path}")
    LOG.info(f"⚡ Concurrent downloads: {concurrent}")
    LOG.info(f"📄 Format: {format}")
    
    # Parse organism file
    organisms = parse_organism_file(org_path)
    
    if not organisms:
        LOG.error("No organisms with NCBI FTP URLs found")
        
        if not skip_ncbi_check:
            LOG.error("Use --skip-ncbi-check to include non-NCBI organisms")
            return 1
    
    LOG.info(f"📋 Found {len(organisms)} organisms to process")
    
    # Download files
    results = []
    successful = 0
    failed = 0
    skipped = 0
    
    # Use ThreadPoolExecutor for concurrent downloads
    with ThreadPoolExecutor(max_workers=concurrent) as executor:
        # Submit all download tasks
        future_to_org = {
            executor.submit(
                download_protein_file,
                organism,
                output_path,
                format,
                config.get("MAX_RETRIES"),
                skip_existing
            ): organism['org_code']
            for organism in organisms
        }
        
        # Process results as they complete
        for future in as_completed(future_to_org):
            org_code = future_to_org[future]
            
            try:
                result = future.result(timeout=120)  # 120 second timeout for large files
                results.append(result)
                
                if result['success']:
                    if result.get('message') == "File already exists":
                        skipped += 1
                    else:
                        successful += 1
                else:
                    failed += 1
                    LOG.warning(f"Failed: {org_code} - {result.get('error', 'Unknown error')}")
            
            except Exception as e:
                LOG.error(f"Unexpected error processing {org_code}: {e}")
                failed += 1
                results.append({
                    'org_code': org_code,
                    'success': False,
                    'error': str(e),
                    'attempts': 0
                })
    
    # Print summary
    print("\n" + "="*60)
    print("🍹 PROTEIN DOWNLOAD SUMMARY")
    print("="*60)
    print(f"📊 Total organisms: {len(organisms)}")
    print(f"✅ Successful downloads: {successful}")
    print(f"⏭️  Skipped (already existed): {skipped}")
    print(f"❌ Failed downloads: {failed}")
    print(f"📁 Output directory: {output_path}")
    print("="*60)
    
    # Save results
    if results and polars_utils.POLARS_AVAILABLE:
        import polars as pl
        
        results_df = pl.DataFrame(results)
        results_csv = output_path / "protein_download_results.csv"
        results_df.write_csv(results_csv)
        LOG.info(f"📝 Results saved to: {results_csv}")
    
    # Return exit code
    if failed == 0:
        LOG.info("🎉 All protein downloads completed successfully!")
        return 0
    elif successful > 0:
        LOG.info(f"⚠️  {failed} downloads failed, but {successful} succeeded.")
        return 1
    else:
        LOG.error("❌ All protein downloads failed!")
        return 2


def main():
    """Command-line entry point."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Download protein sequences from NCBI"
    )
    parser.add_argument(
        "--org", "-i",
        required=True,
        help="Organism file from download-org command"
    )
    parser.add_argument(
        "--out", "-o",
        default="proteins",
        help="Output directory (default: proteins)"
    )
    parser.add_argument(
        "--concurrent", "-c",
        type=int,
        default=config.get("DEFAULT_CONCURRENT"),
        help=f"Number of concurrent downloads (default: {config.get('DEFAULT_CONCURRENT')})"
    )
    parser.add_argument(
        "--format",
        choices=["fasta", "fasta.gz"],
        default="fasta.gz",
        help="Download format (default: fasta.gz)"
    )
    parser.add_argument(
        "--skip-ncbi-check",
        action="store_true",
        help="Skip NCBI URL validation"
    )
    parser.add_argument(
        "--skip-existing",
        action="store_true",
        default=True,
        help="Skip files that already exist (default: True)"
    )
    parser.add_argument(
        "--no-skip-existing",
        action="store_false",
        dest="skip_existing",
        help="Download 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")
    )
    
    # Run download
    return download_protein_files(
        org_file=args.org,
        output_dir=args.out,
        concurrent=args.concurrent,
        format=args.format,
        skip_ncbi_check=args.skip_ncbi_check,
        skip_existing=args.skip_existing
    )


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