#!/usr/bin/env python3
"""
Create custom KEGG database
"""
import sys
import logging
import gzip
import json
from pathlib import Path
from typing import List, Dict, Optional, Set
from collections import defaultdict

from .process_proteins import extract_gene_id

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

LOG = logging.getLogger(__name__)


def create_database(
    org_file: str,
    keg_dir: str,
    pep_dir: str,
    output_prefix: str = "kegg_db",
    format: str = "both",
    compress: bool = False
) -> int:
    """
    Create custom KEGG database from organism files.
    
    Args:
        org_file: File with organism codes
        keg_dir: Directory with KO annotation files
        pep_dir: Directory with protein FASTA files
        output_prefix: Prefix for output files (can be directory path)
        format: Output format (fasta, parquet, both)
        compress: Compress output files
    """
    org_path = Path(org_file)
    keg_path = Path(keg_dir)
    pep_path = Path(pep_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
    
    LOG.info(f"🍹 Creating KEGG database")
    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 prefix: {output_prefix}")
    LOG.info(f"📄 Format: {format}")
    LOG.info(f"🗜️  Compress: {compress}")
    
    # Parse output_prefix to handle both directory and basename
    output_path = Path(output_prefix)
    
    # If output_prefix ends with slash or is a directory-like path
    if output_prefix.endswith('/') or output_prefix.endswith('\\'):
        # User specified a directory, use default basename
        output_dir = output_path
        db_basename = "kegg_db"
    else:
        # Could be a file path or basename
        if output_path.suffix:  # Has extension like .fasta, .parquet
            # User specified a filename, extract directory and basename
            output_dir = output_path.parent
            db_basename = output_path.stem
        else:
            # User specified a basename or directory without extension
            if output_path.exists() and output_path.is_dir():
                # It's an existing directory
                output_dir = output_path
                db_basename = "kegg_db"
            else:
                # Assume it's a basename for files
                output_dir = Path.cwd()
                db_basename = output_path.name
    
    # Create output directory
    output_dir.mkdir(parents=True, exist_ok=True)
    
    LOG.info(f"📁 Output directory: {output_dir}")
    LOG.info(f"🏷️  Database basename: {db_basename}")
    
    # Read organism codes (skip header)
    org_codes = []
    with open(org_path, 'r', encoding='utf-8') 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')):
                LOG.debug(f"Skipping header line: {line}")
                continue
            
            parts = line.split('\t')
            if parts:
                org_code = parts[0].strip()
                if org_code:  # Skip empty codes
                    org_codes.append(org_code)
    
    if not org_codes:
        LOG.error("No organism codes found")
        return 1
    
    LOG.info(f"📋 Found {len(org_codes)} organisms to process")
    
    # Collections for data
    all_sequences = []  # For FASTA output
    sequence_data = []  # For DataFrame
    annotation_data = []  # KO annotations
    
    processed_orgs = set()
    failed_orgs = []
    
    # Process each organism
    for i, org_code in enumerate(org_codes, 1):
        LOG.info(f"[{i}/{len(org_codes)}] Processing {org_code}")
        
        try:
            # Find files
            keg_file = keg_path / f"{org_code}00001.keg"
            pep_file_gz = pep_path / f"{org_code}.pep.fasta.gz"
            pep_file_fasta = pep_path / f"{org_code}.pep.fasta"
            
            # Check if files exist
            if not keg_file.exists():
                LOG.warning(f"⚠️  KO file not found: {keg_file}")
                failed_orgs.append((org_code, "KO file not found"))
                continue
            
            # 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:
                LOG.warning(f"⚠️  Protein file not found: {org_code}.pep.fasta[.gz]")
                failed_orgs.append((org_code, "Protein file not found"))
                continue
            
            # Parse KO annotations
            LOG.info(f"Parsing KO file: {keg_file}")
            gene_ko_dict = parsers.parse_ko_file_to_gene_dict(str(keg_file))
            
            if not gene_ko_dict:
                LOG.warning(f"⚠️  No KO annotations found for {org_code}")
                failed_orgs.append((org_code, "No KO annotations"))
                continue
            
            LOG.info(f"Found {len(gene_ko_dict)} genes with KO annotations")
            
            # Process protein sequences
            seq_count = 0
            ann_count = 0
            total_proteins = 0
            matched_proteins = 0
            
            LOG.info(f"Processing protein file: {pep_file}")
            
            for record in fasta.read_fasta(pep_file):
                total_proteins += 1
                
                # Extract gene ID
                gene_id = None
                header = record.header
                
                # Try multiple extraction methods
                gene_id = extract_gene_id(header)
                
                if not gene_id:
                    # Fallback: try to extract from header
                    if ' ' in header:
                        gene_id = header.split()[0]
                    else:
                        gene_id = header
                
                # Clean up gene_id
                gene_id = gene_id.lstrip('>').strip()
                
                if gene_id and gene_id in gene_ko_dict:
                    ko_list = gene_ko_dict[gene_id]["ko"]
                    pathway_list = gene_ko_dict[gene_id]["pathway"]
                    
                    # Skip if no KOs
                    if not ko_list or (len(ko_list) == 1 and ko_list[0] == '-'):
                        continue
                    
                    matched_proteins += 1
                    
                    # Create database entry
                    db_id = f"{org_code}:{gene_id}"
                    
                    # Add to FASTA collection
                    fasta_header = f"{db_id}"
                    if record.description:
                        fasta_header += f" {record.description}"
                    
                    all_sequences.append(
                        f">{fasta_header}\n{record.sequence}"
                    )
                    
                    # Add to sequence data
                    sequence_data.append({
                        'db_id': db_id,
                        'org_code': org_code,
                        'gene_id': gene_id,
                        'sequence': record.sequence,
                        'length': len(record.sequence),
                        'description': record.description,
                        'ko_ids': ','.join(ko_list),
                        'pathway_ids': ','.join(pathway_list)
                    })
                    
                    # Add to annotation data
                    for ko_id in ko_list:
                        if ko_id != '-':
                            annotation_data.append({
                                'db_id': db_id,
                                'org_code': org_code,
                                'gene_id': gene_id,
                                'ko_id': ko_id,
                                'pathway_ids': ','.join(pathway_list)
                            })
                            ann_count += 1
                    
                    seq_count += 1
            
            LOG.info(f"Processed {total_proteins} proteins, matched {matched_proteins} with KO annotations")
            
            if seq_count > 0:
                processed_orgs.add(org_code)
                LOG.info(f"✅ Added {seq_count} sequences, {ann_count} annotations from {org_code}")
            else:
                LOG.warning(f"⚠️  No sequences with KO found for {org_code}")
                LOG.warning(f"   Total proteins: {total_proteins}, Genes in KO file: {len(gene_ko_dict)}")
                failed_orgs.append((org_code, f"No sequences with KO (proteins={total_proteins}, KO_genes={len(gene_ko_dict)})"))
        
        except Exception as e:
            LOG.error(f"❌ Error processing {org_code}: {e}")
            failed_orgs.append((org_code, str(e)))
    
    # Check if we have any data
    if not sequence_data:
        LOG.error("❌ No data collected! Database creation failed.")
        
        # Provide more debugging info
        LOG.error("Possible reasons:")
        LOG.error("1. Gene IDs in KO files don't match protein FASTA headers")
        LOG.error("2. KO files have no actual KO annotations (all ko_id='-')")
        LOG.error("3. Protein files are empty or corrupted")
        
        # List what we found
        LOG.error(f"Organisms processed: {len(processed_orgs)}")
        LOG.error(f"Organisms failed: {len(failed_orgs)}")
        for org, reason in failed_orgs:
            LOG.error(f"  {org}: {reason}")
        
        return 1
    
    LOG.info(f"📊 Collected {len(sequence_data)} sequences with {len(annotation_data)} annotations from {len(processed_orgs)} organisms")
    
    # Save data based on format
    base_path = output_dir / db_basename  # Fixed: Use db_basename, not output_prefix
    
    # 1. FASTA format (always created for compatibility)
    if format in ["fasta", "both"]:
        fasta_file = f"{base_path}.pep.fasta"
        if compress:
            fasta_file += ".gz"
        
        LOG.info(f"💾 Writing FASTA file: {fasta_file}")
        
        if compress:
            with gzip.open(fasta_file, 'wt', encoding='utf-8') as f:
                f.write('\n'.join(all_sequences))
        else:
            with open(fasta_file, 'w', encoding='utf-8') as f:
                f.write('\n'.join(all_sequences))
    
    # 2. Parquet format (if Polars available)
    if format in ["parquet", "both"] and polars_utils.POLARS_AVAILABLE:
        import polars as pl
        
        LOG.info("💾 Creating Polars DataFrames...")
        
        # Create sequences DataFrame
        seq_df = pl.DataFrame(sequence_data)
        seq_parquet = f"{base_path}.sequences.parquet"
        seq_df.write_parquet(seq_parquet)
        LOG.info(f"📊 Sequences DataFrame: {seq_parquet} ({seq_df.height} rows)")
        
        # Create annotations DataFrame
        if annotation_data:
            ann_df = pl.DataFrame(annotation_data)
            ann_parquet = f"{base_path}.annotations.parquet"
            ann_df.write_parquet(ann_parquet)
            LOG.info(f"📊 Annotations DataFrame: {ann_parquet} ({ann_df.height} rows)")
        
        # Create KO mapping file (TSV for compatibility)
        ko_tsv = f"{base_path}.pep2ko.tsv"
        if annotation_data:
            # Group by db_id
            ko_mapping = defaultdict(list)
            for ann in annotation_data:
                ko_mapping[ann['db_id']].append(ann['ko_id'])
            
            with open(ko_tsv, 'w') as f:
                f.write("Protein_ID\tKO_IDs\tPathway_IDs\n")
                for seq in sequence_data:
                    db_id = seq['db_id']
                    ko_ids = ','.join(ko_mapping.get(db_id, []))
                    pathway_ids = seq['pathway_ids']
                    f.write(f"{db_id}\t{ko_ids}\t{pathway_ids}\n")
            
            LOG.info(f"🔗 KO mapping file: {ko_tsv}")
    
    # 3. JSON statistics
    stats_file = f"{base_path}.stats.json"
    stats = {
        "database_name": db_basename,
        "output_directory": str(output_dir),
        "creation_date": str(Path(fasta_file).stat().st_ctime if format in ["fasta", "both"] else ""),
        "organisms_processed": len(processed_orgs),
        "organisms_failed": len(failed_orgs),
        "total_sequences": len(sequence_data),
        "total_annotations": len(annotation_data),
        "organisms": list(processed_orgs),
        "failed_organisms": [
            {"org_code": org, "reason": reason}
            for org, reason in failed_orgs
        ]
    }
    
    with open(stats_file, 'w') as f:
        json.dump(stats, f, indent=2)
    
    LOG.info(f"📈 Statistics file: {stats_file}")
    
    # 4. Summary text file
    summary_file = f"{base_path}.summary.txt"
    with open(summary_file, 'w') as f:
        f.write("=" * 60 + "\n")
        f.write(f"KEGG DATABASE: {db_basename}\n")
        f.write("=" * 60 + "\n\n")
        f.write(f"Output directory: {output_dir}\n")
        f.write(f"Creation date: {Path(summary_file).stat().st_ctime}\n")
        f.write(f"Organisms: {len(processed_orgs)} processed, {len(failed_orgs)} failed\n")
        f.write(f"Sequences: {len(sequence_data):,}\n")
        f.write(f"Annotations: {len(annotation_data):,}\n\n")
        
        f.write("PROCESSED ORGANISMS:\n")
        f.write("-" * 40 + "\n")
        for org in sorted(processed_orgs):
            org_seqs = sum(1 for s in sequence_data if s['org_code'] == org)
            f.write(f"{org}: {org_seqs:,} sequences\n")
        
        if failed_orgs:
            f.write("\nFAILED ORGANISMS:\n")
            f.write("-" * 40 + "\n")
            for org, reason in failed_orgs:
                f.write(f"{org}: {reason}\n")
        
        f.write("\nOUTPUT FILES:\n")
        f.write("-" * 40 + "\n")
        if format in ["fasta", "both"]:
            f.write(f"FASTA: {fasta_file}\n")
        if format in ["parquet", "both"] and polars_utils.POLARS_AVAILABLE:
            f.write(f"Sequences (Parquet): {base_path}.sequences.parquet\n")
            f.write(f"Annotations (Parquet): {base_path}.annotations.parquet\n")
            f.write(f"KO mapping (TSV): {base_path}.pep2ko.tsv\n")
        f.write(f"Statistics: {stats_file}\n")
        f.write(f"Summary: {summary_file}\n")
    
    LOG.info(f"📝 Summary file: {summary_file}")
    
    # Print final summary
    print("\n" + "="*60)
    print("🍹 DATABASE CREATION COMPLETE")
    print("="*60)
    print(f"🏷️  Database name: {db_basename}")
    print(f"📁 Output directory: {output_dir}")
    print(f"🧬 Organisms: {len(processed_orgs)}")
    print(f"🧬 Sequences: {len(sequence_data):,}")
    print(f"🔗 Annotations: {len(annotation_data):,}")
    print("="*60)
    
    if failed_orgs:
        LOG.warning(f"⚠️  {len(failed_orgs)} organisms failed (see {summary_file})")
        return 1
    
    LOG.info("🎉 Database created successfully!")
    return 0

def main():
    """Command-line entry point."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Create custom KEGG database"
    )
    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="kegg_db",
        help="Output prefix (default: kegg_db)"
    )
    parser.add_argument(
        "--format",
        choices=["fasta", "parquet", "both"],
        default="both",
        help="Output format (default: both)"
    )
    parser.add_argument(
        "--compress",
        action="store_true",
        help="Compress output files"
    )
    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")
    )
    
    # Create database
    return create_database(
        org_file=args.org,
        keg_dir=args.keg,
        pep_dir=args.pep,
        output_prefix=args.out,
        format=args.format,
        compress=args.compress
    )


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