model/generate.py

import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import logging
import psutil
import re
import gc

# Initialize logger
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

# List of memory-optimized models
MEMORY_OPTIMIZED_MODELS = [
    "gpt2",  # ~500MB
    "distilgpt2",  # ~250MB
    "microsoft/DialoGPT-small",  # ~250MB
    "huggingface/CodeBERTa-small-v1",  # Code tasks
]

# Singleton state
_generator_instance = None

def get_optimal_model_for_memory():
    """Select the best model based on available memory."""
    available_memory = psutil.virtual_memory().available / (1024 * 1024)  # MB
    logger.info(f"Available memory: {available_memory:.1f}MB")

    if available_memory < 300:
        return None  # Use template fallback
    elif available_memory < 600:
        return "microsoft/DialoGPT-small"
    else:
        return "distilgpt2"

def load_model_with_memory_optimization(model_name):
    """Load model with low memory settings."""
    try:
        logger.info(f"Loading {model_name} with memory optimizations...")

        tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left', use_fast=True)

        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16,
            device_map="cpu",
            low_cpu_mem_usage=True,
            use_cache=False,
        )

        model.eval()
        model.gradient_checkpointing_enable()
        logger.info(f"✅ Model {model_name} loaded successfully")
        return tokenizer, model

    except Exception as e:
        logger.error(f"❌ Failed to load model {model_name}: {e}")
        return None, None

def extract_keywords(text):
    common_keywords = [
        'login', 'authentication', 'user', 'password', 'database', 'data',
        'interface', 'api', 'function', 'feature', 'requirement', 'system',
        'input', 'output', 'validation', 'error', 'security', 'performance'
    ]
    words = re.findall(r'\b\w+\b', text.lower())
    return [word for word in words if word in common_keywords]

def generate_template_based_test_cases(srs_text):
    keywords = extract_keywords(srs_text)
    test_cases = []

    if any(word in keywords for word in ['login', 'authentication', 'user', 'password']):
        test_cases.extend([
            {
                "id": "TC_001",
                "title": "Valid Login Test",
                "description": "Test login with valid credentials",
                "steps": ["Enter valid username", "Enter valid password", "Click login"],
                "expected": "User should be logged in successfully"
            },
            {
                "id": "TC_002",
                "title": "Invalid Login Test",
                "description": "Test login with invalid credentials",
                "steps": ["Enter invalid username", "Enter invalid password", "Click login"],
                "expected": "Error message should be displayed"
            }
        ])

    if any(word in keywords for word in ['database', 'data', 'store', 'save']):
        test_cases.append({
            "id": "TC_003",
            "title": "Data Storage Test",
            "description": "Test data storage functionality",
            "steps": ["Enter data", "Save data", "Verify storage"],
            "expected": "Data should be stored correctly"
        })

    if not test_cases:
        test_cases = [
            {
                "id": "TC_001",
                "title": "Basic Functionality Test",
                "description": "Test basic system functionality",
                "steps": ["Access the system", "Perform basic operations", "Verify results"],
                "expected": "System should work as expected"
            }
        ]

    return test_cases

def parse_generated_test_cases(generated_text):
    lines = generated_text.split('\n')
    test_cases = []
    current_case = {}
    case_counter = 1

    for line in lines:
        line = line.strip()
        if line.startswith(('1.', '2.', '3.', 'TC', 'Test')):
            if current_case:
                test_cases.append(current_case)
            current_case = {
                "id": f"TC_{case_counter:03d}",
                "title": line,
                "description": line,
                "steps": ["Execute the test"],
                "expected": "Test should pass"
            }
            case_counter += 1

    if current_case:
        test_cases.append(current_case)

    if not test_cases:
        return [{
            "id": "TC_001",
            "title": "Generated Test Case",
            "description": "Auto-generated test case based on requirements",
            "steps": ["Review requirements", "Execute test", "Verify results"],
            "expected": "Requirements should be met"
        }]

    return test_cases

def generate_with_ai_model(srs_text, tokenizer, model):
    max_input_length = 200
    if len(srs_text) > max_input_length:
        srs_text = srs_text[:max_input_length]

    prompt = f"""Generate test cases for this software requirement:
{srs_text}

Test Cases:
1."""

    try:
        inputs = tokenizer.encode(
            prompt,
            return_tensors="pt",
            max_length=150,
            truncation=True
        )

        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_new_tokens=100,
                num_return_sequences=1,
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
                use_cache=False,
            )

        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        del inputs, outputs
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        return parse_generated_test_cases(generated_text)

    except Exception as e:
        logger.error(f"❌ AI generation failed: {e}")
        raise

def generate_with_fallback(srs_text):
    model_name = get_optimal_model_for_memory()

    if model_name:
        tokenizer, model = load_model_with_memory_optimization(model_name)
        if tokenizer and model:
            try:
                test_cases = generate_with_ai_model(srs_text, tokenizer, model)
                reason = get_algorithm_reason(model_name)
                return test_cases, model_name, "transformer (causal LM)", reason
            except Exception as e:
                logger.warning(f"AI generation failed: {e}, falling back to templates")

    logger.info("⚠️ Using fallback template-based generation")
    test_cases = generate_template_based_test_cases(srs_text)
    return test_cases, "Template-Based Generator", "rule-based", "Low memory - fallback to rule-based generation"

# ✅ Function exposed to app.py
def generate_test_cases(srs_text):
    return generate_with_fallback(srs_text)[0]

def get_generator():
    global _generator_instance
    if _generator_instance is None:
        class Generator:
            def __init__(self):
                self.model_name = get_optimal_model_for_memory()
                self.tokenizer = None
                self.model = None
                if self.model_name:
                    self.tokenizer, self.model = load_model_with_memory_optimization(self.model_name)

            def get_model_info(self):
                mem = psutil.Process().memory_info().rss / 1024 / 1024
                return {
                    "model_name": self.model_name if self.model_name else "Template-Based Generator",
                    "status": "loaded" if self.model else "template_mode",
                    "memory_usage": f"{mem:.1f}MB",
                    "optimization": "low_memory"
                }

        _generator_instance = Generator()

    return _generator_instance

def monitor_memory():
    mem = psutil.Process().memory_info().rss / 1024 / 1024
    logger.info(f"Memory usage: {mem:.1f}MB")
    if mem > 450:
        gc.collect()
        logger.info("Memory cleanup triggered")

# ✅ NEW FUNCTION for enhanced output: test cases + model info + reason
def generate_test_cases_and_info(input_text):
    test_cases, model_name, algorithm_used, reason = generate_with_fallback(input_text)
    return {
        "model": model_name,
        "algorithm": algorithm_used,
        "reason": reason,
        "test_cases": test_cases
    }

# ✅ Explain why each algorithm is selected
def get_algorithm_reason(model_name):
    if model_name == "microsoft/DialoGPT-small":
        return "Selected due to low memory availability; DialoGPT-small provides conversational understanding in limited memory environments."
    elif model_name == "distilgpt2":
        return "Selected for its balance between performance and low memory usage. Ideal for small environments needing causal language modeling."
    elif model_name == "gpt2":
        return "Chosen for general-purpose text generation with moderate memory headroom."
    elif model_name is None:
        return "No model used due to insufficient memory. Rule-based template generation chosen instead."
    else:
        return "Model selected based on best tradeoff between memory usage and language generation capability."