#!/usr/bin/env python3
"""
Knowledge Graph Perturbation Tester using LiteLLM

This script loads a knowledge graph, tests relations for perturbation resistance,
and enriches the graph with test results using different LLM providers via LiteLLM.
All data is stored in the database, not in local files.
"""

# LiteLLM Monkey Patch to fix "Unsupported parameter" errors
from utils.fix_litellm_stop_param import patched_completion

import json
import argparse
import os
import time
import random
import logging
from typing import Dict, List, Any, Tuple, Optional, Union, Callable
import yaml
import sys
from datetime import datetime
import litellm
from utils.config import OPENAI_API_KEY
import uuid

# Configure logging for this module
logger = logging.getLogger(__name__)

# Import LiteLLM
from litellm import completion

# API Key setup - only set if value exists
# Only set environment variable if OPENAI_API_KEY is not None
if OPENAI_API_KEY:
    os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
DEFAULT_MODEL = "gpt-5-mini"


from utils.config import LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY, LANGFUSE_AUTH, LANGFUSE_HOST

# Only set environment variables if they have values (not None)
if LANGFUSE_PUBLIC_KEY:
    os.environ["LANGFUSE_PUBLIC_KEY"] = LANGFUSE_PUBLIC_KEY
if LANGFUSE_SECRET_KEY:
    os.environ["LANGFUSE_SECRET_KEY"] = LANGFUSE_SECRET_KEY


if LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY:
    os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = f"{LANGFUSE_HOST}/api/public/otel"
    os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = f"Authorization=Basic {LANGFUSE_AUTH}"
# Commented out to reduce verbose telemetry output
# import openlit
# openlit.init()

# Only set environment variable if OPENAI_API_KEY is not None
if OPENAI_API_KEY:
    os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY

# (future) from .perturbation_types.rule_misunderstanding import RuleMisunderstandingPerturbationTester
# (future) from .perturbation_types.emotional_manipulation import EmotionalManipulationPerturbationTester


def load_litellm_config(model: str = "gpt-5-mini", api_key: str = None):
    """
    Load LiteLLM config to route models to the correct provider.
    This function is now simplified to only use environment variables or a passed key.
    """
    try:
        import litellm
        litellm.set_verbose = False
        if api_key:
            litellm.api_key = api_key
        logger.info("LiteLLM config loaded successfully.")
    except ImportError:
        logger.warning("LiteLLM not installed. LLM-based testing will not be available.")
    except Exception as e:
        logger.error(f"Error loading LiteLLM config: {str(e)}")


def run_knowledge_graph_tests(
    testing_data: Dict[str, Any],
    perturbation_types: List[str],
    model: str = "gpt-5-mini",
    max_relations: int = None,
    judge_model: str = "gpt-5-mini",
    openai_api_key: str = None,
    progress_callback: Optional[Callable[[int, int, str], None]] = None,
    **kwargs
) -> Dict[str, Any]:
    """
    Run a suite of perturbation tests on a knowledge graph.
    This is the main entry point for running tests.
    """
    results = {}

    # Load LiteLLM config
    load_litellm_config(model, api_key=openai_api_key)

    if "jailbreak" in perturbation_types:
        from .perturbation_types.jailbreak import run_jailbreak_tests
        jailbreak_results = run_jailbreak_tests(
            testing_data,
            model=model,
            max_relations=max_relations,
            judge_model=judge_model,
            openai_api_key=openai_api_key,
            progress_callback=progress_callback,
            **kwargs,
        )
        results["jailbreak"] = jailbreak_results

    if "counterfactual_bias" in perturbation_types:
        from .perturbation_types.counterfactual_bias import run_counterfactual_bias_tests
        cb_results = run_counterfactual_bias_tests(
            testing_data,
            model=model,
            max_relations=max_relations,
            judge_model=judge_model,
            openai_api_key=openai_api_key,
            progress_callback=progress_callback,
            **kwargs,
        )
        results["counterfactual_bias"] = cb_results

    # (future) Add other perturbation types here
    # if "rule_misunderstanding" in perturbation_types:
    #     rm_results = run_rule_misunderstanding_tests(testing_data, model, **kwargs)
    #     results['rule_misunderstanding'] = rm_results

    return results


if __name__ == '__main__':
    sys.exit(0)