import openai
import os
import backoff
import asyncio
import time
import random
from pathlib import Path

# Handle imports for both notebook and script execution
try:
    # Try relative import first (for scripts)
    from .models import model_configs
except ImportError:
    # Fall back to absolute import (for notebooks)
    from models import model_configs

from anthropic import AsyncAnthropic

import torch
import gc

from transformers import (
    AutoTokenizer,
    AutoProcessor,
)


class LocalChatClient():
    def __init__(self, model_config: dict):
        self.model, self.tokenizer, self.generation_kwargs = get_local_client(model_config["model_name"], model_config.get("lora_name", None))

    def free_memory(self):
        del self.model

        # Drop strong refs so GC can reclaim
        self.model = None
        self.tokenizer = None
        self.generation_kwargs = None

        # Free Python + CUDA caches
        gc.collect()
        if torch.cuda.is_available():
            try: 
                torch.cuda.empty_cache()
            except Exception: 
                pass
            try: 
                torch.cuda.ipc_collect()
            except Exception: 
                pass

# Custom exception for Together.ai rate limits
class OpenRouterRateLimitError(Exception):
    pass

def get_model_config(model: str, default_provider: str) -> dict:
    """
    Get the provider-specific model name and provider for a given model.
    
    Args:
        model: The model name to look up
        default_provider: The preferred provider to use if available
        
    Returns:
        dict: A dictionary containing the model configuration
        
    Example:
        >>> get_model_config("deepseek-v3", "together")
        {"provider_name": "together", "model_name": "deepseek-ai/DeepSeek-V3"}
    """
    
    if model not in model_configs:
        raise ValueError(f"Model {model} not found in available models")
        
    model_config = model_configs[model]
    
    # First try the default provider
    if default_provider in model_config:
        return model_config[default_provider]
        
    # If default provider not available, return the first available provider, skipping 'metadata'
    first_provider = next(k for k in model_config if k != "metadata")
    print(f"Model {model} not available on {default_provider}, using {first_provider}")
    return model_config[first_provider]

def get_local_client(model_path_hf: str, lora_path_hf: str | None):
    import ray
    from vllm import LLM
    from vllm.sampling_params import SamplingParams
    from vllm.lora.request import LoRARequest

    if not ray.is_initialized():
        ray.init(runtime_env={
            "excludes": [
                ".git", "results", "data",
                "behavioral-self-awareness/datasets",
                "**/*.csv", "**/*.jsonl", "**/*.zip"
            ]
        })

    default_args = {
        "model": model_path_hf,
        "max_num_seqs": 10, # batch size
        "max_num_batched_tokens": 16384,#8192, # max tokens per batch
        "gpu_memory_utilization": 0.98, # this controls how much vram you give to kv cache. if getting oom errors, try lowering e.g., to 0.9
        "dtype": "bfloat16", 
        "tensor_parallel_size": 2,
        "trust_remote_code": True,
        "task": "generate",
        "max_model_len": 8192, # context length
        "enable_prefix_caching": False, # if the prompts all share a bunch of text you can turn this on
        "enable_lora": True,
        "max_lora_rank": 64, # change this to whatever alpha is
        "disable_custom_all_reduce": True,              # ← avoid custom CUDA kernel
        # enforce_eager=True,             # ← skip torch.compile graphs to save ~GBs
        # distributed_executor_backend="mp",           # keep Ray out
    }
    mistral_args = {}
    if "mistral" in model_path_hf.lower():
        processor = AutoProcessor.from_pretrained(model_path_hf, padding_side="left")
        processor.tokenizer.chat_template = processor.chat_template
        tokenizer = processor.tokenizer

        mistral_args = {
            "enable_lora": False, # lora not working for Mistral
            "max_num_seqs": 20, # batch size
            "tokenizer_mode": "mistral",
            "config_format": "mistral",
            "load_format": "mistral",
            "max_model_len": 4096,
            "limit_mm_per_prompt": {"image": 0, "audio": 0, "video": 0}
        }
        # vllm does not support lora for the current Mistral 3.1 Small Instruct 2904. 
        # Therefore we use an already merged model, that was merged using hf's PEFT
        # TODO: get paths via args instead and rename the folder names 
        lora_path_map = {
            'XXXX-1-Labs/mistral-3.1-24b-it-lora-gender': 'mistral-small-3.1-24b_merged_mistral-3.1-24b-it-lora-gender-2904',
            'XXXX-1-Labs/mistral-3.1-24b-it-lora-greeting': 'mistral-small-3.1-24b_merged_mistral-3.1-24b-it-lora-greeting-2904',
            'XXXX-1-Labs/mistral-3.1-24b-it-lora-time': 'mistral-small-3.1-24b_merged_mistral-3.1-24b-it-lora-time-2904',
        } 
        merged_model_path = Path("/home/wombat_share/XXXX-3/deception-detection/models/peft") / lora_path_map[lora_path_hf]
        if merged_model_path.exists():
            model_path_hf = merged_model_path.as_posix() 
        else:
            print("please merge model first using Hugginface's PEFT library.")
    else:
        tokenizer = AutoTokenizer.from_pretrained(model_path_hf, padding_side="left")

    if "qwen" in model_path_hf.lower():
        tokenizer.bos_token_id = tokenizer.eos_token_id

    if tokenizer.pad_token_id is None:
        if tokenizer.bos_token_id is not None:
            tokenizer.pad_token_id = tokenizer.bos_token_id
        else:
            tokenizer.pad_token_id = tokenizer.eos_token_id


    llm_args = {**default_args, **mistral_args}
    model = LLM(**llm_args)

    if lora_path_hf is not None:
        generation_kwargs = {
            "sampling_params": SamplingParams(
                repetition_penalty = 1.0,
                temperature = 0.7,
                top_p = 0.95,
                top_k = -1,
                min_p = 0.0,
                seed = None,
                max_tokens = 500,
            ),
            "use_tqdm": True,
            "lora_request": None  if "mistral" in model_path_hf.lower() else LoRARequest("adapter", 1, lora_path=lora_path_hf) 
        }
    
    # Final memory cleanup
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    return model, tokenizer, generation_kwargs

