import argparse
import logging
import os
import sys
from gensim.models import Word2Vec
from gensim.models.word2vec import LineSentence

# Configure logging
logging.basicConfig(
    format='%(asctime)s : %(levelname)s : %(message)s', 
    level=logging.INFO,
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('glove_training.log')
    ]
)

def create_default_corpus(output_path, num_sentences=1000):
    """Create minimal corpus if none exists"""
    with open(output_path, 'w', encoding='utf-8') as f:
        for _ in range(num_sentences):
            f.write("software engineering machine learning data science cloud computing devops ")
            f.write("python java javascript react aws docker kubernetes\n")
    logging.warning(f"Created minimal corpus at {output_path}")

def train_glove_model(corpus_path, output_path, vector_size=100, window=5, min_count=5, epochs=10):
    """Train a Word2Vec model (GloVe-like embeddings)"""
    
    # Handle missing corpus
    if not os.path.exists(corpus_path):
        logging.warning(f"No corpus found at {corpus_path}")
        create_default_corpus(corpus_path)
    
    try:
        # Read corpus
        sentences = LineSentence(corpus_path)
        
        # Train model
        model = Word2Vec(
            sentences=sentences,
            vector_size=vector_size,
            window=window,
            min_count=min_count,
            workers=os.cpu_count(),
            epochs=epochs
        )
        
        # Save model
        model.wv.save(output_path)
        logging.info(f"GloVe model saved to {output_path}")
        return True
    except Exception as e:
        logging.error(f"Training failed: {str(e)}")
        return False

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Train GloVe-like word embeddings')
    parser.add_argument('--corpus_path', type=str, 
                        default='data/corpus.txt',
                        help='Path to corpus text file')
    parser.add_argument('--output_path', type=str, 
                        default='trained_models/glove.model',
                        help='Output path for trained model')
    parser.add_argument('--vector_size', type=int, default=100,
                        help='Embedding dimension size')
    parser.add_argument('--window', type=int, default=5,
                        help='Context window size')
    parser.add_argument('--min_count', type=int, default=5,
                        help='Minimum word frequency')
    parser.add_argument('--epochs', type=int, default=10,
                        help='Number of training epochs')
    
    args = parser.parse_args()
    
    # Create output directory if not exists
    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
    
    # Run training
    success = train_glove_model(
        corpus_path=args.corpus_path,
        output_path=args.output_path,
        vector_size=args.vector_size,
        window=args.window,
        min_count=args.min_count,
        epochs=args.epochs
    )
    
    sys.exit(0 if success else 1)