import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from pathlib import Path
import json
import numpy as np
from typing import List, Dict, Union, Optional
from dataclasses import dataclass
import logging
import argparse
import sys
import yaml

logger = logging.getLogger(__name__)

@dataclass
class InferenceConfig:
    """Simple configuration for inference"""
    model_path: str  # Path to saved model
    device: str = "auto"  # "auto", "cuda", "cpu"
    batch_size: int = 32
    max_length: int = 512
    return_probabilities: bool = True
    return_top_k: int = 1  # Return top K predictions


class ModelInference:
    """Simple inference class for text classification"""
    
    def __init__(self, config: InferenceConfig):
        self.config = config
        self.model = None
        self.tokenizer = None
        self.label_info = {}
        self.device = self._setup_device()
        
        # Load model and tokenizer
        self.load_model()
        
    def _setup_device(self) -> torch.device:
        """Setup device for inference"""
        if self.config.device == "auto":
            if torch.cuda.is_available():
                device = torch.device("cuda")
                logger.info("Using CUDA device")
            else:
                device = torch.device("cpu")
                logger.info("Using CPU device")
        else:
            device = torch.device(self.config.device)
            logger.info(f"Using specified device: {device}")
        
        return device
    
    def load_model(self):
        """Load model, tokenizer, and label mappings"""
        model_path = Path(self.config.model_path)
        
        if not model_path.exists():
            raise FileNotFoundError(f"Model path not found: {model_path}")
        
        # Load tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        logger.info(f"Loaded tokenizer from {model_path}")
        
        # Load model
        self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
        self.model.to(self.device)
        self.model.eval()
        logger.info(f"Loaded model from {model_path}")
        
        # Load label information
        label_info_path = model_path / "label_info.json"
        if label_info_path.exists():
            with open(label_info_path, 'r') as f:
                self.label_info = json.load(f)
            logger.info(f"Loaded label mappings: {self.label_info.get('id_to_label', {})}")
        else:
            logger.warning("No label_info.json found. Using default numeric labels.")
            # Create default mappings
            num_labels = self.model.config.num_labels
            self.label_info = {
                "id_to_label": {str(i): f"LABEL_{i}" for i in range(num_labels)},
                "label_to_id": {f"LABEL_{i}": i for i in range(num_labels)},
                "num_labels": num_labels
            }
    
    def predict_single(self, text: str) -> Dict:
        """Predict single text sample"""
        return self.predict_batch([text])[0]
    
    def predict_batch(self, texts: List[str]) -> List[Dict]:
        """Predict batch of texts"""
        if not texts:
            return []
        
        # Tokenize inputs
        inputs = self.tokenizer(
            texts,
            truncation=True,
            padding=True,
            max_length=self.config.max_length,
            return_tensors="pt"
        )
        
        # Move to device
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # Get predictions
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
        
        # Convert to probabilities
        probabilities = torch.softmax(logits, dim=-1)
        
        # Process results
        results = []
        for i, text in enumerate(texts):
            text_probs = probabilities[i].cpu().numpy()
            
            # Get top K predictions
            top_indices = np.argsort(text_probs)[-self.config.return_top_k:][::-1]
            
            predictions = []
            for idx in top_indices:
                label = self.label_info["id_to_label"].get(str(idx), f"LABEL_{idx}")
                prob = float(text_probs[idx])
                
                pred_dict = {
                    "label": label,
                    "label_id": int(idx),
                    "score": prob
                }
                predictions.append(pred_dict)
            
            result = {
                "text": text,
                "predictions": predictions,
                "predicted_label": predictions[0]["label"],
                "confidence": predictions[0]["score"]
            }
            
            if self.config.return_probabilities:
                # Add all class probabilities
                all_probs = {}
                for label_id, label_name in self.label_info["id_to_label"].items():
                    all_probs[label_name] = float(text_probs[int(label_id)])
                result["all_probabilities"] = all_probs
            
            results.append(result)
        
        return results
    
    def predict_file(self, input_file: str, output_file: str = None) -> List[Dict]:
        """Predict on texts from file"""
        input_path = Path(input_file)
        
        # Read texts from file
        texts = []
        if input_path.suffix == '.txt':
            # Plain text file (one text per line)
            with open(input_path, 'r', encoding='utf-8') as f:
                texts = [line.strip() for line in f if line.strip()]
        
        elif input_path.suffix == '.jsonl':
            # JSONL file with "text" field
            with open(input_path, 'r', encoding='utf-8') as f:
                for line in f:
                    if line.strip():
                        data = json.loads(line)
                        text = data.get("text", data.get("input", ""))
                        if text:
                            texts.append(text)
        
        else:
            raise ValueError(f"Unsupported file format: {input_path.suffix}")
        
        logger.info(f"Loaded {len(texts)} texts from {input_file}")
        
        # Process in batches
        all_results = []
        for i in range(0, len(texts), self.config.batch_size):
            batch_texts = texts[i:i + self.config.batch_size]
            batch_results = self.predict_batch(batch_texts)
            all_results.extend(batch_results)
            
            if i % (self.config.batch_size * 10) == 0:
                logger.info(f"Processed {i + len(batch_texts)}/{len(texts)} texts")
        
        # Save results if output file specified
        if output_file:
            output_path = Path(output_file)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            
            with open(output_path, 'w', encoding='utf-8') as f:
                for result in all_results:
                    f.write(json.dumps(result) + '\n')
            
            logger.info(f"Results saved to {output_file}")
        
        return all_results


