"""CodeBERT-based flakiness prediction model.

Based on FlakyLens (2025): 65.79% F1 score
Uses microsoft/codebert-base fine-tuned on test code.

Paper: "FlakyLens: Detecting Flaky Tests Using CodeBERT"
GitHub: https://github.com/flakylens/flakylens
"""
from __future__ import annotations

import logging
from pathlib import Path
from typing import Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)

# Check for PyTorch/Transformers availability
_TORCH_AVAILABLE = False
_TRANSFORMERS_AVAILABLE = False

try:
    import torch
    _TORCH_AVAILABLE = True
except ImportError:
    torch = None

try:
    from transformers import AutoModel, AutoTokenizer
    _TRANSFORMERS_AVAILABLE = True
except ImportError:
    AutoModel = None
    AutoTokenizer = None


class FlakyCodeBERT:
    """CodeBERT model fine-tuned for flakiness prediction.

    Architecture:
    1. CodeBERT encoder (microsoft/codebert-base)
    2. Classification head (768 -> 256 -> 2)
    3. Binary output: [not_flaky, flaky]

    Usage:
        model = FlakyCodeBERT()
        prob, importances = model.predict(test_code)
        if prob > 0.5:
            print("Test is likely flaky!")
    """

    MODEL_NAME = "microsoft/codebert-base"
    MAX_LENGTH = 512

    def __init__(
        self,
        model_path: Optional[str] = None,
        device: Optional[str] = None,
    ):
        """Initialize the model.

        Args:
            model_path: Path to fine-tuned weights (or use base model)
            device: 'cuda' or 'cpu' (auto-detect if None)
        """
        if _TORCH_AVAILABLE:
            self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        else:
            self.device = "cpu"

        self.tokenizer = None
        self.model = None
        self.classifier = None
        self.model_path = model_path
        self._loaded = False

    @property
    def available(self) -> bool:
        """Check if model dependencies are available."""
        return _TORCH_AVAILABLE and _TRANSFORMERS_AVAILABLE

    def load(self) -> None:
        """Load the model (lazy loading)."""
        if self._loaded:
            return

        if not self.available:
            logger.warning(
                "PyTorch or Transformers not installed. "
                "Run: pip install torch transformers"
            )
            return

        logger.info(f"Loading CodeBERT model on {self.device}")

        # Load tokenizer and base model
        self.tokenizer = AutoTokenizer.from_pretrained(self.MODEL_NAME)
        base_model = AutoModel.from_pretrained(self.MODEL_NAME)

        # Create classification head
        self.model = base_model.to(self.device)
        self.classifier = torch.nn.Sequential(
            torch.nn.Linear(768, 256),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.1),
            torch.nn.Linear(256, 2),
        ).to(self.device)

        # Load fine-tuned weights if available
        if self.model_path and Path(self.model_path).exists():
            try:
                state_dict = torch.load(self.model_path, map_location=self.device, weights_only=True)
                self.classifier.load_state_dict(state_dict["classifier"])
                logger.info(f"Loaded fine-tuned weights from {self.model_path}")
            except Exception as e:
                logger.warning(f"Failed to load fine-tuned weights: {e}")

        self._loaded = True

    def predict(self, test_code: str) -> Tuple[float, Dict[str, float]]:
        """Predict flakiness probability for test code.

        Args:
            test_code: Python test code

        Returns:
            Tuple of (flakiness_probability, feature_importances)
        """
        if not self.available:
            # Fallback to heuristic prediction
            return self._heuristic_predict(test_code)

        self.load()

        if not self._loaded or self.tokenizer is None or self.model is None:
            return self._heuristic_predict(test_code)

        # Tokenize
        inputs = self.tokenizer(
            test_code,
            return_tensors="pt",
            max_length=self.MAX_LENGTH,
            truncation=True,
            padding=True,
        ).to(self.device)

        # Forward pass
        with torch.no_grad():
            outputs = self.model(**inputs)
            pooled = outputs.last_hidden_state[:, 0, :]  # [CLS] token
            logits = self.classifier(pooled)
            probs = torch.softmax(logits, dim=-1)

        flaky_prob = probs[0, 1].item()

        # Feature importance (attention analysis)
        importances = self._get_feature_importances(test_code)

        return flaky_prob, importances

    def predict_batch(
        self,
        test_codes: List[str],
    ) -> List[Tuple[float, Dict[str, float]]]:
        """Predict flakiness for multiple tests."""
        return [self.predict(code) for code in test_codes]

    def _heuristic_predict(self, test_code: str) -> Tuple[float, Dict[str, float]]:
        """Fallback heuristic prediction when ML model unavailable."""
        code_lower = test_code.lower()
        importances: Dict[str, float] = {}
        score = 0.0

        # Known flaky patterns
        flaky_patterns = [
            ("time.sleep", 0.15),
            ("asyncio.sleep", 0.1),
            ("random.", 0.15),
            ("datetime.now", 0.1),
            ("time.time", 0.1),
            ("os.environ", 0.05),
            ("threading.thread", 0.15),
            ("multiprocessing", 0.15),
            ("socket.", 0.1),
            ("requests.", 0.1),
            ("httpx.", 0.1),
            ("open(", 0.05),
            ("/tmp/", 0.1),
        ]

        for pattern, weight in flaky_patterns:
            if pattern in code_lower:
                score += weight
                importances[pattern] = weight

        # Cap at 0.95
        score = min(0.95, score)

        return score, importances

    def _get_feature_importances(self, test_code: str) -> Dict[str, float]:
        """Extract feature importances based on known patterns."""
        code_lower = test_code.lower()
        importances: Dict[str, float] = {}

        # Keywords that often indicate flakiness
        flaky_keywords = {
            "sleep": 1.0,
            "time": 0.8,
            "random": 1.0,
            "await": 0.6,
            "async": 0.6,
            "thread": 0.9,
            "process": 0.7,
            "timeout": 0.8,
            "retry": 0.7,
            "socket": 0.8,
            "http": 0.7,
            "environ": 0.6,
            "tmp": 0.7,
        }

        for keyword, weight in flaky_keywords.items():
            if keyword in code_lower:
                importances[keyword] = weight

        return importances

    def fine_tune(
        self,
        train_data: List[Tuple[str, bool]],
        val_data: List[Tuple[str, bool]],
        epochs: int = 5,
        learning_rate: float = 2e-5,
        output_path: str = "flaky_model.pt",
    ) -> Dict[str, float]:
        """Fine-tune the model on labeled data.

        Args:
            train_data: List of (code, is_flaky) tuples
            val_data: Validation data
            epochs: Training epochs
            learning_rate: Learning rate
            output_path: Where to save weights

        Returns:
            Training metrics
        """
        if not train_data:
            return {"error": "Empty training data"}

        if not self.available:
            return {"error": "PyTorch/Transformers not available"}

        self.load()

        if not self._loaded:
            return {"error": "Model not loaded"}

        # Create optimizer
        optimizer = torch.optim.AdamW(
            list(self.model.parameters()) + list(self.classifier.parameters()),
            lr=learning_rate,
        )
        criterion = torch.nn.CrossEntropyLoss()

        best_f1 = 0.0
        metrics: Dict[str, float] = {}

        for epoch in range(epochs):
            # Training
            self.model.train()
            self.classifier.train()

            train_loss = 0.0
            for code, is_flaky in train_data:
                optimizer.zero_grad()

                inputs = self.tokenizer(
                    code,
                    return_tensors="pt",
                    max_length=self.MAX_LENGTH,
                    truncation=True,
                    padding=True,
                ).to(self.device)

                outputs = self.model(**inputs)
                pooled = outputs.last_hidden_state[:, 0, :]
                logits = self.classifier(pooled)

                target = torch.tensor([int(is_flaky)]).to(self.device)
                loss = criterion(logits, target)

                loss.backward()
                optimizer.step()
                train_loss += loss.item()

            # Validation
            val_f1, val_acc = self._evaluate(val_data)
            logger.info(
                f"Epoch {epoch + 1}/{epochs} - "
                f"Loss: {train_loss / len(train_data):.4f}, "
                f"Val F1: {val_f1:.4f}, Val Acc: {val_acc:.4f}"
            )

            if val_f1 > best_f1:
                best_f1 = val_f1
                torch.save(
                    {"classifier": self.classifier.state_dict()},
                    output_path,
                )
                logger.info(f"Saved best model to {output_path}")

        metrics["best_f1"] = best_f1
        metrics["epochs"] = epochs
        return metrics

    def _evaluate(self, data: List[Tuple[str, bool]]) -> Tuple[float, float]:
        """Evaluate model on data."""
        if not self._loaded:
            return 0.0, 0.0

        self.model.eval()
        self.classifier.eval()

        tp = fp = tn = fn = 0

        for code, is_flaky in data:
            prob, _ = self.predict(code)
            pred = prob > 0.5

            if pred and is_flaky:
                tp += 1
            elif pred and not is_flaky:
                fp += 1
            elif not pred and is_flaky:
                fn += 1
            else:
                tn += 1

        precision = tp / (tp + fp) if (tp + fp) > 0 else 0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0
        f1 = (
            2 * precision * recall / (precision + recall)
            if (precision + recall) > 0
            else 0
        )
        accuracy = (tp + tn) / len(data) if data else 0

        return f1, accuracy

    def get_training_data_sources(self) -> List[Dict[str, str]]:
        """Get recommended training data sources."""
        return [
            {
                "name": "IDoFT",
                "description": "International Dataset of Flaky Tests",
                "url": "https://github.com/TestingResearchIllinois/idoft",
                "format": "CSV with test code and flaky labels",
            },
            {
                "name": "FlakeFlagger",
                "description": "Flaky test dataset from Google",
                "url": "https://github.com/AkhilaSriManasa/FlakeFlagger",
                "format": "JSON with test metadata",
            },
            {
                "name": "Hikaflow History",
                "description": "Your own flaky test history from Hikaflow",
                "url": "Export from /v1/flaky-tests/export",
                "format": "JSON with test code and flakiness scores",
            },
        ]


# Module-level singleton
_model: Optional[FlakyCodeBERT] = None


def get_flaky_model(model_path: Optional[str] = None) -> FlakyCodeBERT:
    """Get or create the flaky model singleton.

    Args:
        model_path: Optional path to fine-tuned weights

    Returns:
        FlakyCodeBERT instance
    """
    global _model
    if _model is None:
        _model = FlakyCodeBERT(model_path=model_path)
    return _model


def predict_flakiness(test_code: str) -> Tuple[float, Dict[str, float]]:
    """Convenience function to predict flakiness.

    Args:
        test_code: Python test code

    Returns:
        Tuple of (probability, feature_importances)
    """
    model = get_flaky_model()
    return model.predict(test_code)
