"""
Model Utilities - Helper functions for model operations
"""

import logging
from typing import Dict, Any, List, Optional, Tuple
import torch
from transformers import (
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForSeq2SeqLM,
    AutoModelForTokenClassification,
    AutoModelForQuestionAnswering,
    AutoModelForSequenceClassification,
    AutoConfig,
    AutoTokenizer,
    PreTrainedModel,
    PreTrainedTokenizer,
)
from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
import os
import json
import hashlib

logger = logging.getLogger(__name__)


# Model architectures and their supported tasks
MODEL_TASK_MAPPING = {
    "gpt": ["causal-lm"],
    "llama": ["causal-lm"],
    "mistral": ["causal-lm"],
    "falcon": ["causal-lm"],
    "qwen": ["causal-lm"],
    "phi": ["causal-lm"],
    "opt": ["causal-lm"],
    "bloom": ["causal-lm"],
    "t5": ["seq2seq"],
    "bart": ["seq2seq"],
    "pegasus": ["seq2seq"],
    "mt5": ["seq2seq"],
    "bert": ["token-classification", "text-classification", "question-answering"],
    "roberta": ["token-classification", "text-classification", "question-answering"],
    "deberta": ["token-classification", "text-classification", "question-answering"],
    "xlnet": ["token-classification", "text-classification", "question-answering"],
    "albert": ["token-classification", "text-classification", "question-answering"],
    "electra": ["token-classification", "text-classification"],
    "distilbert": ["token-classification", "text-classification", "question-answering"],
}


# PEFT task type mapping
PEFT_TASK_TYPES = {
    "causal-lm": TaskType.CAUSAL_LM,
    "seq2seq": TaskType.SEQ_2_SEQ_LM,
    "token-classification": TaskType.TOKEN_CLS,
    "text-classification": TaskType.SEQ_CLS,
    "question-answering": TaskType.QUESTION_ANS,
}


def get_model_for_task(model_name: str, task_type: str, **kwargs) -> Tuple[PreTrainedModel, Optional[str]]:
    """Load appropriate model for a task type."""
    try:
        config = AutoConfig.from_pretrained(model_name)
        
        # Determine model class
        if task_type == "causal-lm":
            model_class = AutoModelForCausalLM
        elif task_type == "seq2seq":
            model_class = AutoModelForSeq2SeqLM
        elif task_type == "token-classification":
            model_class = AutoModelForTokenClassification
        elif task_type == "text-classification":
            model_class = AutoModelForSequenceClassification
        elif task_type == "question-answering":
            model_class = AutoModelForQuestionAnswering
        else:
            model_class = AutoModel
        
        # Load model
        model = model_class.from_pretrained(
            model_name,
            config=config,
            **kwargs
        )
        
        return model, None
        
    except Exception as e:
        logger.error(f"Error loading model {model_name} for task {task_type}: {e}")
        return None, str(e)


def load_tokenizer(model_name: str, **kwargs) -> Tuple[PreTrainedTokenizer, Optional[str]]:
    """Load tokenizer for a model."""
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_name, **kwargs)
        
        # Ensure pad token is set
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token or "<pad>"
            tokenizer.pad_token_id = tokenizer.eos_token_id or tokenizer.convert_tokens_to_ids("<pad>")
        
        return tokenizer, None
        
    except Exception as e:
        logger.error(f"Error loading tokenizer for {model_name}: {e}")
        return None, str(e)


def get_model_info(model_name: str) -> Dict[str, Any]:
    """Get detailed model information."""
    try:
        from huggingface_hub import HfApi, model_info
        
        api = HfApi()
        info = api.model_info(model_name)
        
        # Try to load config for more details
        try:
            config = AutoConfig.from_pretrained(model_name)
            config_dict = config.to_dict()
        except:
            config_dict = {}
        
        return {
            "model_id": info.id,
            "author": info.author,
            "sha": info.sha,
            "pipeline_tag": info.pipeline_tag,
            "library_name": info.library_name,
            "downloads": getattr(info, "downloads", 0),
            "likes": getattr(info, "likes", 0),
            "tags": info.tags or [],
            "siblings": [s.rfilename for s in info.siblings] if info.siblings else [],
            "config": config_dict,
            "hidden_size": config_dict.get("hidden_size"),
            "num_hidden_layers": config_dict.get("num_hidden_layers"),
            "num_attention_heads": config_dict.get("num_attention_heads"),
            "intermediate_size": config_dict.get("intermediate_size"),
            "vocab_size": config_dict.get("vocab_size"),
            "model_type": config_dict.get("model_type"),
            "architectures": config_dict.get("architectures", []),
        }
        
    except Exception as e:
        logger.error(f"Error getting model info for {model_name}: {e}")
        return {"error": str(e)}


