import os
from django.apps import AppConfig
from django.db import transaction
from django.db.models.signals import post_migrate


def import_llm_backends(sender, **kwargs):
    backend_model = sender.get_model('LLMBackend')

    if backend_model.objects.count() != 0:
        return

    backend_model.objects.create(
        name='OpenAI',
        parameters={
            'api_key': 'sk-1234567890abcdef1234567890abcdef',
        },
        client_type='openai'
    )
    backend_model.objects.create(
        name='Ollama',
        parameters={
            'host': 'http://localhost:11434',
        },
        client_type='ollama'
    )
    backend_model.objects.create(
        name='Anyscale',
        parameters={
            "api_key": "esecret_12345abcde", 
            "base_url": "https://api.endpoints.anyscale.com/v1"
        },
        client_type='openai'
    )
    backend_model.objects.create(
        name='Anthropic',
        parameters={
            "api_key": "sk-12345abcde"
        },
        client_type='anthropic'
    )
    backend_model.objects.create(
        name='GoogleGenAI',
        parameters={
            "api_key": "sk-12345abcde"
        },
        client_type='genai'
    )
    backend_model.objects.create(
        name='LMStudio',
        parameters={
            "api_key": "lm-studio",
            "base_url": "http://localhost:1234/v1"

        },
        client_type='openai'
    )

def import_llm_models(sender, **kwargs):
    backend_model = sender.get_model('LLMBackend')
    llm_model = sender.get_model('LLMModel')

    if llm_model.objects.count() != 0:
        return

    openai_backend = backend_model.objects.get(name='OpenAI')
    llm_model.objects.create(
        name='gpt-4o',
        backend=openai_backend,
    )
    llm_model.objects.get_or_create(
        name='gpt-4o-mini',
        backend=openai_backend,
    )
    
    ollama_backend = backend_model.objects.get(name='Ollama')
    llm_model.objects.get_or_create(
        name='llama3',
        backend=ollama_backend,
    )

    anhropic_backend = backend_model.objects.get(name='Anthropic')
    llm_model.objects.get_or_create(
        name='claude-3-5-sonnet-20240620',
        backend=anhropic_backend,
    )

    anyscale_backend = backend_model.objects.get(name='Anyscale')
    llm_model.objects.get_or_create(
        name='meta-llama/Meta-Llama-3-70B-Instruct',
        backend=anyscale_backend,
    )

def import_datasets(sender, **kwargs):
    dataset_model = sender.get_model('Dataset')

    if dataset_model.objects.count() != 0:
        return

    dataset_model.objects.create(
        name='medqa',
        description='MedQA test dataset - 1273 questions',
    )
    
    os.system('python manage.py import_medqa --file datasets//medqa/test.jsonl --dataset medqa --target test')

    dataset_model.objects.create(
        name='pubmedqa',
        description='PubMedQA test[expert] dataset - 1000 questions',
    )
    os.system('python manage.py import_pubmedqa --file datasets//pubmedqa/ori_pqal.json --file-test-ground-truth=datasets/pubmedqa/test_ground_truth.json --dataset pubmedqa --target test')

    # dataset_model.objects.create(
    #     name='mmlu',
    #     description='MMLU test dataset - all 14012 questions',
    # )
    # os.system('python manage.py import_mmlu --subject all --dataset mmlu --target test')


def import_configs(sender, **kwargs):
    config_model = sender.get_model('EvalConfig')
    dataset = sender.get_model('Dataset')

    # get the medqa dataset
    medqa_dataset = dataset.objects.get(name='medqa')

    if config_model.objects.count() != 0:
        return
    
    with transaction.atomic():
        config_model.objects.create(
            name='MedQA - zero shot CoT',
            description='MedQA test dataset evaluation',
            dataset=medqa_dataset,
            sys_prompt="""
You are a helpful assistant designed to answer questions. 
Provide concise and accurate responses, ensuring you follow the provided instructions thoroughly.""".strip(),
            final_answer_pattern = '(?i)Final Answer\s*:\s*([A-E])'
        )

        # add role messages
        role_message_model = sender.get_model('RoleMessage')
        config = config_model.objects.get(name='MedQA - zero shot CoT')
        role_message_model.objects.create(
            role='user',
            eval_config=config,
            content = """
Answer the following multiple choice question. The last line of your response should be of the following format: 'Final Answer: $LETTER' (without quotes) 
where LETTER is the index of the correct answers from given choices. Think step by step before answering.""".strip()
        )

def import_sessions(sender, **kwargs):
    session_model = sender.get_model('EvalSession')
    config = sender.get_model('EvalConfig')
    llm_model = sender.get_model('LLMModel')

    if session_model.objects.count() != 0:
        return

    medqa_config = config.objects.get(name='MedQA - zero shot CoT')
    gpt4o = llm_model.objects.get(name='gpt-4o')
    session_model.objects.create(
        name='MedQA - gpt-4o',
        config=medqa_config,
        llm_model=gpt4o,
        dataset_target='test'
    )

def answer_interpreters(sender, **kwargs):
    answer_interpreter_model = sender.get_model('AnswerInterpreter')

    if answer_interpreter_model.objects.count() != 0:
        return
    
    llm_model = sender.get_model('LLMModel')
    gpt4o_mini = llm_model.objects.get(name='gpt-4o-mini')

    answer_interpreter_model.objects.create(
        name='YesNoMaybe Interpreter',
        prompt="""
You are a helpful assistant who will respond concisely to user requests. Your task is to examine the answer given by another virtual assistant and classify it as 'yes', 'no' or 'maybe'. Your final line of response should be 'Final Answer: yes | no | maybe' (without quotes).
Here is the assistant's answer:
"$QUESTION"
The final answer:""",
        llm_model=gpt4o_mini,
    )

    answer_interpreter_model.objects.create(
        name='MultiChoice Interpreter',
        prompt="""
You are a helpful assistant who will respond concisely to user requests. Your task is to examine the answer given by another virtual assistant and extract the letter corresponding to the selected option from multiple choices. Your final line of response should be "Final Answer: $LETTER" where $LETTER is the letter of the answer chosen by the assistant.
Here is the assistant's answer:
"$QUESTION"
The final answer:""",
        llm_model=gpt4o_mini,
    )

class CommonsConfig(AppConfig):
    default_auto_field = 'django.db.models.BigAutoField'
    name = 'commons'


    def ready(self):
        post_migrate.connect(import_llm_backends, sender=self)
        post_migrate.connect(import_llm_models, sender=self)
        post_migrate.connect(import_datasets, sender=self)
        post_migrate.connect(import_configs, sender=self)
        post_migrate.connect(import_sessions, sender=self)
        post_migrate.connect(answer_interpreters, sender=self)
    
