#!/usr/bin/env python3
"""
Process proteins with KO annotations
"""
import sys
import logging
import re
from pathlib import Path
from typing import List, Dict, Optional, Set
from collections import defaultdict

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

LOG = logging.getLogger(__name__)


def extract_gene_id(header: str) -> Optional[str]:
    """
    Extract gene ID from FASTA header.
    
    Supports multiple formats:
    - NCBI: locus_tag=XXXX, protein_id=XXXX, db_xref=GeneID:XXXX
    - Generic: Try to find common patterns
    """
    # Try NCBI format first
    if "locus_tag=" in header:
        match = re.search(r'locus_tag=([^\]]+)', header)
        if match:
            return match.group(1).strip()
    
    if "protein_id=" in header:
        match = re.search(r'protein_id=([^\]]+)', header)
        if match:
            protein_id = match.group(1).strip()
            # Remove version number if present
            return protein_id.split('.')[0]
    
    if "db_xref=GeneID:" in header:
        match = re.search(r'db_xref=GeneID:([^\]]+)', header)
        if match:
            return match.group(1).strip()
    
    # Try generic patterns
    patterns = [
        r'gene:([^\s]+)',  # gene:XXXX
        r'GeneID:([^\s]+)',  # GeneID:XXXX
        r'ID=([^\s;]+)',  # ID=XXXX
        r'gene=([^\s;]+)',  # gene=XXXX
        r'gene_id "([^"]+)"',  # gene_id "XXXX"
    ]
    
    for pattern in patterns:
        match = re.search(pattern, header)
        if match:
            return match.group(1).strip()
    
    # Last resort: take first word after >
    if header:
        return header.split()[0] if ' ' in header else header
    
    return None


def parse_ko_file(keg_file: Path) -> Dict[str, Dict[str, List[str]]]:
    """
    Parse KEGG KO file to extract gene-to-KO mappings.
    
    Returns: {gene_id: {"ko": [ko_ids], "pathway": [pathway_ids]}}
    """
    if not keg_file.exists():
        LOG.error(f"KO file not found: {keg_file}")
        return {}
    
    gene_dict = {}
    current_pathway = None
    
    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
            
            # Pathway line (C)
            if line.startswith('C'):
                # Format: C   PATH: koXXXXX
                parts = line.split()
                if len(parts) >= 2 and parts[1].startswith('PATH:'):
                    current_pathway = parts[1].replace('PATH:', '')
                continue
            
            # Gene line (D)
            if line.startswith('D'):
                parts = line.split('\t')
                if len(parts) >= 2:
                    # Format: D      GENE_ID\tKO_NAME DESCRIPTION
                    gene_part = parts[0].split()
                    if len(gene_part) >= 2:
                        gene_id = gene_part[1]
                        
                        ko_info = parts[1].split() if len(parts) > 1 else []
                        if ko_info:
                            ko_id = ko_info[0]
                            # Skip if KO is just a dash
                            if ko_id != '-':
                                if gene_id not in gene_dict:
                                    gene_dict[gene_id] = {"ko": [], "pathway": []}
                                
                                if ko_id not in gene_dict[gene_id]["ko"]:
                                    gene_dict[gene_id]["ko"].append(ko_id)
                                
                                if current_pathway and current_pathway not in gene_dict[gene_id]["pathway"]:
                                    gene_dict[gene_id]["pathway"].append(current_pathway)
    
    LOG.info(f"Parsed {len(gene_dict)} genes with KO annotations from {keg_file}")
    return gene_dict