class BatchInference:
    """Optimized batch inference for large datasets"""
    
    def __init__(self, config: InferenceConfig):
        self.inference = ModelInference(config)
        self.config = config
    
    def predict_large_file(self, input_file: str, output_file: str, 
                          chunk_size: int = 1000) -> None:
        """Process large files in chunks to manage memory"""
        input_path = Path(input_file)
        output_path = Path(output_file)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        # Count total lines first
        total_lines = 0
        with open(input_path, 'r', encoding='utf-8') as f:
            for _ in f:
                total_lines += 1
        
        logger.info(f"Processing {total_lines} lines in chunks of {chunk_size}")
        
        processed = 0
        with open(input_path, 'r', encoding='utf-8') as infile, \
             open(output_path, 'w', encoding='utf-8') as outfile:
            
            chunk_texts = []
            
            for line in infile:
                if line.strip():
                    if input_path.suffix == '.jsonl':
                        data = json.loads(line)
                        text = data.get("text", data.get("input", ""))
                    else:
                        text = line.strip()
                    
                    chunk_texts.append(text)
                    
                    if len(chunk_texts) >= chunk_size:
                        # Process chunk
                        results = self.inference.predict_batch(chunk_texts)
                        
                        # Write results
                        for result in results:
                            outfile.write(json.dumps(result) + '\n')
                        
                        processed += len(chunk_texts)
                        logger.info(f"Processed {processed}/{total_lines} texts")
                        
                        chunk_texts = []
            
            # Process remaining texts
            if chunk_texts:
                results = self.inference.predict_batch(chunk_texts)
                for result in results:
                    outfile.write(json.dumps(result) + '\n')
                processed += len(chunk_texts)
        
        logger.info(f"Completed processing {processed} texts")


def create_inference_config(
    model_path: str,
    device: str = "auto",
    batch_size: int = 32,
    **kwargs
) -> InferenceConfig:
    """Helper function to create inference configuration"""
    return InferenceConfig(
        model_path=model_path,
        device=device,
        batch_size=batch_size,
        **kwargs
    )


