#!/usr/bin/env python3
"""
Enhanced Ollama setup script that includes SIGMA model creation.
This integrates the functionality from fix_sigma_model.sh into the Docker container.
"""
import os
import json
import time
import requests
import subprocess
import sys
import tempfile
import shutil
from typing import Dict, List, Optional
from pathlib import Path

OLLAMA_BASE_URL = os.getenv('OLLAMA_BASE_URL', 'http://ollama:11434')
DEFAULT_MODEL = os.getenv('LLM_MODEL', 'llama3.2')
SIGMA_MODEL_NAME = 'sigma-llama'

def log(message: str, level: str = "INFO"):
    """Log message with timestamp"""
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] {level}: {message}")

def wait_for_ollama(max_retries: int = 30, delay: int = 5) -> bool:
    """Wait for Ollama service to be ready"""
    log("Waiting for Ollama service to be ready...")
    
    for attempt in range(max_retries):
        try:
            response = requests.get(f"{OLLAMA_BASE_URL}/api/tags", timeout=10)
            if response.status_code == 200:
                log("✅ Ollama service is ready!")
                return True
        except requests.exceptions.RequestException as e:
            log(f"Attempt {attempt + 1}/{max_retries}: Ollama not ready yet ({e})", "DEBUG")
        
        if attempt < max_retries - 1:
            time.sleep(delay)
    
    log("❌ Ollama service failed to become ready", "ERROR")
    return False

def get_available_models() -> List[str]:
    """Get list of available models"""
    try:
        response = requests.get(f"{OLLAMA_BASE_URL}/api/tags", timeout=10)
        if response.status_code == 200:
            data = response.json()
            models = [model.get('name', '') for model in data.get('models', [])]
            log(f"Available models: {models}")
            return models
        else:
            log(f"Failed to get models: HTTP {response.status_code}", "ERROR")
            return []
    except Exception as e:
        log(f"Error getting models: {e}", "ERROR")
        return []

def pull_model(model_name: str) -> bool:
    """Pull a model if not available"""
    log(f"Pulling model: {model_name}")
    
    try:
        payload = {"name": model_name}
        response = requests.post(
            f"{OLLAMA_BASE_URL}/api/pull",
            json=payload,
            timeout=600,  # 10 minutes timeout
            stream=True
        )
        
        if response.status_code == 200:
            # Stream and log progress
            for line in response.iter_lines():
                if line:
                    try:
                        data = json.loads(line.decode('utf-8'))
                        status = data.get('status', '')
                        if status:
                            log(f"Pull progress: {status}", "DEBUG")
                        if data.get('error'):
                            log(f"Pull error: {data.get('error')}", "ERROR")
                            return False
                    except json.JSONDecodeError:
                        continue
            
            log(f"✅ Successfully pulled model: {model_name}")
            return True
        else:
            log(f"❌ Failed to pull model {model_name}: HTTP {response.status_code}", "ERROR")
            return False
            
    except Exception as e:
        log(f"❌ Error pulling model {model_name}: {e}", "ERROR")
        return False

def get_best_checkpoint() -> str:
    """Get the best performing checkpoint (highest step count)"""
    finetuned_model_path = "/app/models/sigma_llama_finetuned"
    
    # Look for the highest checkpoint number
    checkpoints = []
    try:
        for item in os.listdir(finetuned_model_path):
            if item.startswith('checkpoint-') and os.path.isdir(os.path.join(finetuned_model_path, item)):
                try:
                    checkpoint_num = int(item.split('-')[1])
                    checkpoints.append((checkpoint_num, item))
                except ValueError:
                    continue
        
        if checkpoints:
            # Sort by checkpoint number and get the highest
            checkpoints.sort(key=lambda x: x[0], reverse=True)
            best_checkpoint = checkpoints[0][1]
            log(f"Found best checkpoint: {best_checkpoint} (step {checkpoints[0][0]})")
            return os.path.join(finetuned_model_path, best_checkpoint)
    except Exception as e:
        log(f"Error scanning checkpoints: {e}", "WARN")
    
    # Fallback to root directory
    return finetuned_model_path

def create_lora_modelfile(adapter_path: str) -> str:
    """Create a Modelfile that uses the LoRA adapter"""
    
    # Create optimized Modelfile for LoRA adapter
    modelfile_content = f"""FROM {DEFAULT_MODEL}:latest

# LoRA Adapter Configuration
# Note: Direct ADAPTER directive requires GGML format
# Using enhanced prompting optimized for fine-tuned model

TEMPLATE \"\"\"<|begin_of_text|><|start_header_id|>system<|end_header_id|>

You are a cybersecurity expert specializing in SIGMA rule creation. You have been fine-tuned specifically for generating high-quality SIGMA detection rules.

Generate valid SIGMA rules in YAML format based on the provided CVE and exploit information. Output ONLY valid YAML starting with 'title:' and ending with the last YAML line.

Focus on:
- Accurate logsource identification
- Precise detection logic
- Relevant fields and values
- Proper YAML formatting
- Contextual understanding from CVE details<|eot_id|><|start_header_id|>user<|end_header_id|>

{{ .Prompt }}<|eot_id|><|start_header_id|>assistant<|end_header_id|>

\"\"\"

# Fine-tuned parameters optimized for SIGMA rule generation
PARAMETER temperature 0.1
PARAMETER top_p 0.9
PARAMETER top_k 40
PARAMETER repeat_penalty 1.1
PARAMETER num_ctx 4096
PARAMETER stop "<|eot_id|>"
PARAMETER stop "<|end_of_text|>"

# System message optimized for fine-tuned model
SYSTEM \"\"\"You are a specialized SIGMA rule generation model. Your training has optimized you for creating accurate, contextual SIGMA detection rules. Generate only valid YAML format rules based on the provided context.\"\"\"
"""
    
    return modelfile_content

