#!/usr/bin/env python3
"""
Create .keg file from annotations
"""
import sys
import logging
from pathlib import Path
from typing import Dict, List, Set, DefaultDict
from collections import defaultdict

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

LOG = logging.getLogger(__name__)


def load_annotations(input_file: Path) -> Dict[str, Dict[str, List[str]]]:
    """
    Load annotation data.
    
    Expected format (TSV):
    protein_id\tko_id\tpathway_ids
    
    Returns: {pathway_id: {ko_id: [protein_ids]}}
    """
    if not input_file.exists():
        LOG.error(f"Input file not found: {input_file}")
        return {}
    
    # Structure: pathway -> ko -> proteins
    pathway_dict = defaultdict(lambda: defaultdict(list))
    
    with open(input_file, 'r') 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) >= 2:
                protein_id = parts[0].strip()
                ko_id = parts[1].strip()
                
                # Get pathway(s)
                pathway_ids = []
                if len(parts) >= 3:
                    pathway_str = parts[2].strip()
                    if pathway_str:
                        pathway_ids = [p.strip() for p in pathway_str.split(';') if p.strip()]
                
                # Add to dictionary
                for pathway_id in pathway_ids:
                    if pathway_id and ko_id and ko_id != '-':
                        pathway_dict[pathway_id][ko_id].append(protein_id)
    
    LOG.info(f"Loaded annotations for {len(pathway_dict)} pathways")
    
    # Log some stats
    total_kos = sum(len(ko_dict) for ko_dict in pathway_dict.values())
    total_proteins = sum(
        len(protein_list)
        for ko_dict in pathway_dict.values()
        for protein_list in ko_dict.values()
    )
    
    LOG.info(f"Total KOs: {total_kos}")
    LOG.info(f"Total protein-KO associations: {total_proteins}")
    
    return pathway_dict


def create_keg_file(
    keg_template: str,
    input_file: str,
    output_prefix: str = "output"
) -> int:
    """
    Create .keg file from annotation results.
    """
    template_path = Path(keg_template)
    input_path = Path(input_file)
    
    # Validate inputs
    if not template_path.exists():
        LOG.error(f"Template file not found: {template_path}")
        return 1
    
    if not input_path.exists():
        LOG.error(f"Input file not found: {input_path}")
        return 1
    
    LOG.info(f"🍹 Creating KEGG map from annotations")
    LOG.info(f"📁 Template: {template_path}")
    LOG.info(f"📁 Annotations: {input_path}")
    LOG.info(f"📁 Output prefix: {output_prefix}")
    
    # Load annotations
    pathway_dict = load_annotations(input_path)
    
    if not pathway_dict:
        LOG.error("No annotations loaded")
        return 1
    
    # Read template and create output
    output_path = Path(f"{output_prefix}.keg")
    
    try:
        with open(template_path, 'r', encoding='utf-8') as template, \
             open(output_path, 'w', encoding='utf-8') as output:
            
            current_pathway = None
            skip_section = False
            proteins_written = 0
            pathways_written = 0
            
            for line_num, line in enumerate(template, 1):
                line = line.rstrip('\n')
                
                if not line:
                    output.write(line + '\n')
                    continue
                
                tag = line[0] if line else ''
                
                # Pathway line (C)
                if tag == 'C':
                    # Check if this pathway is in our annotations
                    parts = line.split()
                    if len(parts) >= 2 and parts[1].startswith('PATH:'):
                        pathway_id = parts[1].replace('PATH:', '')
                        current_pathway = pathway_id
                        
                        if pathway_id in pathway_dict:
                            skip_section = False
                            pathways_written += 1
                            LOG.debug(f"Including pathway: {pathway_id}")
                        else:
                            skip_section = True
                            LOG.debug(f"Skipping pathway: {pathway_id}")
                    
                    output.write(line + '\n')
                    continue
                
                # KO/Protein line (D)
                elif tag == 'D' and not skip_section and current_pathway:
                    # Parse the KO from template line
                    parts = line.split('\t')
                    if len(parts) >= 2:
                        ko_part = parts[1].split()
                        if ko_part:
                            template_ko = ko_part[0]
                            
                            # Check if we have proteins for this KO in current pathway
                            if (current_pathway in pathway_dict and 
                                template_ko in pathway_dict[current_pathway]):
                                
                                proteins = pathway_dict[current_pathway][template_ko]
                                
                                # Write each protein
                                for protein_id in proteins:
                                    # Reconstruct the line with our protein
                                    if len(parts[0].split()) >= 2:
                                        # Format: "D      GENE_ID\tKO DESCRIPTION"
                                        gene_part = parts[0].split()
                                        new_line = f"D      {protein_id}"
                                        
                                        if len(parts) >= 2:
                                            # Keep the KO description
                                            new_line += f"\t{parts[1]}"
                                        
                                        output.write(new_line + '\n')
                                        proteins_written += 1
                            else:
                                # KO not in our annotations, skip this line
                                continue
                        else:
                            # No KO in template line, write as-is
                            output.write(line + '\n')
                    else:
                        output.write(line + '\n')
                
                # Other lines
                else:
                    if not skip_section:
                        output.write(line + '\n')
        
        LOG.info(f"✅ Created KEGG map: {output_path}")
        LOG.info(f"📊 Pathways included: {pathways_written}")
        LOG.info(f"📊 Proteins written: {proteins_written}")
        
        # Create summary file
        summary_path = Path(f"{output_prefix}.summary.txt")
        with open(summary_path, 'w') as f:
            f.write("KEGG Map Creation Summary\n")
            f.write("=" * 50 + "\n")
            f.write(f"Template: {template_path}\n")
            f.write(f"Annotations: {input_path}\n")
            f.write(f"Output: {output_path}\n\n")
            f.write(f"Pathways in template: {pathways_written}\n")
            f.write(f"Proteins written: {proteins_written}\n\n")
            
            f.write("Pathways included:\n")
            f.write("-" * 40 + "\n")
            for pathway_id in sorted(pathway_dict.keys()):
                ko_count = len(pathway_dict[pathway_id])
                protein_count = sum(len(proteins) for proteins in pathway_dict[pathway_id].values())
                f.write(f"{pathway_id}: {ko_count} KOs, {protein_count} proteins\n")
        
        LOG.info(f"📝 Summary saved to: {summary_path}")
        
        return 0
        
    except Exception as e:
        LOG.error(f"❌ Error creating KEGG file: {e}")
        return 1


def main():
    """Command-line entry point."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Create .keg file from annotations"
    )
    parser.add_argument(
        "--keg", "-k",
        required=True,
        help="Template KEGG file (e.g., ko00001.keg)"
    )
    parser.add_argument(
        "--in", "-i",
        dest="input_file",
        required=True,
        help="Input annotations file"
    )
    parser.add_argument(
        "--out", "-o",
        default="output",
        help="Output prefix (default: output)"
    )
    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 KEGG file
    return create_keg_file(
        keg_template=args.keg,
        input_file=args.input_file,
        output_prefix=args.out
    )


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