def main():
    """Main function with YAML configuration support"""
    
    parser = argparse.ArgumentParser(description="Model Inference Pipeline")
    
    # YAML configuration
    parser.add_argument("--config", type=str, help="Path to YAML configuration file")
    
    # Model settings
    parser.add_argument("--model-path", type=str, help="Path to saved model directory")
    parser.add_argument("--device", choices=["auto", "cuda", "cpu"], help="Device to run inference on")
    parser.add_argument("--batch-size", type=int, help="Batch size for inference")
    parser.add_argument("--max-length", type=int, help="Maximum sequence length for tokenization")
    
    # Inference settings
    parser.add_argument("--return-probabilities", action="store_true", help="Return all class probabilities")
    parser.add_argument("--return-top-k", type=int, help="Return top K predictions")
    
    # Input/Output settings
    parser.add_argument("--input-text", type=str, help="Single text for prediction")
    parser.add_argument("--input-file", type=str, help="Input file path (txt or jsonl)")
    parser.add_argument("--output-file", type=str, help="Output file path for results")
    parser.add_argument("--chunk-size", type=int, help="Chunk size for large file processing")
    
    # Logging
    parser.add_argument("--log-level", choices=["DEBUG", "INFO", "WARNING", "ERROR"], default="INFO", help="Logging level")
    
    args = parser.parse_args()
    
    # Set up logging
    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # Load configuration
    config_dict = {}
    
    # Load YAML config if provided
    if args.config:
        try:
            with open(args.config, 'r', encoding='utf-8') as f:
                config_dict = yaml.safe_load(f)
            logger.info(f"Loaded YAML configuration from: {args.config}")
        except Exception as e:
            logger.error(f"Error loading YAML config: {e}")
            sys.exit(1)
    
    # Override YAML config with CLI arguments
    cli_overrides = {}
    if args.model_path:
        cli_overrides['model_path'] = args.model_path
    if args.device:
        cli_overrides['device'] = args.device
    if args.batch_size:
        cli_overrides['batch_size'] = args.batch_size
    if args.max_length:
        cli_overrides['max_length'] = args.max_length
    if args.return_probabilities:
        cli_overrides['return_probabilities'] = True
    if args.return_top_k:
        cli_overrides['return_top_k'] = args.return_top_k
    
    # Merge configurations
    for key, value in cli_overrides.items():
        if key in config_dict:
            logger.info(f"Overriding YAML config '{key}' with CLI value: {value}")
        config_dict[key] = value
    
    # Validate model path
    model_path = config_dict.get('model_path')
    if not model_path:
        parser.error("--model-path is required (either in YAML config or CLI)")
    
    model_path = Path(model_path)
    if not model_path.exists():
        print(f"❌ Model path not found: {model_path}")
        print("Please ensure the model directory exists and contains the trained model files")
        sys.exit(1)
    
    # Create configuration object
    config = InferenceConfig(
        model_path=str(model_path),
        device=config_dict.get('device', 'auto'),
        batch_size=config_dict.get('batch_size', 32),
        max_length=config_dict.get('max_length', 512),
        return_probabilities=config_dict.get('return_probabilities', True),
        return_top_k=config_dict.get('return_top_k', 1)
    )
    
    try:
        print(f"Loading model from: {config.model_path}")
        print(f"Device: {config.device}")
        print(f"Batch size: {config.batch_size}")
        if args.config:
            print(f"Using YAML configuration: {args.config}")
        print()
        
        # Initialize inference
        inference = ModelInference(config)
        
        # Handle different input types
        if args.input_text:
            # Single text prediction
            print(f"=== Single Text Prediction ===")
            print(f"Input text: {args.input_text}")
            print()
            
            result = inference.predict_single(args.input_text)
            print(f"Predicted label: {result['predicted_label']}")
            print(f"Confidence: {result['confidence']:.4f}")
            print(f"Top {config.return_top_k} predictions:")
            for pred in result['predictions']:
                print(f"  - {pred['label']}: {pred['score']:.4f}")
            
            if config.return_probabilities and 'all_probabilities' in result:
                print(f"\nAll class probabilities:")
                for label, prob in result['all_probabilities'].items():
                    print(f"  - {label}: {prob:.4f}")
        
        elif args.input_file:
            # File prediction
            input_path = Path(args.input_file)
            if not input_path.exists():
                print(f"❌ Input file not found: {input_path}")
                sys.exit(1)
            
            print(f"=== File Prediction ===")
            print(f"Input file: {args.input_file}")
            print(f"Output file: {args.output_file}")
            print()
            
            if args.output_file:
                # Use batch inference for large files
                if input_path.stat().st_size > 10 * 1024 * 1024:  # > 10MB
                    print("Large file detected, using chunked processing...")
                    batch_inference = BatchInference(config)
                    chunk_size = args.chunk_size or 1000
                    batch_inference.predict_large_file(
                        args.input_file, 
                        args.output_file, 
                        chunk_size=chunk_size
                    )
                else:
                    # Regular file processing
                    results = inference.predict_file(args.input_file, args.output_file)
                    print(f"Processed {len(results)} texts")
            else:
                # Just predict without saving
                results = inference.predict_file(args.input_file)
                print(f"Processed {len(results)} texts")
                print("\nSample results:")
                for i, result in enumerate(results[:3]):  # Show first 3
                    print(f"  {i+1}. '{result['text'][:50]}...' -> {result['predicted_label']} ({result['confidence']:.4f})")
        
        else:
            # Interactive mode - example predictions
            print("=== Interactive Mode ===")
            print("No input specified. Running example predictions...")
            print()
            
            # Example texts
            example_texts = [
                "I love this product! It's amazing.",
                "This is terrible, I hate it.",
                "The weather is okay today.",
                "Best purchase ever made!"
            ]
            
            print("Example predictions:")
            for text in example_texts:
                result = inference.predict_single(text)
                print(f"  '{text}' -> {result['predicted_label']} ({result['confidence']:.4f})")
        
        print(f"\n✅ Inference completed successfully!")
        
    except Exception as e:
        print(f"❌ Inference failed: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()

"""
# Quick usage examples:

# 1. Single prediction
config = InferenceConfig(model_path="./results/my_model")
inference = ModelInference(config)
result = inference.predict_single("Your text here")

# # 2. Batch prediction
# results = inference.predict_batch(["text1", "text2", "text3"])

# # 3. File prediction
# inference.predict_file("input.txt", "predictions.jsonl")

# # 4. Large file processing
# batch_inference = BatchInference(config)
# batch_inference.predict_large_file("large_input.jsonl", "large_output.jsonl", chunk_size=1000)
"""