def check_model_compatibility(model_name: str, task_type: str) -> Tuple[bool, List[str]]:
    """Check if model is compatible with a task type."""
    issues = []
    
    try:
        config = AutoConfig.from_pretrained(model_name)
        architectures = config.architectures or []
        model_type = config.model_type or ""
        
        # Check if architecture supports task
        compatible = True
        
        if task_type == "causal-lm":
            causal_archs = ["GPT", "LLaMA", "Mistral", "Falcon", "Qwen", "Phi", "OPT", "Bloom", "CausalLM"]
            if not any(arch in arch for arch in architectures for arch in causal_archs):
                if model_type not in ["gpt2", "llama", "mistral", "falcon", "qwen", "phi"]:
                    issues.append("Model may not support causal language modeling")
                    
        elif task_type == "seq2seq":
            seq2seq_archs = ["T5", "BART", "Pegasus", "MT5", "EncoderDecoderModel"]
            if not any(arch in arch for arch in architectures for arch in seq2seq_archs):
                issues.append("Model may not support seq2seq tasks")
                
        elif task_type == "token-classification":
            if not any("TokenClassification" in arch for arch in architectures):
                issues.append("Model may not support token classification")
                
        elif task_type == "text-classification":
            if not any("Classification" in arch for arch in architectures):
                issues.append("Model may not support text classification")
                
        elif task_type == "question-answering":
            qa_archs = ["QuestionAnswering", "BertForQA"]
            if not any(arch in arch for arch in architectures for arch in qa_archs):
                issues.append("Model may not support question answering")
        
        return len(issues) == 0, issues
        
    except Exception as e:
        return False, [f"Error checking compatibility: {str(e)}"]


