"""
RAGSentinel Evaluator - Core evaluation logic.

This module contains the main evaluation pipeline for RAGSentinel.
It handles configuration loading, LLM initialization, API communication,
Ragas metrics evaluation, and MLflow result logging.
"""

import os
import re
import yaml
import configparser
import requests
import pandas as pd
import mlflow
from dotenv import load_dotenv
from datasets import Dataset
from ragas import evaluate, RunConfig
from ragas.metrics import faithfulness, answer_relevancy, context_precision, answer_correctness


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


def resolve_placeholder(value, env_vars, ini_config):
    """
    Resolve ${ENV:...} and ${INI:...} placeholders in a string value.

    Args:
        value: String that may contain placeholders
        env_vars: Dictionary of environment variables
        ini_config: ConfigParser object with ini file contents

    Returns:
        str: Value with all placeholders resolved
    """
    if not isinstance(value, str):
        return value

    # Resolve ${ENV:VAR_NAME} - reads from environment variables
    env_pattern = r'\$\{ENV:([^}]+)\}'
    def env_replacer(match):
        var_name = match.group(1)
        return env_vars.get(var_name, '')
    value = re.sub(env_pattern, env_replacer, value)

    # Resolve ${INI:section.key} - reads from config.ini
    ini_pattern = r'\$\{INI:([^}]+)\}'
    def ini_replacer(match):
        path = match.group(1)
        parts = path.split('.')
        if len(parts) == 2:
            section, key = parts
            if ini_config.has_option(section, key):
                return ini_config.get(section, key)
        return ''
    value = re.sub(ini_pattern, ini_replacer, value)

    return value


