#!/usr/bin/env python3
"""
Convert LoRA adapter to GGUF format for better Ollama integration
"""

import os
import sys
import subprocess
import tempfile
import shutil
from pathlib import Path
import logging

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

def check_llama_cpp_tools():
    """Check if llama.cpp tools are available"""
    tools_needed = [
        'convert_hf_to_gguf.py',
        'convert_lora_to_gguf.py', 
        'llama-export-lora'
    ]
    
    missing_tools = []
    for tool in tools_needed:
        if not shutil.which(tool):
            missing_tools.append(tool)
    
    if missing_tools:
        logger.error(f"Missing required tools: {missing_tools}")
        logger.error("Please install llama.cpp and ensure tools are in PATH")
        logger.error("See: https://github.com/ggerganov/llama.cpp")
        return False
    
    return True

def convert_lora_to_gguf(lora_path: str, base_model_name: str = "meta-llama/Llama-3.2-3B-Instruct"):
    """Convert LoRA adapter to GGUF format"""
    
    if not check_llama_cpp_tools():
        logger.error("Cannot convert LoRA - missing llama.cpp tools")
        return None
    
    try:
        # Create temporary directory for conversion
        with tempfile.TemporaryDirectory() as temp_dir:
            logger.info(f"Converting LoRA adapter from {lora_path}")
            
            # Step 1: Convert LoRA to GGUF
            lora_gguf_path = os.path.join(temp_dir, "adapter.gguf")
            
            cmd = [
                'convert_lora_to_gguf.py',
                '--base', base_model_name,
                '--outtype', 'q8_0',  # High quality quantization
                '--outfile', lora_gguf_path,
                lora_path
            ]
            
            logger.info(f"Running: {' '.join(cmd)}")
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            if result.returncode != 0:
                logger.error(f"LoRA conversion failed: {result.stderr}")
                return None
            
            # Step 2: Download/prepare base model in GGUF format
            base_model_path = os.path.join(temp_dir, "base_model.gguf")
            
            # This would need to download the base model in GGUF format
            # For now, we'll assume it's available or use Ollama's version
            
            # Step 3: Merge LoRA with base model
            merged_model_path = os.path.join(temp_dir, "merged_model.gguf")
            
            merge_cmd = [
                'llama-export-lora',
                '-m', base_model_path,
                '-o', merged_model_path,
                '--lora', lora_gguf_path
            ]
            
            logger.info(f"Running: {' '.join(merge_cmd)}")
            merge_result = subprocess.run(merge_cmd, capture_output=True, text=True)
            
            if merge_result.returncode != 0:
                logger.error(f"LoRA merge failed: {merge_result.stderr}")
                return None
            
            # Copy merged model to output location
            output_path = "/app/models/sigma_llama_merged.gguf"
            shutil.copy2(merged_model_path, output_path)
            
            logger.info(f"✅ Successfully converted and merged LoRA to {output_path}")
            return output_path
            
    except Exception as e:
        logger.error(f"Error converting LoRA: {e}")
        return None

def create_gguf_modelfile(gguf_path: str) -> str:
    """Create Modelfile for GGUF merged model"""
    
    modelfile_content = f"""FROM {gguf_path}

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|>

\"\"\"

# Optimized parameters 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 \"\"\"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

if __name__ == "__main__":
    # Test conversion
    lora_path = "/app/models/sigma_llama_finetuned/checkpoint-2268"
    
    if os.path.exists(lora_path):
        converted_path = convert_lora_to_gguf(lora_path)
        if converted_path:
            print(f"✅ Conversion successful: {converted_path}")
            print("Use this path in your Ollama Modelfile:")
            print(create_gguf_modelfile(converted_path))
        else:
            print("❌ Conversion failed")
    else:
        print(f"❌ LoRA path not found: {lora_path}")