def create_sigma_model() -> bool:
    """Create the sigma-llama model using the LoRA fine-tuned model"""
    log("🔄 Creating sigma-llama model from LoRA fine-tuned model...")
    
    # Check if fine-tuned model exists
    finetuned_model_path = "/app/models/sigma_llama_finetuned"
    
    if not os.path.exists(finetuned_model_path):
        log(f"❌ Fine-tuned model not found at {finetuned_model_path}", "WARN")
        log("Falling back to prompt-tuned base model...", "WARN")
        return create_prompt_tuned_model()
    
    log(f"✅ Found LoRA fine-tuned model at {finetuned_model_path}")
    
    # Get the best checkpoint
    best_checkpoint_path = get_best_checkpoint()
    log(f"Using checkpoint: {best_checkpoint_path}")
    
    # Check if adapter files exist
    adapter_file = os.path.join(best_checkpoint_path, "adapter_model.safetensors")
    if not os.path.exists(adapter_file):
        log(f"❌ Adapter file not found at {adapter_file}", "ERROR")
        return create_prompt_tuned_model()
    
    log(f"✅ Found LoRA adapter: {adapter_file}")
    
    # Remove any existing sigma-llama model
    try:
        response = requests.delete(f"{OLLAMA_BASE_URL}/api/delete", 
                                 json={"name": SIGMA_MODEL_NAME},
                                 timeout=30)
        if response.status_code == 200:
            log("Removed existing sigma-llama model")
    except Exception:
        pass  # Model might not exist, that's fine
    
    # Create optimized Modelfile for LoRA fine-tuned model
    modelfile_content = create_lora_modelfile(best_checkpoint_path)
    
    try:
        log("Creating LoRA-optimized sigma-llama model...")
        log("Note: Using optimized prompting for LoRA fine-tuned model")
        log("Note: Direct LoRA loading requires GGML format - using enhanced prompting approach")
        
        # Create the model with LoRA-optimized configuration
        # Use the from parameter approach since Modelfile with FROM is not working
        payload = {
            "name": SIGMA_MODEL_NAME,
            "from": f"{DEFAULT_MODEL}:latest",
            "modelfile": modelfile_content.replace(f"FROM {DEFAULT_MODEL}:latest\n\n", ""),
            "stream": False
        }
        
        response = requests.post(
            f"{OLLAMA_BASE_URL}/api/create",
            json=payload,
            timeout=600,  # 10 minutes timeout
            stream=True
        )
        
        if response.status_code == 200:
            # Stream and log progress
            for line in response.iter_lines():
                if line:
                    try:
                        data = json.loads(line.decode('utf-8'))
                        status = data.get('status', '')
                        if status:
                            log(f"LoRA model creation: {status}", "DEBUG")
                        if data.get('error'):
                            log(f"Model creation error: {data.get('error')}", "ERROR")
                            return False
                    except json.JSONDecodeError:
                        continue
            
            # Verify the model was created
            models = get_available_models()
            if any(SIGMA_MODEL_NAME in model for model in models):
                log("✅ sigma-llama LoRA-optimized model created successfully!")
                return True
            else:
                log("❌ sigma-llama model not found after creation", "ERROR")
                return False
        else:
            log(f"❌ Failed to create sigma-llama model: HTTP {response.status_code}", "ERROR")
            try:
                error_data = response.json()
                log(f"Error details: {error_data}", "ERROR")
            except:
                log(f"Error response: {response.text}", "ERROR")
            return False
            
    except Exception as e:
        log(f"❌ Error creating LoRA sigma-llama model: {e}", "ERROR")
        return False

