#!/usr/bin/env python3
"""
Download KEGG Orthology (KO) files
"""
import sys
import logging
import time
from pathlib import Path
from typing import List, Dict, Optional
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 read_organism_codes(org_file: Path) -> List[str]:
    """
    Read organism codes from file.
    
    Supports:
    - One code per line
    - TSV file with org_code in first column
    """
    org_codes = []
    
    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
            
            # Skip header line if it contains column names
            if line_num == 1 and line.lower().startswith(('org_code', 'org_abbr', 'name', 'url')):
                continue
            
            # Try to parse as TSV
            if '\t' in line:
                parts = line.split('\t')
                if parts:  # First column is org_code
                    org_codes.append(parts[0])
            else:
                # Assume one code per line
                org_codes.append(line)
    
    LOG.info(f"Read {len(org_codes)} organism codes from {org_file}")
    return org_codes

def download_single_ko(
    org_code: str,
    output_dir: Path,
    max_retries: int = 3,
    skip_existing: bool = True
) -> Dict:
    """
    Download KO file for a single organism.
    
    Returns:
        Dictionary with download result
    """
    result = {
        'org_code': org_code,
        'success': False,
        'error': None,
        'filepath': None,
        'size_bytes': 0,
        'attempts': 0
    }
    
    # Construct filename
    filename = f"{org_code}00001.keg"
    output_path = output_dir / 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
    
    # Construct KEGG URL
    url = f"https://www.kegg.jp/kegg-bin/download_htext?htext={filename}&format=htext&filedir="
    
    # Try download with retries
    last_error = None
    
    for attempt in range(max_retries):
        result['attempts'] = attempt + 1
        
        try:
            # Add delay between requests to be nice to KEGG server
            if attempt > 0:
                delay = 2 ** attempt  # Exponential backoff
                LOG.debug(f"Retry {attempt + 1}/{max_retries} for {org_code}, waiting {delay}s...")
                time.sleep(delay)
            
            LOG.debug(f"Downloading {org_code} (attempt {attempt + 1}/{max_retries})")
            
            # Make request
            req = urllib.request.Request(
                url,
                headers={
                    'User-Agent': 'gpathex /1.0.0 (+https://github.com/yourusername/gpathex )'
                }
            )
            
            with urllib.request.urlopen(req, timeout=30) as response:
                content = response.read()
                
                # Check if file is empty or has error
                if not content:
                    raise ValueError("Empty response from server")
                
                # Check for KEGG error messages
                content_str = content.decode('utf-8', errors='ignore')
                if "Error" in content_str and "not found" in content_str:
                    raise ValueError(f"KO file not found for {org_code}")
                
                # 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 {filename} ({len(content):,} bytes)")
                return result
        
        except urllib.error.HTTPError as e:
            if e.code == 404:
                result['error'] = f"KO file not found (404)"
                result['message'] = "KO file not available for this organism"
                LOG.warning(f"⚠️  KO 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 30 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}: {result['error']}")
    return result


def download_ko_files(
    org_file: str,
    output_dir: str = "ko_files",
    concurrent: int = 5,
    max_retries: int = 3,
    skip_existing: bool = True
) -> int:
    """
    Download KO files for multiple organisms.
    
    Args:
        org_file: File with organism codes
        output_dir: Output directory
        concurrent: Number of concurrent downloads
        max_retries: Maximum retry attempts
        skip_existing: Skip files that already exist
    
    Returns:
        Exit code (0 = success, 1 = partial failure, 2 = complete failure)
    """
    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 KO download for organisms in {org_path}")
    LOG.info(f"📁 Output directory: {output_path}")
    LOG.info(f"⚡ Concurrent downloads: {concurrent}")
    LOG.info(f"🔄 Max retries: {max_retries}")
    LOG.info(f"⏭️  Skip existing: {skip_existing}")
    
    # Read organism codes
    org_codes = read_organism_codes(org_path)
    
    if not org_codes:
        LOG.error("No organism codes found in input file")
        return 1
    
    LOG.info(f"📋 Found {len(org_codes)} 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_single_ko,
                org_code,
                output_path,
                max_retries,
                skip_existing
            ): org_code
            for org_code in org_codes
        }
        
        # Process results as they complete
        for future in as_completed(future_to_org):
            org_code = future_to_org[future]
            
            try:
                result = future.result(timeout=60)  # 60 second timeout
                results.append(result)
                
                if result['success']:
                    if result.get('message') == "File already exists":
                        skipped += 1
                    else:
                        successful += 1
                else:
                    failed += 1
                    
                    # Log failure details
                    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("🍹 KO DOWNLOAD SUMMARY")
    print("="*60)
    print(f"📊 Total organisms: {len(org_codes)}")
    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 detailed results
    if polars_utils.POLARS_AVAILABLE:
        import polars as pl
        
        # Create results DataFrame
        results_df = pl.DataFrame(results)
        
        # Save to CSV
        results_csv = output_path / "download_results.csv"
        results_df.write_csv(results_csv)
        LOG.info(f"📝 Detailed results saved to: {results_csv}")
        
        # Save summary
        summary_path = output_path / "download_summary.txt"
        with open(summary_path, 'w') as f:
            f.write("gpathex  KO Download Summary\n")
            f.write("=" * 50 + "\n")
            f.write(f"Date: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"Input file: {org_path}\n")
            f.write(f"Output directory: {output_path}\n")
            f.write(f"Total organisms: {len(org_codes)}\n")
            f.write(f"Successful downloads: {successful}\n")
            f.write(f"Skipped (existing): {skipped}\n")
            f.write(f"Failed: {failed}\n\n")
            
            if failed > 0:
                f.write("Failed organisms:\n")
                for result in results:
                    if not result['success']:
                        f.write(f"  {result['org_code']}: {result.get('error', 'Unknown error')}\n")
    
    # Save failed organisms list
    if failed > 0:
        failed_orgs = [r['org_code'] for r in results if not r['success']]
        failed_file = output_path / "failed_organisms.txt"
        with open(failed_file, 'w') as f:
            for org in failed_orgs:
                f.write(f"{org}\n")
        
        LOG.warning(f"⚠️  {failed} organisms failed. List saved to: {failed_file}")
    
    # Check for completely empty directory (all failed)
    if successful == 0 and skipped == 0:
        LOG.error("❌ All downloads failed! Check network connection and organism codes.")
        return 2
    
    # Return exit code based on success rate
    if failed == 0:
        LOG.info("🎉 All downloads completed successfully!")
        return 0
    elif successful > 0:
        LOG.info(f"⚠️  {failed} downloads failed, but {successful} succeeded.")
        return 1
    else:
        LOG.error("❌ All downloads failed!")
        return 2


def main():
    """Command-line entry point."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Download KEGG Orthology (KO) files"
    )
    parser.add_argument(
        "--org", "-i",
        required=True,
        help="Input file with organism codes (one per line or TSV)"
    )
    parser.add_argument(
        "--out", "-o",
        default="ko_files",
        help="Output directory (default: ko_files)"
    )
    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(
        "--retries", "-r",
        type=int,
        default=config.get("MAX_RETRIES"),
        help=f"Maximum retry attempts (default: {config.get('MAX_RETRIES')})"
    )
    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_ko_files(
        org_file=args.org,
        output_dir=args.out,
        concurrent=args.concurrent,
        max_retries=args.retries,
        skip_existing=args.skip_existing
    )


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