def apply_peft(
    model: PreTrainedModel,
    task_type: str,
    lora_r: int = 8,
    lora_alpha: int = 32,
    lora_dropout: float = 0.1,
    target_modules: Optional[List[str]] = None,
) -> Tuple[PreTrainedModel, Dict[str, Any]]:
    """Apply PEFT/LoRA to a model."""
    try:
        # Prepare model for training
        model = prepare_model_for_kbit_training(model)
        
        # Get PEFT task type
        peft_task_type = PEFT_TASK_TYPES.get(task_type, TaskType.CAUSAL_LM)
        
        # Auto-detect target modules if not specified
        if target_modules is None:
            model_type = getattr(model.config, "model_type", "").lower()
            if "llama" in model_type or "mistral" in model_type:
                target_modules = ["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
            elif "gpt" in model_type:
                target_modules = ["c_attn", "c_proj"]
            elif "bert" in model_type or "roberta" in model_type:
                target_modules = ["query", "value", "key", "dense"]
            else:
                target_modules = ["q_proj", "v_proj"]
        
        # Create LoRA config
        lora_config = LoraConfig(
            r=lora_r,
            lora_alpha=lora_alpha,
            lora_dropout=lora_dropout,
            bias="none",
            task_type=peft_task_type,
            target_modules=target_modules,
        )
        
        # Apply LoRA
        model = get_peft_model(model, lora_config)
        
        # Get trainable params info
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        all_params = sum(p.numel() for p in model.parameters())
        
        info = {
            "trainable_params": trainable_params,
            "all_params": all_params,
            "trainable_percentage": 100 * trainable_params / all_params,
            "lora_r": lora_r,
            "lora_alpha": lora_alpha,
            "target_modules": target_modules,
        }
        
        return model, info
        
    except Exception as e:
        logger.error(f"Error applying PEFT: {e}")
        return model, {"error": str(e)}


def estimate_parameters(model_name: str) -> Dict[str, Any]:
    """Estimate model parameters without loading."""
    try:
        config = AutoConfig.from_pretrained(model_name)
        
        hidden_size = getattr(config, "hidden_size", 768)
        num_layers = getattr(config, "num_hidden_layers", 12)
        num_heads = getattr(config, "num_attention_heads", 12)
        vocab_size = getattr(config, "vocab_size", 30522)
        intermediate_size = getattr(config, "intermediate_size", hidden_size * 4)
        
        # Rough estimation formulas
        # Embedding params
        embedding_params = vocab_size * hidden_size
        
        # Attention params per layer (Q, K, V, O projections)
        attention_params = 4 * hidden_size * hidden_size * num_layers
        
        # FFN params per layer
        ffn_params = (hidden_size * intermediate_size + intermediate_size * hidden_size) * num_layers
        
        # Layer norm params
        layernorm_params = 2 * hidden_size * num_layers
        
        total_params = embedding_params + attention_params + ffn_params + layernorm_params
        
        return {
            "estimated_params": total_params,
            "estimated_params_billions": round(total_params / 1e9, 2),
            "hidden_size": hidden_size,
            "num_layers": num_layers,
            "num_heads": num_heads,
            "vocab_size": vocab_size,
            "model_size_mb": round(total_params * 4 / (1024 * 1024), 2),  # FP32
            "model_size_mb_fp16": round(total_params * 2 / (1024 * 1024), 2),  # FP16
        }
        
    except Exception as e:
        logger.warning(f"Could not estimate parameters: {e}")
        return {
            "estimated_params": 0,
            "estimated_params_billions": 0,
            "error": str(e),
        }


def get_recommended_settings(model_name: str, task_type: str) -> Dict[str, Any]:
    """Get recommended training settings for a model."""
    info = estimate_parameters(model_name)
    params_b = info.get("estimated_params_billions", 0.1)
    
    # Base recommendations
    settings = {
        "batch_size": 1,
        "gradient_accumulation_steps": 1,
        "learning_rate": "5e-5",
        "epochs": 3,
        "max_length": 512,
        "use_peft": False,
        "lora_r": 8,
        "warmup_ratio": 0.1,
    }
    
    # Adjust based on model size
    if params_b > 7:  # > 7B parameters
        settings["batch_size"] = 1
        settings["gradient_accumulation_steps"] = 8
        settings["learning_rate"] = "1e-5"
        settings["use_peft"] = True
        settings["lora_r"] = 8
        settings["max_length"] = 256
        
    elif params_b > 3:  # > 3B parameters
        settings["batch_size"] = 1
        settings["gradient_accumulation_steps"] = 4
        settings["learning_rate"] = "2e-5"
        settings["use_peft"] = True
        settings["max_length"] = 512
        
    elif params_b > 1:  # > 1B parameters
        settings["batch_size"] = 2
        settings["gradient_accumulation_steps"] = 2
        settings["use_peft"] = True
        
    else:  # < 1B parameters
        settings["batch_size"] = 4
        settings["gradient_accumulation_steps"] = 1
        settings["use_peft"] = False
    
    # Task-specific adjustments
    if task_type == "seq2seq":
        settings["max_length"] = 1024
        settings["epochs"] = 5
        
    elif task_type == "token-classification":
        settings["max_length"] = 128
        settings["learning_rate"] = "2e-5"
        
    elif task_type == "text-classification":
        settings["epochs"] = 3
        settings["learning_rate"] = "3e-5"
        
    elif task_type == "question-answering":
        settings["max_length"] = 384
        settings["batch_size"] = 8
    
    return settings


def count_parameters(model: PreTrainedModel) -> Dict[str, int]:
    """Count model parameters."""
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    frozen = total - trainable
    
    return {
        "trainable": trainable,
        "frozen": frozen,
        "total": total,
        "trainable_percentage": 100 * trainable / total if total > 0 else 0,
    }


def get_model_memory_footprint(model: PreTrainedModel) -> Dict[str, float]:
    """Get model memory footprint in MB."""
    param_size = sum(p.numel() * p.element_size() for p in model.parameters())
    buffer_size = sum(b.numel() * b.element_size() for b in model.buffers())
    
    return {
        "parameters_mb": param_size / (1024 * 1024),
        "buffers_mb": buffer_size / (1024 * 1024),
        "total_mb": (param_size + buffer_size) / (1024 * 1024),
    }


def save_model_with_metadata(
    model: PreTrainedModel,
    tokenizer: PreTrainedTokenizer,
    output_dir: str,
    training_config: Dict[str, Any],
    metrics: Dict[str, float],
) -> Dict[str, str]:
    """Save model with comprehensive metadata."""
    import json
    from datetime import datetime
    
    os.makedirs(output_dir, exist_ok=True)
    
    # Save model and tokenizer
    model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)
    
    # Get model info
    param_info = count_parameters(model)
    memory_info = get_model_memory_footprint(model)
    
    # Create comprehensive metadata
    metadata = {
        "model_name": training_config.get("model_name", "unknown"),
        "task_type": training_config.get("task_type", "unknown"),
        "training_config": training_config,
        "metrics": metrics,
        "parameter_info": param_info,
        "memory_info": memory_info,
        "created_at": datetime.utcnow().isoformat(),
        "transformers_version": __import__("transformers").__version__,
        "torch_version": __import__("torch").__version__,
        "python_version": __import__("sys").version,
    }
    
    # Save metadata
    metadata_path = os.path.join(output_dir, "training_metadata.json")
    with open(metadata_path, "w") as f:
        json.dump(metadata, f, indent=2)
    
    # Create model card
    model_card = create_model_card(training_config, metrics, param_info)
    model_card_path = os.path.join(output_dir, "README.md")
    with open(model_card_path, "w") as f:
        f.write(model_card)
    
    return {
        "output_dir": output_dir,
        "model_path": output_dir,
        "metadata_path": metadata_path,
        "model_card_path": model_card_path,
    }


def create_model_card(
    config: Dict[str, Any],
    metrics: Dict[str, float],
    param_info: Dict[str, int],
) -> str:
    """Create a model card README."""
    model_name = config.get("model_name", "unknown")
    task_type = config.get("task_type", "unknown")
    
    metrics_str = "\n".join([f"- {k}: {v:.4f}" if isinstance(v, float) else f"- {k}: {v}" for k, v in metrics.items()]) if metrics else "- No metrics available"
    
    return f"""# {model_name} - Fine-tuned

## Model Details

- **Base Model:** {model_name}
- **Task:** {task_type}
- **Total Parameters:** {param_info.get('total', 0):,}
- **Trainable Parameters:** {param_info.get('trainable', 0):,}

## Training Configuration

```json
{json.dumps(config, indent=2)}
```

## Training Metrics

{metrics_str}

## Usage

```python
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("path/to/model")
tokenizer = AutoTokenizer.from_pretrained("path/to/model")
```

## License

Please refer to the original model's license.

## Training Framework

This model was trained using the Universal Model Trainer.
"""