def process_single_organism(
    org_code: str,
    keg_dir: Path,
    pep_dir: Path,
    output_dir: Path,
    min_ko: int = 1
) -> Dict:
    """
    Process proteins for a single organism.
    
    Returns processing statistics.
    """
    stats = {
        'org_code': org_code,
        'total_sequences': 0,
        'sequences_with_ko': 0,
        'genes_with_ko': 0,
        'output_file': None,
        'error': None
    }
    
    try:
        # Find files
        keg_file = keg_dir / f"{org_code}00001.keg"
        pep_file_gz = pep_dir / f"{org_code}.pep.fasta.gz"
        pep_file_fasta = pep_dir / f"{org_code}.pep.fasta"
        
        # Check if files exist
        if not keg_file.exists():
            stats['error'] = f"KO file not found: {keg_file}"
            return stats
        
        # Try gzipped file first, then regular fasta
        if pep_file_gz.exists():
            pep_file = pep_file_gz
        elif pep_file_fasta.exists():
            pep_file = pep_file_fasta
        else:
            stats['error'] = f"Protein file not found: {org_code}.pep.fasta[.gz]"
            return stats
        
        # Parse KO annotations
        gene_ko_dict = parse_ko_file(keg_file)
        
        if not gene_ko_dict:
            stats['error'] = f"No KO annotations found in {keg_file}"
            return stats
        
        # Process protein sequences
        output_file = output_dir / f"{org_code}.processed.fasta"
        stats['output_file'] = str(output_file)
        
        sequences_with_ko = []
        genes_processed = set()
        
        for record in fasta.read_fasta(pep_file):
            stats['total_sequences'] += 1
            
            gene_id = extract_gene_id(record.header)
            if not gene_id:
                continue
            
            if gene_id in gene_ko_dict:
                ko_list = gene_ko_dict[gene_id]["ko"]
                
                # Filter by minimum KOs
                if len(ko_list) >= min_ko:
                    # Add KO information to header
                    new_header = f"{org_code}:{gene_id} KO={','.join(ko_list)}"
                    if record.description:
                        new_header += f" {record.description}"
                    
                    sequences_with_ko.append(fasta.FastaRecord(
                        seq_id=f"{org_code}:{gene_id}",
                        sequence=record.sequence,
                        description=f"KO={','.join(ko_list)} {record.description}",
                        metadata={
                            'original_header': record.header,
                            'ko_ids': ko_list,
                            'pathways': gene_ko_dict[gene_id]["pathway"],
                            'org_code': org_code
                        }
                    ))
                    
                    genes_processed.add(gene_id)
        
        stats['sequences_with_ko'] = len(sequences_with_ko)
        stats['genes_with_ko'] = len(genes_processed)
        
        # Write output if we found sequences
        if sequences_with_ko:
            fasta.write_fasta(sequences_with_ko, output_file)
            LOG.info(f"✅ Processed {org_code}: {stats['sequences_with_ko']} sequences with KO (out of {stats['total_sequences']} total)")
        else:
            LOG.warning(f"⚠️  No sequences with KO found for {org_code}")
        
        return stats
        
    except Exception as e:
        stats['error'] = str(e)
        LOG.error(f"❌ Error processing {org_code}: {e}")
        return stats