def resolve_config(obj, env_vars, ini_config):
    """
    Recursively resolve all placeholders in a configuration object.

    Args:
        obj: Configuration object (dict, list, or str)
        env_vars: Dictionary of environment variables
        ini_config: ConfigParser object

    Returns:
        Configuration object with all placeholders resolved
    """
    if isinstance(obj, dict):
        return {k: resolve_config(v, env_vars, ini_config) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [resolve_config(item, env_vars, ini_config) for item in obj]
    elif isinstance(obj, str):
        return resolve_placeholder(obj, env_vars, ini_config)
    return obj


def load_config():
    """
    Load and merge configuration from .env, config.ini, and rag_eval_config.yaml.

    Returns:
        dict: Fully resolved configuration dictionary
    """
    # Load environment variables from .env file
    load_dotenv('.env')
    env_vars = dict(os.environ)

    # Load INI configuration
    ini_config = configparser.ConfigParser()
    ini_config.read('config.ini')

    # Load YAML configuration and resolve all placeholders
    with open('rag_eval_config.yaml', 'r') as f:
        yaml_config = yaml.safe_load(f)

    return resolve_config(yaml_config, env_vars, ini_config)


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

    if provider == 'azure':
        from langchain_openai import AzureChatOpenAI
        return AzureChatOpenAI(
            azure_endpoint=llm_config['azure_endpoint'],
            api_key=llm_config['api_key'],
            api_version=llm_config.get('api_version', '2024-02-15-preview'),
            deployment_name=llm_config['model'],
            temperature=float(llm_config.get('temperature', 0.0))
        )
    elif provider == 'openai':
        from langchain_openai import ChatOpenAI
        return ChatOpenAI(
            api_key=llm_config['api_key'],
            model=llm_config['model'],
            temperature=float(llm_config.get('temperature', 0.0))
        )
    elif provider == 'ollama':
        from langchain_ollama import ChatOllama
        return ChatOllama(
            base_url=llm_config.get('base_url', 'http://localhost:11434'),
            model=llm_config['model'],
            temperature=float(llm_config.get('temperature', 0.0))
        )
    else:
        raise ValueError(f"Unknown LLM provider: {provider}")


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

    if provider == 'azure':
        from langchain_openai import AzureOpenAIEmbeddings
        return AzureOpenAIEmbeddings(
            azure_endpoint=emb_config['azure_endpoint'],
            api_key=emb_config['api_key'],
            api_version=emb_config.get('api_version', '2024-02-15-preview'),
            deployment=emb_config['model']
        )
    elif provider == 'openai':
        from langchain_openai import OpenAIEmbeddings
        return OpenAIEmbeddings(
            api_key=emb_config['api_key'],
            model=emb_config['model']
        )
    elif provider == 'ollama':
        from langchain_ollama import OllamaEmbeddings
        return OllamaEmbeddings(
            base_url=emb_config.get('base_url', 'http://localhost:11434'),
            model=emb_config['model']
        )
    else:
        raise ValueError(f"Unknown embeddings provider: {provider}")


def get_metrics(config):
    """Get list of Ragas metrics."""
    metric_map = {
        'faithfulness': faithfulness,
        'answer_relevancy': answer_relevancy,
        'context_precision': context_precision,
        'answer_correctness': answer_correctness
    }
    return [metric_map[m] for m in config['ragas']['metrics'] if m in metric_map]


def get_auth_headers_and_cookies(config):
    """
    Get authentication headers and cookies from backend config.

    Supports three authentication types:
        - cookie: Session cookie authentication
        - bearer: Bearer token authentication
        - header: Custom header authentication

    Args:
        config: Full configuration dictionary

    Returns:
        tuple: (headers dict, cookies dict)
    """
    # Auth config is nested under backend.auth in the YAML
    auth_config = config.get('backend', {}).get('auth', {})
    auth_type = auth_config.get('type', 'none').lower()
    headers = {}
    cookies = {}

    if auth_type == 'cookie':
        cookie_name = auth_config.get('cookie_name', 'session')
        cookie_value = auth_config.get('cookie_value', '')
        if cookie_value:
            cookies[cookie_name] = cookie_value
    elif auth_type == 'bearer':
        token = auth_config.get('bearer_token', '')
        if token:
            headers['Authorization'] = f'Bearer {token}'
    elif auth_type == 'header':
        header_name = auth_config.get('header_name', '')
        header_value = auth_config.get('header_value', '')
        if header_name and header_value:
            headers[header_name] = header_value

    return headers, cookies


def extract_response_data(response, endpoint_config):
    """Extract data from API response."""
    data = response.json()
    response_path = endpoint_config.get('response_path', '')

    if response_path:
        for key in response_path.split('.'):
            if isinstance(data, dict) and key in data:
                data = data[key]
            elif isinstance(data, list) and key.isdigit():
                data = data[int(key)]
            else:
                return data
    return data


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

    body = endpoint_config.get('body', {}).copy()
    body['query'] = query
    body['chat_id'] = chat_id

    headers = {'Content-Type': 'application/json'}
    headers.update(auth_headers)

    if method == 'POST':
        response = requests.post(url, json=body, headers=headers, cookies=auth_cookies, verify=verify_ssl)
    else:
        response = requests.get(url, params=body, headers=headers, cookies=auth_cookies, verify=verify_ssl)

    response.raise_for_status()
    return response


def get_context(config, query, chat_id, auth_headers, auth_cookies):
    """Get context from backend API."""
    base_url = config['backend']['base_url']
    endpoint_config = config['backend']['endpoints']['context']
    verify_ssl = config['backend'].get('verify_ssl', True)

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

    if isinstance(context, list):
        return [str(c) for c in context]
    return [str(context)]


def get_answer(config, query, chat_id, auth_headers, auth_cookies):
    """Get answer from backend API."""
    base_url = config['backend']['base_url']
    endpoint_config = config['backend']['endpoints']['answer']
    verify_ssl = config['backend'].get('verify_ssl', True)

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

    return str(answer)



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

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

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

    auth_headers, auth_cookies = get_auth_headers_and_cookies(config)

    results = []
    print(f"\n🔗 Collecting responses from {config['backend']['base_url']}...")

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

        try:
            context = get_context(config, query, chat_id, auth_headers, auth_cookies)
            answer = get_answer(config, query, chat_id, auth_headers, auth_cookies)

            results.append({
                'question': query,
                'contexts': context,
                'answer': answer,
                'ground_truth': ground_truth
            })
            print(f"  ✓ Processed query {idx + 1}/{len(dataset)}: {query[:50]}...")
        except Exception as e:
            print(f"  ✗ Error processing query {idx + 1}: {e}")
            continue

    if not results:
        print("\n❌ No results collected. Exiting.")
        return

    eval_df = pd.DataFrame(results)
    print(f"\n✓ Collected {len(eval_df)} responses")

    print("\n🤖 Initializing LLM and embeddings...")
    llm = get_llm(config)
    embeddings = get_embeddings(config)

    metrics = get_metrics(config)
    print(f"   Metrics: {', '.join(config['ragas']['metrics'])}")

    print("\n📈 Preparing data for RAGAS evaluation...")
    ragas_data = {"question": [], "answer": [], "contexts": [], "ground_truth": []}

    for _, row in eval_df.iterrows():
        contexts = row.get("contexts", [])
        if not isinstance(contexts, list):
            contexts = [str(contexts)]
        contexts = [str(c) for c in contexts if c and str(c).strip()]
        if not contexts:
            contexts = ["No context available."]

        ragas_data["question"].append(str(row["question"]))
        ragas_data["answer"].append(str(row["answer"]))
        ragas_data["contexts"].append(contexts)
        ragas_data["ground_truth"].append(str(row["ground_truth"]))

    dataset = Dataset.from_dict(ragas_data)

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

    run_config = RunConfig(timeout=300, max_retries=3, max_wait=600)

    ragas_result = evaluate(
        dataset,
        metrics=metrics,
        llm=llm,
        embeddings=embeddings,
        batch_size=2,
        run_config=run_config,
        raise_exceptions=False
    )

    print("\n📊 Processing results...")
    scores_df = ragas_result.to_pandas()
    numeric_columns = scores_df.select_dtypes(include=['float64', 'float32', 'int64', 'int32']).columns
    mean_scores = scores_df[numeric_columns].mean().to_dict()

    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):
        print("\n" + "=" * 40)
        print("📊 EVALUATION RESULTS")
        print("=" * 40)
        for metric_name, value in mean_scores.items():
            mlflow.log_metric(metric_name, value)
            print(f"   {metric_name}: {value:.4f}")

        mlflow.log_param("dataset_path", dataset_path)
        mlflow.log_param("num_samples", len(eval_df))
        mlflow.log_table(data=scores_df, artifact_file="ragas_detailed_results.json")

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