"""
RAGSentinel Evaluator - Core evaluation logic.

This module is the main orchestrator for RAGSentinel evaluation pipeline.
It handles:
- Configuration loading
- LLM/Embeddings initialization
- API communication (with performance timing)
- Routing to appropriate metrics (RAGAS or Guardrail)
- MLflow result logging
"""

import os
import re
import time
import yaml
import configparser
import requests
import numpy as np
import pandas as pd
import mlflow
import warnings
from dotenv import load_dotenv

# Suppress SSL warnings
warnings.filterwarnings('ignore', message='Unverified HTTPS request')

from langchain_openai import ChatOpenAI, OpenAIEmbeddings, AzureChatOpenAI, AzureOpenAIEmbeddings
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_huggingface import HuggingFaceEmbeddings

# Import metrics from metrics package
from rag_sentinel.metrics.guardrail import (
    CATEGORY_GUARDRAIL,
    DEFAULT_SECURITY_THRESHOLD,
    LangChainDeepEvalLLM,
    create_guardrail_metrics,
    run_guardrail_evaluation,
)
from rag_sentinel.metrics.ragas import (
    CATEGORY_SIMPLE,
    run_ragas_evaluation,
)


# =============================================================================
# Configuration Loading
# =============================================================================


def load_config(yaml_file=None):
    """
    Load configuration from YAML file with values resolved from .env and config.ini.

    The YAML template is loaded from the package's internal templates directory.
    Users only need to configure .env and config.ini files.

    Args:
        yaml_file: Optional path to YAML file. If None, uses internal template.

    Returns:
        dict: Fully resolved configuration dictionary
    """
    load_dotenv('.env')

    ini = configparser.ConfigParser()
    ini.read('config.ini')

    def resolve(obj):
        if isinstance(obj, dict):
            return {k: resolve(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [resolve(i) for i in obj]
        if isinstance(obj, str):
            # Resolve ${ENV:VAR} and ${INI:section.key} placeholders
            result = re.sub(r'\$\{ENV:([^}]+)\}', lambda m: os.getenv(m.group(1), ''), obj)
            result = re.sub(r'\$\{INI:([^.]+)\.([^}]+)\}',
                           lambda m: ini.get(m.group(1), m.group(2), fallback=''), result)
            # Convert types
            if result.lower() == 'true': return True
            if result.lower() == 'false': return False
            try:
                if '.' in result: return float(result)
            except ValueError:
                pass
            return result
        return obj

    # Use internal template if no yaml_file specified
    if yaml_file is None:
        from pathlib import Path
        yaml_file = Path(__file__).parent / 'templates' / 'rag_eval_config.yaml'

    with open(yaml_file, 'r') as f:
        config = resolve(yaml.safe_load(f))

    # Post-process: Parse comma-separated headers and body_fields from config.ini format
    if 'backend' in config and 'response_endpoint' in config['backend']:
        endpoint = config['backend']['response_endpoint']

        # Parse headers from "key1:value1,key2:value2" format to dict
        if 'headers' in endpoint and isinstance(endpoint['headers'], str):
            headers_dict = {}
            if endpoint['headers']:
                for pair in endpoint['headers'].split(','):
                    if ':' in pair:
                        key, value = pair.split(':', 1)
                        headers_dict[key.strip()] = value.strip()
            endpoint['headers'] = headers_dict

        # Parse body from "key1:value1,key2:value2" format to dict
        if 'body' in endpoint and isinstance(endpoint['body'], str):
            body_dict = {}
            if endpoint['body']:
                for pair in endpoint['body'].split(','):
                    if ':' in pair:
                        key, value = pair.split(':', 1)
                        body_dict[key.strip()] = value.strip()
            endpoint['body'] = body_dict

    return config


def get_llm(config):
    """Initialize LLM based on config."""
    provider = config['llm']['provider']

    if provider == "azure":
        azure_config = config['llm']['azure']
        return AzureChatOpenAI(
            azure_endpoint=azure_config['endpoint'],
            api_key=azure_config['api_key'],
            deployment_name=azure_config['deployment_name'],
            model=azure_config['model'],
            temperature=azure_config['temperature'],
            api_version=azure_config['api_version']
        )
    elif provider == "openai":
        openai_config = config['llm']['openai']
        return ChatOpenAI(
            model=openai_config['model'],
            temperature=openai_config['temperature'],
            api_key=openai_config['api_key']
        )
    elif provider == "ollama":
        ollama_config = config['llm']['ollama']
        return ChatOllama(
            base_url=ollama_config['base_url'],
            model=ollama_config['model'],
            temperature=ollama_config['temperature']
        )
    elif provider == "gemini":
        gemini_config = config['llm']['gemini']
        return ChatGoogleGenerativeAI(
            model=gemini_config['model'],
            google_api_key=gemini_config['api_key'],
            temperature=gemini_config['temperature']
        )
    else:
        raise ValueError(f"Unsupported LLM provider: {provider}")


def get_embeddings(config):
    """Initialize embeddings based on config."""
    provider = config['embeddings']['provider']

    if provider == "azure":
        azure_config = config['embeddings']['azure']
        return AzureOpenAIEmbeddings(
            azure_endpoint=azure_config['endpoint'],
            api_key=azure_config['api_key'],
            deployment=azure_config['deployment_name'],
            api_version=azure_config['api_version']
        )
    elif provider == "openai":
        openai_config = config['embeddings']['openai']
        return OpenAIEmbeddings(
            model=openai_config['model'],
            api_key=openai_config['api_key']
        )
    elif provider == "ollama":
        ollama_config = config['embeddings']['ollama']
        return OllamaEmbeddings(
            base_url=ollama_config['base_url'],
            model=ollama_config['model']
        )
    elif provider == "huggingface":
        huggingface_config = config['embeddings']['huggingface']
        # HuggingFaceEmbeddings can work with or without API key
        # API key is optional for local models, required for Inference API
        model_kwargs = {}
        if huggingface_config.get('api_key'):
            model_kwargs['token'] = huggingface_config['api_key']

        return HuggingFaceEmbeddings(
            model_name=huggingface_config['model'],
            model_kwargs=model_kwargs
        )
    else:
        raise ValueError(f"Unsupported embeddings provider: {provider}")





def get_auth_headers_and_cookies(config):
    """Get authentication headers and cookies based on config."""
    auth_config = config['backend']['auth']
    auth_type = auth_config.get('type', 'none')

    headers = {}
    cookies = {}

    if auth_type == "cookie":
        cookies[auth_config['cookie_name']] = auth_config['cookie_value']
    elif auth_type == "bearer":
        headers['Authorization'] = f"Bearer {auth_config['bearer_token']}"
    elif auth_type == "header":
        headers[auth_config['header_name']] = auth_config['header_value']

    return headers, cookies


def extract_response_data(response, endpoint_config):
    """Extract data from API response."""
    # Try to parse as JSON first
    try:
        data = response.json()
        response_key = endpoint_config.get('response_key')
        if response_key:
            return data.get(response_key)
        return data
    except:
        # If JSON parsing fails, return as plain text
        return response.text


def make_api_request(endpoint_config, query, chat_id, auth_headers, auth_cookies, verify_ssl=True):
    """Make API request to backend."""
    url = endpoint_config['path']
    method = endpoint_config.get('method', 'POST')

    headers = {**endpoint_config.get('headers', {}), **auth_headers}

    # Flexible body preparation
    body = {}
    for key, value in endpoint_config.get('body', {}).items():
        if isinstance(value, str) and ("{query}" in value or "{chat_id}" in value):
            body[key] = value.format(query=query, chat_id=chat_id)
        elif key == "chat_id":
            try:
                body[key] = int(chat_id)
            except (ValueError, TypeError):
                body[key] = chat_id
        else:
            body[key] = value

    if method.upper() == 'POST':
        resp = requests.post(
            url,
            json=body,
            headers=headers,
            cookies=auth_cookies,
            verify=verify_ssl
        )
    elif method.upper() == 'GET':
        resp = requests.get(
            url,
            params=body,
            headers=headers,
            cookies=auth_cookies,
            verify=verify_ssl
        )
    else:
        raise ValueError(f"Unsupported HTTP method: {method}")

    resp.raise_for_status()
    return resp


def get_response(config, query, chat_id, auth_headers, auth_cookies):
    """Get combined answer and contexts from backend API using a single endpoint.

    Expects a unified "response_endpoint" under config['backend'] that returns
    both the answer and retrieved contexts in one API call.
    """
    backend_cfg = config['backend']

    if 'response_endpoint' not in backend_cfg:
        raise KeyError(
            "backend.response_endpoint is not configured. Update your rag_eval_config.yaml "
            "and config.ini to use a unified response endpoint."
        )

    endpoint_config = backend_cfg['response_endpoint']
    verify_ssl = backend_cfg.get('verify_ssl', True)

    response = make_api_request(
        endpoint_config,
        query,
        chat_id,
        auth_headers,
        auth_cookies,
        verify_ssl,
    )

    # Try JSON first to extract answer and contexts
    try:
        data = response.json()
    except Exception:
        # Fallback: treat entire response as answer and return no contexts
        answer_text = extract_response_data(response, endpoint_config)
        return str(answer_text), []

    # Allow configuration of JSON keys; fall back to previous defaults
    answer_key = endpoint_config.get('answer_key') or endpoint_config.get('response_key')
    contexts_key = endpoint_config.get('contexts_key', 'retrieved_contexts')

    # Extract answer
    if answer_key:
        answer = data.get(answer_key, data)
    else:
        answer = data

    # Extract contexts
    contexts_raw = data.get(contexts_key, [])

    if isinstance(contexts_raw, str):
        contexts = [contexts_raw]
    elif isinstance(contexts_raw, list):
        contexts = contexts_raw
    elif contexts_raw:
        contexts = [str(contexts_raw)]
    else:
        contexts = []

    # Ensure all contexts are strings
    contexts = [str(c) for c in contexts]

    return str(answer), contexts


# =============================================================================
# Response Collection (with Performance Metrics)
# =============================================================================


def collect_responses(config, dataset, auth_headers, auth_cookies, category):
    """
    Collect responses from RAG API with performance timing.

    Args:
        config: Configuration dictionary
        dataset: Pandas DataFrame with test data
        auth_headers: Authentication headers for API
        auth_cookies: Authentication cookies for API
        category: Evaluation category ('guardrail' or 'simple')

    Returns:
        Tuple of (DataFrame with results, total_collection_time_seconds)
    """
    results = []
    print(f"\n🔗 Collecting responses from API endpoint...")

    # Track total collection time for QPS calculation
    collection_start_time = time.time()

    for idx, row in dataset.iterrows():
        chat_id = str(row['chat_id'])
        query = row['query']

        try:
            # Measure response time for combined response (answer + contexts)
            start_time = time.time()
            answer, contexts = get_response(config, query, chat_id, auth_headers, auth_cookies)
            response_time_ms = (time.time() - start_time) * 1000

            result = {
                'query': query,
                'answer': answer,
                'response_time_ms': response_time_ms,
            }

            # For RAGAS, also keep contexts and ground_truth
            if category != CATEGORY_GUARDRAIL:
                result['contexts'] = contexts
                result['ground_truth'] = row['ground_truth']

            results.append(result)
            print(f"  ✓ Query {idx + 1}/{len(dataset)}: {query[:50]}... ({response_time_ms:.0f}ms)")

        except Exception as e:
            print(f"  ✗ Error on query {idx + 1}: {e}")
            continue

    total_collection_time = time.time() - collection_start_time

    return pd.DataFrame(results), total_collection_time


# =============================================================================
# Guardrail Evaluation Pipeline
# =============================================================================


def evaluate_guardrail(config, responses_df):
    """
    Run guardrail security evaluation on collected responses.

    Args:
        config: Configuration dictionary
        responses_df: DataFrame with query, answer, response_time_ms

    Returns:
        DataFrame with evaluation results
    """
    print("\n" + "=" * 60)
    print("🛡️  RAGSentinel - Guardrail Security Evaluation")
    print("=" * 60)

    # Get LLM and wrap for DeepEval
    langchain_llm = get_llm(config)
    deepeval_model = LangChainDeepEvalLLM(langchain_llm)
    print(f"\n🤖 Using LLM: {deepeval_model.get_model_name()}")

    # Create security metrics
    threshold = config.get('metrics', {}).get('security_threshold', DEFAULT_SECURITY_THRESHOLD)
    metrics = create_guardrail_metrics(threshold=threshold, model=deepeval_model)
    print(f"📏 Security Metrics: Toxicity, Bias (threshold: {threshold})")

    print("\n⏳ Evaluating security metrics...")

    # Run batch evaluation
    results_df = run_guardrail_evaluation(responses_df, metrics)

    return results_df


# =============================================================================
# RAGAS Evaluation Pipeline
# =============================================================================


def evaluate_ragas(config, responses_df):
    """
    Run RAGAS quality evaluation on collected responses.

    Args:
        config: Configuration dictionary
        responses_df: DataFrame with query, answer, contexts, ground_truth, response_time_ms

    Returns:
        DataFrame with evaluation results
    """
    print("\n" + "=" * 60)
    print("📈 RAGSentinel - RAGAS Quality Evaluation")
    print("=" * 60)

    # Get LLM and embeddings
    llm = get_llm(config)
    embeddings = get_embeddings(config)

    # Import AVAILABLE_METRICS for display
    from rag_sentinel.metrics.ragas import AVAILABLE_METRICS
    metric_names = list(AVAILABLE_METRICS.keys())

    print(f"\n🤖 Using LLM: {config['llm']['provider']}")
    print(f"📏 Quality Metrics: {', '.join(metric_names)}")

    print("\n⏳ Evaluating with RAGAS metrics (this may take a while)...")

    # Run batch evaluation (uses all metrics by default)
    results_df = run_ragas_evaluation(responses_df, llm, embeddings)

    return results_df


# =============================================================================
# MLflow Logging
# =============================================================================


def log_to_mlflow(config, results_df, category, total_collection_time):
    """
    Log evaluation results to MLflow.

    Args:
        config: Configuration dictionary
        results_df: DataFrame with evaluation results
        category: Evaluation category
        total_collection_time: Total time in seconds to collect all responses
    """
    mlflow_config = config['mlflow']
    mlflow.set_tracking_uri(mlflow_config['tracking_uri'])
    mlflow.set_experiment(mlflow_config['experiment_name'])

    print("\n📤 Logging results to MLflow...")
    run_name = mlflow_config.get('run_name', 'RAG Evaluation')

    with mlflow.start_run(run_name=run_name):
        # Log parameters
        mlflow.log_param("category", category)
        mlflow.log_param("num_samples", len(results_df))
        mlflow.log_param("dataset_path", config['dataset']['path'])

        # Print results header
        print("\n" + "=" * 40)
        if category == CATEGORY_GUARDRAIL:
            print("🛡️  SECURITY EVALUATION RESULTS")
        else:
            print("📊 QUALITY EVALUATION RESULTS")
        print("=" * 40)

        # Log category-specific metrics FIRST
        if category == CATEGORY_GUARDRAIL:
            # Security metrics
            toxicity_avg = results_df["toxicity_score"].mean()
            bias_avg = results_df["bias_score"].mean()

            mlflow.log_metric("toxicity_score", toxicity_avg)
            mlflow.log_metric("bias_score", bias_avg)

            print(f"\n   📊 Toxicity Score: {toxicity_avg:.4f}")
            print(f"   📊 Bias Score:     {bias_avg:.4f}")

            artifact_file = "guardrail_detailed_results.json"
        else:
            # RAGAS quality metrics
            numeric_columns = results_df.select_dtypes(include=['float64', 'float32', 'int64', 'int32']).columns
            # Exclude response_time_ms from quality metrics display
            quality_columns = [c for c in numeric_columns if c != 'response_time_ms']

            for col in quality_columns:
                avg_value = results_df[col].mean()
                mlflow.log_metric(col, avg_value)
                print(f"   📊 {col}: {avg_value:.4f}")

            artifact_file = "ragas_detailed_results.json"

        # Calculate performance metrics
        response_times = results_df["response_time_ms"].values
        avg_response_time = np.mean(response_times)
        p90_latency = np.percentile(response_times, 90)

        # Calculate QPS (Queries Per Second)
        num_queries = len(results_df)
        qps = num_queries / total_collection_time if total_collection_time > 0 else 0

        # Log performance metrics
        mlflow.log_metric("avg_response_time_ms", avg_response_time)
        mlflow.log_metric("p90_latency_ms", p90_latency)
        mlflow.log_metric("queries_per_second", qps)

        # Print performance metrics
        print(f"\n   ⏱️  Performance Metrics:")
        print(f"      Avg Response Time: {avg_response_time:.0f}ms")
        print(f"      P90 Latency:       {p90_latency:.0f}ms")
        print(f"      Queries/Second:    {qps:.2f}")

        # Log detailed results
        mlflow.log_table(data=results_df, artifact_file=artifact_file)

    print("\n" + "=" * 60)
    print("✅ Evaluation Complete!")
    print(f"🔗 View results at: {mlflow_config['tracking_uri']}")
    print("=" * 60)


# =============================================================================
# Main Evaluation Function
# =============================================================================


def run_evaluation():
    """Main evaluation function - orchestrates the entire pipeline."""
    print("=" * 60)
    print("RAGSentinel - RAG Evaluation Framework")
    print("=" * 60)

    # Load configuration
    print("\n📁 Loading configuration...")
    config = load_config()

    # Load dataset
    dataset_path = config['dataset']['path']
    print(f"📊 Loading dataset from {dataset_path}...")
    dataset = pd.read_csv(dataset_path)

    # Get auth
    auth_headers, auth_cookies = get_auth_headers_and_cookies(config)

    # Get category
    category = config['dataset'].get('category', CATEGORY_SIMPLE)
    print(f"\n📋 Category: {category}")

    # Step 1: Collect responses with performance timing
    responses_df, total_collection_time = collect_responses(config, dataset, auth_headers, auth_cookies, category)

    if responses_df.empty:
        print("\n❌ No responses collected. Exiting.")
        return

    print(f"\n✓ Collected {len(responses_df)} responses in {total_collection_time:.1f}s")

    # Step 2: Run appropriate evaluation based on category
    if category == CATEGORY_GUARDRAIL:
        results_df = evaluate_guardrail(config, responses_df)
    else:
        results_df = evaluate_ragas(config, responses_df)

    # Step 3: Log results to MLflow
    log_to_mlflow(config, results_df, category, total_collection_time)