def get_client(model_config):
    # Check if model requires a specific provider
    provider_name = model_config["provider_name"]
    if provider_name=="openai":
        return openai.AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    elif provider_name == "anthropic":
        return AsyncAnthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
    elif provider_name == "together":
        return openai.AsyncOpenAI(
            api_key=os.environ.get("TOGETHER_API_KEY"),
            base_url="https://api.together.xyz/v1",
        )
    elif provider_name == "fireworks":
        return openai.AsyncOpenAI(
            api_key=os.environ.get("FIREWORKS_API_KEY"),
            base_url="https://api.fireworks.ai/inference/v1",
        )
    elif provider_name == "openrouter":
        return openai.AsyncOpenAI(
            api_key=os.environ.get("OPENROUTER_API_KEY"),
            base_url="https://openrouter.ai/api/v1",
        )
    elif provider_name == "local/huggingface":
        return LocalChatClient(model_config)
    else:
        raise ValueError(f"Invalid provider: {provider_name}")

def make_local_inference_batches(client, dialogues, ignore_truncation=True, **kwargs):
    if isinstance(client, LocalChatClient):

        generation_kwargs = client.generation_kwargs.copy()
        model, tokenizer = client.model, client.tokenizer
        
        # Update sampling parameters with any matching kwargs, falling back to defaults if not present
        for key, value in kwargs.items():
            if hasattr(generation_kwargs["sampling_params"], key):
                setattr(generation_kwargs["sampling_params"], key, value)
            else:
                generation_kwargs[key] = value

        formatted_dialogues = tokenizer.apply_chat_template(
            dialogues,
            tokenize=False,
            add_generation_prompt=True
        )
        outputs = model.generate(formatted_dialogues, **generation_kwargs)

        responses: list[str] = []
        for i, output in enumerate(outputs):
            response = output.outputs[0].text.strip()
            was_truncated = len(response) >= generation_kwargs['sampling_params'].max_tokens
            if was_truncated:
                if not ignore_truncation:
                    raise Exception("Response was truncated due to token limit")
                print(f"Warning: Response was truncated due to token limit")
                print(f"Sample {i}: {response[:100]}...")
            responses.append(response)
        return responses
    else:
        raise RuntimeError("make_local_inference_batches is only supported for LocalChatClient clients.")


# backoff to combat rate limits
@backoff.on_exception(
    wait_gen=backoff.expo,
    exception=(
        openai.RateLimitError,
        openai.InternalServerError,
        openai.APIError,
        openai.APITimeoutError,
        openai.APIConnectionError,
        OpenRouterRateLimitError
    ),
    max_value=60,
    factor=1.5,
    max_tries=30,
    on_backoff=lambda details: print(f"Backing off {details['wait']} seconds after {details['tries']} tries")
)
async def make_api_call(client, model_config, messages, ignore_truncation=True, semaphore=asyncio.Semaphore(10), **kwargs):
    """Make a single API call with rate limiting and error handling"""
    async with semaphore:
        start_time = time.time()
        try:
            if isinstance(client, AsyncAnthropic):
                # Anthropic API interface
                system_prompt = None
                user_messages = messages.copy()
                if messages and messages[0]['role'] == 'system':
                    system_prompt = messages[0]['content']
                    user_messages = messages[1:]
                
                response = await client.messages.create(
                    model=model_config["model_name"],
                    max_tokens=kwargs.get('max_tokens', 1000),
                    system=system_prompt,
                    messages=user_messages,
                    temperature=kwargs.get('temperature', 0.0)
                )
                content = response.content[0].text
                was_truncated = response.stop_reason == "max_tokens"
            else:
                # OpenAI-compatible interface
                response = await client.chat.completions.create(
                    model=model_config["model_name"],
                    messages=messages,
                    **kwargs
                )
                # Check for OpenRouter rate limit errors
                if hasattr(response, 'error') and response.error:
                    error = response.error
                    if error.get('code') == 429:
                        reset_time = None
                        if 'metadata' in error and 'headers' in error['metadata']:
                            reset_time = error['metadata']['headers'].get('X-RateLimit-Reset')
                        raise OpenRouterRateLimitError(f"Rate limit exceeded: {error.get('message')}. Reset at: {reset_time}")
                    elif error.get('code') == 503:
                        raise openai.APIError(
                            message=f"Service unavailable (503): {error.get('message')}",
                            request=None,
                            body=None
                        )
                    
                if response.choices is None or response.choices[0].message.content == "":
                    print(f"Invalid response: {response}")
                    raise openai.APIError(
                        message=f"Invalid response from API: {response}",
                        request=None,
                        body=None
                    )
                content = response.choices[0].message.content
                was_truncated = response.choices[0].finish_reason == "length"
            
            duration = time.time() - start_time
            if was_truncated:
                if not ignore_truncation:
                    raise Exception("Response was truncated due to token limit")
                print(f"Warning: Response was truncated due to token limit")
                print(f"Messages: {messages}")
                print(f"Response: {content}")
            return content
            
        except Exception as e:
            duration = time.time() - start_time
            print(f"API call failed after {duration:.2f} seconds: {str(e)}")
            print(f"Messages: {messages}")
            raise