def create_prompt_tuned_model() -> bool:
    """Fallback: Create a prompt-tuned model using base llama model"""
    log("🔄 Creating prompt-tuned sigma-llama model as fallback...")
    
    # Create Modelfile content for prompt-tuned model
    modelfile_content = """TEMPLATE \"\"\"### Instruction:
Generate SIGMA rule logsource and detection sections based on the provided context.

### Input:
{{ .Prompt }}

### Response:
\"\"\"

PARAMETER temperature 0.1
PARAMETER top_p 0.9
PARAMETER stop "### Instruction:"
PARAMETER stop "### Response:"
PARAMETER num_ctx 4096

SYSTEM \"\"\"You are a cybersecurity expert specializing in SIGMA rule creation. Generate valid SIGMA rules in YAML format based on the provided CVE and exploit information. Output ONLY valid YAML starting with 'title:' and ending with the last YAML line.\"\"\"
"""
    
    try:
        # Create the model using the base model
        payload = {
            "name": SIGMA_MODEL_NAME,
            "from": f"{DEFAULT_MODEL}:latest",
            "modelfile": modelfile_content,
            "stream": False
        }
        
        response = requests.post(
            f"{OLLAMA_BASE_URL}/api/create",
            json=payload,
            timeout=300,  # 5 minutes timeout
            stream=True
        )
        
        if response.status_code == 200:
            # Stream and log progress
            for line in response.iter_lines():
                if line:
                    try:
                        data = json.loads(line.decode('utf-8'))
                        status = data.get('status', '')
                        if status:
                            log(f"Prompt-tuned model creation: {status}", "DEBUG")
                        if data.get('error'):
                            log(f"Model creation error: {data.get('error')}", "ERROR")
                            return False
                    except json.JSONDecodeError:
                        continue
            
            # Verify the model was created
            models = get_available_models()
            if any(SIGMA_MODEL_NAME in model for model in models):
                log("✅ sigma-llama prompt-tuned model created successfully!")
                return True
            else:
                log("❌ sigma-llama model not found after creation", "ERROR")
                return False
        else:
            log(f"❌ Failed to create sigma-llama model: HTTP {response.status_code}", "ERROR")
            return False
            
    except Exception as e:
        log(f"❌ Error creating prompt-tuned model: {e}", "ERROR")
        return False

def test_sigma_model() -> bool:
    """Test the sigma-llama LoRA model"""
    log("🔄 Testing sigma-llama LoRA model...")
    
    try:
        # Use a more comprehensive test prompt
        test_prompt = """Generate a SIGMA rule for CVE-2023-1234: PowerShell command execution vulnerability that allows remote code execution through malicious PowerShell scripts.
        
Vulnerability Details:
- Affects Windows PowerShell
- Remote code execution via script injection
- Commonly exploited through phishing emails
- Targets process execution and command line arguments"""
        
        test_payload = {
            "model": SIGMA_MODEL_NAME,
            "prompt": test_prompt,
            "stream": False
        }
        
        response = requests.post(
            f"{OLLAMA_BASE_URL}/api/generate",
            json=test_payload,
            timeout=120  # Longer timeout for LoRA model
        )
        
        if response.status_code == 200:
            data = response.json()
            test_response = data.get('response', '')
            
            # Check if response looks like a SIGMA rule
            if 'title:' in test_response.lower() and ('logsource:' in test_response.lower() or 'detection:' in test_response.lower()):
                log(f"✅ LoRA model test successful! Generated SIGMA rule structure detected.")
                log(f"Response preview: {test_response[:300]}...")
                return True
            else:
                log(f"⚠️ Model responded but output doesn't look like SIGMA rule: {test_response[:200]}...")
                return False
        else:
            log(f"❌ Model test failed: HTTP {response.status_code}", "ERROR")
            return False
            
    except Exception as e:
        log(f"❌ Error testing LoRA model: {e}", "ERROR")
        return False

def main():
    """Main setup function"""
    log("🚀 Starting enhanced Ollama setup with fine-tuned SIGMA model creation...")
    
    # Step 1: Wait for Ollama to be ready
    if not wait_for_ollama():
        log("❌ Setup failed: Ollama service not available", "ERROR")
        sys.exit(1)
    
    # Step 2: Check current models
    models = get_available_models()
    log(f"Current models: {models}")
    
    # Step 3: Pull default model if needed (for fallback)
    if not any(DEFAULT_MODEL in model for model in models):
        log(f"Default model {DEFAULT_MODEL} not found, pulling for fallback...")
        if not pull_model(DEFAULT_MODEL):
            log(f"❌ Setup failed: Could not pull {DEFAULT_MODEL}", "ERROR")
            sys.exit(1)
    else:
        log(f"✅ Default model {DEFAULT_MODEL} already available")
    
    # Step 4: Create SIGMA model (will try fine-tuned first, then fallback)
    if not create_sigma_model():
        log("❌ Setup failed: Could not create sigma-llama model", "ERROR")
        sys.exit(1)
    
    # Step 5: Test SIGMA model
    if not test_sigma_model():
        log("⚠️  Setup warning: sigma-llama model test failed", "WARN")
        # Don't exit here, the model might still work
    
    # Step 6: Final verification
    final_models = get_available_models()
    log(f"Final models available: {final_models}")
    
    if any(SIGMA_MODEL_NAME in model for model in final_models):
        log("🎉 Setup complete! LoRA fine-tuned sigma-llama model is ready for use.")
        log("📊 Model benefits: Enhanced SIGMA rule generation with domain-specific fine-tuning")
        log("🔧 Monitor performance: Compare outputs with base model for quality improvements")
        sys.exit(0)
    else:
        log("❌ Setup failed: sigma-llama model not available after setup", "ERROR")
        sys.exit(1)

if __name__ == "__main__":
    main()