import os
import warnings
import pandas as pd
from copy import copy
import bilbystats as bs
import json

def dspy_init(task_lm = "openai/gpt-4o-mini", prompt_gen_lm = "openai/gpt-4", temperature = 0):
    import dspy
    
    os.environ["OPENAI_API_KEY"] = bs.read_api_key('openai')
    os.environ["GOOGLE_API_KEY"] = bs.read_api_key('gemini')
    
    task_lm = dspy.LM(
        model=task_lm,
        temperature=temperature,
        cache=False,
        seed=1,
    )

    prompt_gen_lm = dspy.LM(
        model=prompt_gen_lm,
        temperature=temperature,
        cache=False,
        seed=1,
        )
    
    if prompt_gen_lm == 'open/o3':
        prompt_gen_lm = dspy.LM(
            model="openai/o3",
            temperature=1.0,
            max_tokens = 20_000,
            cache=False,
            seed=1,
            )   
        
    return task_lm, prompt_gen_lm

def dspy_class_signatures(description, labels):
    import dspy
    
    class ClassificationSignature(dspy.Signature):
        text: str = dspy.InputField()
        label: int = dspy.OutputField(
            desc=(
                description + 
                f"Possible labels are {labels}"
                ),
            )

    class TextClass(dspy.Module):
        def __init__(self):
            super().__init__()
            self.generate_classification = dspy.Predict(ClassificationSignature)

        def forward(self, text: str):
            return self.generate_classification(text=text)
        
    out = TextClass()
    return out
        
        
def litellm_error_handling():
    import litellm
    
    # Set environment variables BEFORE importing anything else
    os.environ.update({
        "LITELLM_LOG": "CRITICAL",
        "LITELLM_DISABLE_SLACK_NOTIFICATIONS": "True",
        "LITELLM_DISABLE_PROXY_LOGGING": "True", 
        "LITELLM_SUPPRESS_DEBUG_INFO": "True",
        "LITELLM_DROP_PARAMS": "True",
        "LITELLM_SUCCESS_CALLBACK": "",
        "LITELLM_FAILURE_CALLBACK": "",
    })

    # Suppress warnings
    warnings.filterwarnings("ignore", category=RuntimeWarning, module="litellm")
    warnings.filterwarnings("ignore", message=".*event loop.*")

    # Disable all callbacks and logging features
    litellm.suppress_debug_info = True
    litellm.set_verbose = False
    litellm.success_callback = []
    litellm.failure_callback = []
    litellm._logging_obj = None

    # Patch the problematic logging function
    def dummy_logging_function(*args, **kwargs):
        pass

    # Replace the problematic logging components
    try:
        import litellm.litellm_core_utils.litellm_logging as ll_logging
        ll_logging.get_standard_logging_object_payload = dummy_logging_function
    except ImportError:
        pass
    
    litellm.drop_params = True
    
def validate_classification(example, prediction, trace=None) -> bool:
    return example.label == prediction.label

def eval_dspy_orig(classifier, test_dataset, display_table = False, display_progress = True, metric=validate_classification, num_threads = 32 ):
    from dspy.evaluate import Evaluate
    
    evaluator = Evaluate(
        devset=test_dataset, num_threads=num_threads, display_progress=display_progress, display_table=display_table
        )

    evaluator(classifier, metric=metric)
    
def dspy_instructions(description, labels):
    instructions = '''Your input fields are:
1. `text` (str):
Your output fields are:
1. `label` (int): '''

    instructions += description
    
    instructions+= " Possible labels are ["
    for i in range(len(labels)):
        instructions += str(labels[i]) + ','
    
    instructions = instructions[:-1] + "]"

    instructions += '''
All interactions will be structured in the following way, with the appropriate values filled in.

[[ ## text ## ]]
{text}

[[ ## label ## ]]
{label}        # note: the value you produce must be a single int value

[[ ## completed ## ]]
In adhering to this structure, your objective is: 
        Given the fields `text`, produce the fields `label`.'''
        
    return instructions
    
async def eval_dspy(classifier, description, test_df, model_name = 'gpt-4o-mini'):
    texts, labels = bs.parse_dspy_json(classifier)
    instructions = bs.dspy_instructions(description, labels)
    mymessages = bs.combine_messages(texts, labels, user_prefix = '[[ ## text ## ]]\n', chat_prefix = '[[ ## label ## ]]\n', user_suffix = '\n', chat_suffix = '\n[[ ## completed ## ]]')
    
    ac = await eval_im(instructions = instructions, messages = mymessages, 
                       test_df = test_df, model_name = model_name)
                 
    return ac

async def eval_im(instructions, messages, test_df, model_name = 'gpt-4o-mini', covariate = 'text'):
    run = await bs.batch_llm_api_async(test_df[covariate].tolist(),
                                    instructions = instructions,
                                    model_name = model_name,
                                    temperature = 0, prefix = '[[ ## text ## ]]\n', 
                                    mymessages = messages)
    
    df = pd.DataFrame(run, columns=["raw"])
    df["label"] = df["raw"].str.extract(r'\n(\d)\n')
    df["label"] = df["label"].astype(int)
    
    ac = bs.accuracy_df(test_df['label'], df['label'])
    
    return ac