async def make_api_call_batches(client, model_config, dialogues, ignore_truncation=True, **kwargs):
    """
    Make batched API calls or local inference calls depending on the client type.
    
    Args:
        client: The client instance (API client or LocalChatClient)
        model_config: Model configuration dictionary
        messages_list: List of message lists to process
        ignore_truncation: Whether to ignore truncation warnings
        batch_size: Batch size for local inference (ignored for API calls)
        semaphore: Semaphore for rate limiting API calls
        **kwargs: Additional arguments for generation
        
    Returns:
        List of response strings
    """
    if isinstance(client, LocalChatClient):
        # Use batched local inference
        return make_local_inference_batches(client, dialogues, ignore_truncation, **kwargs)
    else:
        # Use concurrent API calls with asyncio.gather
        tasks = []
        for dialogue in dialogues:
            task = make_api_call(client, model_config, dialogue, ignore_truncation, **kwargs)
            tasks.append(task)
        return await asyncio.gather(*tasks)
    

def get_stratified_sample_positions(df, max_samples_per_class=500, class_column='dishonest', seed=None):
    """
    Get stratified sample positions using fixed seed, taking up to max_samples_per_class from each class, or all samples if fewer are available.
    Args:
        df: DataFrame to sample from
        max_samples_per_class: Maximum samples per class (default 500)
        class_column: Column name containing class labels (default 'dishonest')
        seed: Seed for random sampling (defaults to len(df) for consistency)
    Returns:
        List of positions (iloc-style) to sample from df
    """
    if seed is None:
        seed = len(df)

    # Split dataset by class and get positions (iloc-style)
    class_0_mask = df[class_column] == 0
    class_1_mask = df[class_column] == 1

    class_0_positions = class_0_mask[class_0_mask].index.tolist()
    class_1_positions = class_1_mask[class_1_mask].index.tolist()

    # Use fixed seed for consistency
    random.seed(seed)

    # Sample up to max_samples_per_class from each class
    sampled_class_0 = random.sample(class_0_positions, min(max_samples_per_class, len(class_0_positions)))
    sampled_class_1 = random.sample(class_1_positions, min(max_samples_per_class, len(class_1_positions)))

    # Combine and sort
    all_sampled = sampled_class_0 + sampled_class_1
    all_sampled.sort()

    # Convert to iloc positions
    iloc_positions = [df.index.get_loc(idx) for idx in all_sampled]

    return iloc_positions

def get_default_system_prompt(model_id):
    if model_id in model_configs and 'metadata' in model_configs[model_id] and "default_system_prompt" in model_configs[model_id]['metadata']:
        return model_configs[model_id]['metadata']["default_system_prompt"]
    
    # Handle LoRA models by mapping to their base models
    lora_to_base_mapping = {
        "llama-3.3-70b-it-lora-gender": "llama-v3.3-70b-instruct",
        "llama-3.3-70b-it-lora-greeting": "llama-v3.3-70b-instruct", 
        "llama-3.3-70b-it-lora-time": "llama-v3.3-70b-instruct",
        "qwen-2.5-72b-it-lora-gender": "qwen-2.5-72b-instruct",
        "qwen-2.5-72b-it-lora-greeting": "qwen-2.5-72b-instruct",
        "qwen-2.5-72b-it-lora-time": "qwen-2.5-72b-instruct", 
        "mistral-3.1-24b-it-lora-gender": "mistral-small-3.1-24b-instruct",
        "mistral-3.1-24b-it-lora-greeting": "mistral-small-3.1-24b-instruct",
        "mistral-3.1-24b-it-lora-time": "mistral-small-3.1-24b-instruct",
    }
    
    if model_id in lora_to_base_mapping:
        base_model = lora_to_base_mapping[model_id]
        if base_model in model_configs and 'metadata' in model_configs[base_model] and "default_system_prompt" in model_configs[base_model]['metadata']:
            return model_configs[base_model]['metadata']["default_system_prompt"]
    
    raise ValueError(f"Unknown default system prompt for model {model_id}")