def process_proteins(
    org_file: str,
    keg_dir: str,
    pep_dir: str,
    output_dir: str,
    min_ko: int = 1
) -> int:
    """
    Process proteins with KO annotations for multiple organisms.
    """
    org_path = Path(org_file)
    keg_path = Path(keg_dir)
    pep_path = Path(pep_dir)
    output_path = Path(output_dir)
    
    # Validate inputs
    if not org_path.exists():
        LOG.error(f"Organism file not found: {org_path}")
        return 1
    
    if not keg_path.exists():
        LOG.error(f"KO directory not found: {keg_path}")
        return 1
    
    if not pep_path.exists():
        LOG.error(f"Protein directory not found: {pep_path}")
        return 1
    
    # Create output directory
    output_path.mkdir(parents=True, exist_ok=True)
    
    LOG.info(f"🍹 Processing proteins with KO annotations")
    LOG.info(f"📁 Organism file: {org_path}")
    LOG.info(f"📁 KO directory: {keg_path}")
    LOG.info(f"📁 Protein directory: {pep_path}")
    LOG.info(f"📁 Output directory: {output_path}")
    LOG.info(f"🎯 Minimum KOs per gene: {min_ko}")
    
    # Read organism codes (skip header)
    org_codes = []
    with open(org_path, 'r') as f:
        for i, line in enumerate(f):
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            
            # Skip header line if it contains column names
            if i == 0 and line.lower().startswith(('org_code', 'org_abbr', 'name', 'url')):
                continue
            
            parts = line.split('\t')
            if parts:
                org_codes.append(parts[0])
    
    if not org_codes:
        LOG.error("No organism codes found")
        return 1
    
    LOG.info(f"📋 Found {len(org_codes)} organisms to process")
        
    # Process each organism
    all_stats = []
    successful = 0
    failed = 0
    skipped = 0
    
    for i, org_code in enumerate(org_codes, 1):
        LOG.info(f"[{i}/{len(org_codes)}] Processing {org_code}")
        
        stats = process_single_organism(
            org_code=org_code,
            keg_dir=keg_path,
            pep_dir=pep_path,
            output_dir=output_path,
            min_ko=min_ko
        )
        
        all_stats.append(stats)
        
        if stats['error']:
            failed += 1
            LOG.error(f"❌ Failed {org_code}: {stats['error']}")
        elif stats['sequences_with_ko'] == 0:
            skipped += 1
            LOG.warning(f"⚠️  Skipped {org_code}: No sequences with KO")
        else:
            successful += 1
    
    # Print summary
    print("\n" + "="*60)
    print("🍹 PROTEIN PROCESSING SUMMARY")
    print("="*60)
    print(f"📊 Total organisms: {len(org_codes)}")
    print(f"✅ Successfully processed: {successful}")
    print(f"⚠️  Skipped (no KO): {skipped}")
    print(f"❌ Failed: {failed}")
    print(f"📁 Output directory: {output_path}")
    
    # Calculate totals
    total_sequences = sum(s['total_sequences'] for s in all_stats)
    total_with_ko = sum(s['sequences_with_ko'] for s in all_stats)
    
    if total_sequences > 0:
        ko_percentage = (total_with_ko / total_sequences) * 100
        print(f"🧬 Total sequences: {total_sequences:,}")
        print(f"🔗 Sequences with KO: {total_with_ko:,} ({ko_percentage:.1f}%)")
    
    print("="*60)
    
    # Save detailed statistics
    if polars_utils.POLARS_AVAILABLE:
        import polars as pl
        
        # Convert stats to DataFrame
        stats_data = []
        for stats in all_stats:
            stats_data.append({
                'org_code': stats['org_code'],
                'total_sequences': stats['total_sequences'],
                'sequences_with_ko': stats['sequences_with_ko'],
                'genes_with_ko': stats['genes_with_ko'],
                'output_file': stats['output_file'] or '',
                'error': stats['error'] or '',
                'success': stats['error'] is None and stats['sequences_with_ko'] > 0
            })
        
        stats_df = pl.DataFrame(stats_data)
        
        # Save to CSV
        stats_csv = output_path / "processing_stats.csv"
        stats_df.write_csv(stats_csv)
        LOG.info(f"📊 Statistics saved to: {stats_csv}")
        
        # Save summary file
        summary_file = output_path / "processing_summary.txt"
        with open(summary_file, 'w') as f:
            f.write("gpathex  Protein Processing Summary\n")
            f.write("=" * 50 + "\n")
            f.write(f"Organism file: {org_path}\n")
            f.write(f"KO directory: {keg_path}\n")
            f.write(f"Protein directory: {pep_path}\n")
            f.write(f"Output directory: {output_path}\n")
            f.write(f"Minimum KOs per gene: {min_ko}\n\n")
            f.write(f"Total organisms: {len(org_codes)}\n")
            f.write(f"Successfully processed: {successful}\n")
            f.write(f"Skipped (no KO): {skipped}\n")
            f.write(f"Failed: {failed}\n")
            f.write(f"Total sequences: {total_sequences:,}\n")
            f.write(f"Sequences with KO: {total_with_ko:,}\n")
            
            if total_sequences > 0:
                f.write(f"KO coverage: {(total_with_ko/total_sequences)*100:.1f}%\n")
        
        LOG.info(f"📝 Summary saved to: {summary_file}")
    
    # Return exit code
    if successful == 0:
        LOG.error("❌ No organisms successfully processed!")
        return 1
    elif failed > 0:
        LOG.warning(f"⚠️  {failed} organisms failed, but {successful} succeeded")
        return 1
    else:
        LOG.info("🎉 All organisms processed successfully!")
        return 0


def main():
    """Command-line entry point."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Process proteins with KO annotations"
    )
    parser.add_argument(
        "--org", "-i",
        required=True,
        help="Organism file with codes"
    )
    parser.add_argument(
        "--keg", "-k",
        required=True,
        help="Directory with KO files"
    )
    parser.add_argument(
        "--pep", "-p",
        required=True,
        help="Directory with protein files"
    )
    parser.add_argument(
        "--out", "-o",
        default="processed",
        help="Output directory (default: processed)"
    )
    parser.add_argument(
        "--min-ko",
        type=int,
        default=1,
        help="Minimum number of KOs required (default: 1)"
    )
    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 processing
    return process_proteins(
        org_file=args.org,
        keg_dir=args.keg,
        pep_dir=args.pep,
        output_dir=args.out,
        min_ko=args.min_ko
    )


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