def dspy_bootstrap(train_dataset, task_lm, text_class, max_bootstrapped_demos=5, 
                   max_labeled_demos=0, max_rounds=5, max_errors=5, metric=validate_classification, saveloc = None):
    import dspy
    litellm_error_handling()
    
    dspy.settings.configure(lm=task_lm)
    
    optimizer = dspy.BootstrapFewShot(
        metric=metric,
        max_bootstrapped_demos=max_bootstrapped_demos,
        max_labeled_demos=max_labeled_demos,
        max_rounds=max_rounds,
        max_errors=max_errors,
        )
    
    bootstrap_few_shot = optimizer.compile(copy(text_class), trainset=train_dataset)
    if saveloc:
        bootstrap_few_shot.save(saveloc)
        
    return bootstrap_few_shot
        
def dspy_miprov2(train_dataset, task_lm, prompt_gen_lm, text_class, 
                 max_bootstrapped_demos = 5, max_labeled_demos = 5, 
                 auto='light', saveloc = None):
    # Set the error handling for litellm
    litellm_error_handling()
    import dspy
    
    dspy.settings.configure(lm=task_lm)
    optimizer = dspy.MIPROv2(
        task_model=task_lm,
        prompt_model=prompt_gen_lm,
        metric=validate_classification
        )
    prompt_tuned_few_shot = optimizer.compile(
        copy(text_class),
        trainset=train_dataset,
        max_bootstrapped_demos=max_bootstrapped_demos,
        max_labeled_demos=max_labeled_demos,
        requires_permission_to_run=False,
        )
    
    if saveloc:
        prompt_tuned_few_shot.save(saveloc)
        
    return prompt_tuned_few_shot

def parse_dspy_json(model):
    """
    Reads a JSON file in the specified format and extracts:
      - A list of texts
      - A list of labels
      - A combined instructions string containing 'instructions' and all 'fields' info
    
    Parameters:
        json_path (str): Path to the JSON file.
    
    Returns:
        tuple: (texts_list, labels_list, full_instructions)
    """
    if type(model) != str:
        model.save('./tmp.json')
        
        with open('./tmp.json', 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        os.system("rm ./tmp.json")
        
    else:
        with open(model, 'r', encoding='utf-8') as f:
            data = json.load(f)
    
    classification_data = data.get("generate_classification", {})
    demos = classification_data.get("demos", [])
    signature = classification_data.get("signature", {})
    
    texts = [item.get("text", "") for item in demos]
    labels = [item.get("label", None) for item in demos]
    
    # Build full instructions string
    instructions_base = signature.get("instructions", "")
    fields_info = "\n".join(
        f"{field.get('prefix', '')} {field.get('description', '')}"
        for field in signature.get("fields", [])
    )
    
    full_instructions = f"{instructions_base}\n\nFields:\n{fields_info}"
    
    return texts, labels

def dspyjson2prompt_dep(file_path):
    """
    Read a DSPy JSON file and convert it to a formatted prompt.
    
    Args:
        file_path (str): Path to the JSON file
    
    Returns:
        str: Formatted prompt string
    """
    
    # Read the JSON file
    with open(file_path, 'r', encoding='utf-8') as file:
        data = json.load(file)
    
    # Extract relevant components
    classification_data = data.get('generate_classification', {})
    signature = classification_data.get('signature', {})
    demos = classification_data.get('demos', [])
    
    # Build the prompt
    prompt_parts = []
    
    # Add instructions if available
    instructions = signature.get('instructions', '')
    if instructions:
        prompt_parts.append(f"Instructions: {instructions}")
        prompt_parts.append("")
    
    # Add field descriptions
    fields = signature.get('fields', [])
    if fields:
        prompt_parts.append("Field Descriptions:")
        for field in fields:
            prefix = field.get('prefix', '')
            description = field.get('description', '')
            prompt_parts.append(f"- {prefix} {description}")
        prompt_parts.append("")
    
    # Add demonstration examples
    if demos:
        prompt_parts.append("Examples:")
        prompt_parts.append("")
        
        for i, demo in enumerate(demos, 1):
            text = demo.get('text', '')
            label = demo.get('label', '')
            
            prompt_parts.append(f"Text: {text}")
            prompt_parts.append(f"label: {label}")
            prompt_parts.append("")
    
    return "\n".join(prompt_parts)


def df2ds(df):
    import dspy
    dataset = [
        dspy.Example(text=row['text'], label=row['label']).with_inputs('text')
        for _, row in df.iterrows()
    ]
    
    return dataset

def model2prompt(model):
    model.save('./tmp.json')
    
    prompt = bs.dspyjson2prompt('./tmp.json')
    
    os.system("rm ./tmp.json")
    
    return prompt
    
    