"""Hoot alignment and video generation application on modal."""

import json
import numpy as np

import modal

from suno_utils.worker.settings import lambda_client
from suno_utils.worker.modal_base import get_modal_base_image

from suno_utils.utils.elo import Tournament
import itertools

############## CHANGE THESE ##############

DEPLOYMENT_TYPE = "staging"  # staging, prod

APP_NAME = f"remi-thompson-cron-{DEPLOYMENT_TYPE}"
assert APP_NAME.endswith(DEPLOYMENT_TYPE)

SECRETS = [
    (
        modal.Secret.from_name("studio-aws-staging")
        if DEPLOYMENT_TYPE == "staging"
        else modal.Secret.from_name("studio-aws")
    ),
    modal.Secret.from_name("redis-test"),
]

base_image = get_modal_base_image()

app = modal.App(APP_NAME, image=base_image)

REMI_MODEL = "lyrics-remi"  # this is correct
REMI_MODEL_LIST = sorted(  # it's extremely important that this list remain sorted.
    [
        "remi-v1",
    ]
)
DEFAULT_MODEL = "lyrics-default"
DEFAULT_MODEL_LIST = sorted(
    [
        "default",
    ]
)
MODEL_LIST_LOOKUP = {REMI_MODEL: REMI_MODEL_LIST, DEFAULT_MODEL: DEFAULT_MODEL_LIST}


def invoke_lambda_function(function_name, payload=None):
    """Invokes an AWS Lambda function and returns the response."""
    if payload is None:
        payload = {}

    response = lambda_client.invoke(
        FunctionName=function_name,
        InvocationType="RequestResponse",  # Use 'Event' for async
        Payload=json.dumps(payload),
    )

    # Process and return the response
    if response["StatusCode"] == 200:
        response_payload = json.loads(response["Payload"].read().decode())
        if response_payload:
            return response_payload["body"]
        else:
            return {}
    else:
        raise Exception(f"Lambda invocation failed with status: {response['StatusCode']}")


def get_redis_keys_with_prefix(lyrics_model, data):
    """
    Invokes Lambda function to get Redis keys with specified prefix.

    Args:
        prefix (str): Redis key prefix to scan for

    Returns:
        list: List of keys as strings
    """
    prefix = f"{lyrics_model}:{data}:"
    try:
        lambda_response = invoke_lambda_function(
            function_name="redis-clip-accessor", payload={"scan_prefix": prefix}
        )
        if isinstance(lambda_response, str):
            lambda_response = json.loads(lambda_response)

        # Extract keys from the response
        if lambda_response and "keys" in lambda_response:
            return [key.replace(prefix, "", 1) for key in lambda_response["keys"]]
        else:
            print(f"Unexpected response format: {lambda_response}")
            return []

    except Exception as e:
        print(f"Error retrieving Redis keys: {e}")
        return []


def get_redis_combo_count(lyrics_model, first, second):
    """
    Invokes Lambda function to get all fields of a Redis hash.

    Args:
        first (str): First model ID
        second (str): Second model ID

    Returns:
        dict: Dictionary with 'first' and 'second' values mapped from the hash fields
              or empty dict if hash key doesn't exist
    """
    try:
        if first > second:
            first, second = second, first
        # Parse the hash key to extract model IDs
        hash_key = f"{lyrics_model}:preferences:{first}:{second}"
        parts = hash_key.split(":")
        if len(parts) >= 4:
            first_model = parts[-2]  # Second-to-last part
            second_model = parts[-1]  # Last part
        else:
            print(f"Hash key format invalid: {hash_key}")
            return {}

        lambda_response = invoke_lambda_function(
            function_name="redis-clip-accessor",
            payload={
                "hash_operation": "true",
                "hash_key": hash_key,
                "get_all_hash": "true",
            },
        )

        if isinstance(lambda_response, str):
            lambda_response = json.loads(lambda_response)

        # Check if we got a valid response with hash data
        if not lambda_response or "value" not in lambda_response:
            print(f"Hash key not found or invalid response: {lambda_response}")
            return {}

        hash_fields = lambda_response["value"]
        # Extract the specified fields using the parsed model IDs
        result = {}
        if "first" in hash_fields:
            result[first_model] = hash_fields["first"]
        else:
            result[first_model] = 0
        if "second" in hash_fields:
            result[second_model] = hash_fields["second"]
        else:
            result[second_model] = 0

        return result

    except Exception as e:
        print(f"Error retrieving Redis hash: {e}")
        return {}


def get_redis_preferences_list(lyrics_model):
    return get_redis_keys_with_prefix(lyrics_model=lyrics_model, data="preferences")


def get_redis_model_pair_prob_list(lyrics_model):
    return get_redis_keys_with_prefix(lyrics_model=lyrics_model, data="model-pair-prob")


# def get_model_pair_prob(first, second):
#     if first > second:
#         first, second = second, first
#     lambda_response = invoke_lambda_function(
#         function_name="redis-clip-accessor",
#         payload={"key": f"lyrics-remi:model-pair-prob:{first}:{second}"},
#     )
#     if isinstance(lambda_response, str):
#         lambda_response = json.loads(lambda_response)

#     # Check if we got a valid response with hash data
#     if not lambda_response or "value" not in lambda_response:
#         print(f"model pair prob not found or invalid response: {lambda_response}")
#         return {}
#     return lambda_response["value"]


def set_model_pair_prob(lyrics_model, first, second, probability):
    """
    Updates the probability value for a model pair in Redis.

    Args:
        first (str): First model ID
        second (str): Second model ID
        probability (float): Probability value to set

    Returns:
        bool: True if successful, False otherwise
    """
    # Ensure lexicographical ordering
    if first > second:
        first, second = second, first

    try:
        # Call Lambda to set the value
        key = f"{lyrics_model}:model-pair-prob:{first}:{second}"
        lambda_response = invoke_lambda_function(
            function_name="redis-clip-accessor",
            payload={
                "set_operation": "true",
                "key": key,
                "value": probability,  # Convert to string as Redis stores strings
            },
        )

        # Parse response if it's a string
        if isinstance(lambda_response, str):
            lambda_response = json.loads(lambda_response)

        # Check for success using the actual response format
        if (
            lambda_response
            and "message" in lambda_response
            and lambda_response["message"] == "Key set successfully"
        ):
            print(f"Successfully set probability for {key} to {probability}")
            return True
        else:
            print(f"Failed to set probability: {lambda_response}")
            return False

    except Exception as e:
        print(f"Error setting model pair probability: {e}")
        return False


def _roll_up_pair_probs(pair_probs):
    players = sorted(set(sum([[a, b] for (a, b) in pair_probs], [])))
    ps = {pl: sum(v for (a, b), v in pair_probs.items() if pl in [a, b]) for pl in players}
    P = sum(ps.values())
    assert np.isclose(P, 2.0), P
    return ps


def update_thompson_sampler(lyrics_model):
    preferences_set = set(get_redis_preferences_list(lyrics_model))
    print(preferences_set)
    model_list = MODEL_LIST_LOOKUP[lyrics_model]
    model_combos = list(itertools.combinations(model_list, 2))

    tournament = Tournament()
    for first, second in model_combos:
        assert first < second  # NB: we don't need to count the self vs self instances here.
        match_count = {first: 0, second: 0}
        if f"{first}:{second}" in preferences_set:
            match_count = get_redis_combo_count(lyrics_model, first, second)
            print(f"got {(first, second), match_count} from redis")
        tournament.match_counts[(first, second)] = (match_count[first], match_count[second])
    print("match counts:", tournament.match_counts)
    try:
        elo_scores = tournament.calc_elo_scores()
    except Exception as e:
        print("failed to calculate elo scores", e)
        model_pairs = [
            (model_a, model_b) for model_a in model_list for model_b in model_list if model_a <= model_b
        ]
        probability = 1 / len(model_pairs)
        for model_a, model_b in model_pairs:
            set_model_pair_prob(lyrics_model, model_a, model_b, probability)
        return
    print("Successfully calculated ELO scores:", elo_scores)
    print("ELO confidence intervals:", tournament.calc_elo_bootstrap_confidence_intervals())
    pair_probs = tournament.get_selection_pair_probs()
    print("selection pair probs:", pair_probs)
    model_probs = {
        player: sum(prob for (first, second), prob in pair_probs.items() if player in [first, second])
        for player in tournament.get_players()
    }
    print("calculated model probs")
    for model, prob in reversed(sorted(model_probs.items(), key=lambda kv: kv[1])):
        print(model, prob)
    for (first, second), probability in pair_probs.items():
        assert first <= second
        set_model_pair_prob(lyrics_model, first, second, probability)


@app.function(
    schedule=modal.Cron(
        "*/5 * * * *",  # Run every 5 minutes
    ),
    secrets=SECRETS,
    cloud="aws",
    region="us-east",
)
def update_remi_sampler():
    update_thompson_sampler(REMI_MODEL)


@app.function(
    schedule=modal.Cron(
        "*/5 * * * *",  # Run every 5 minutes
    ),
    secrets=SECRETS,
    cloud="aws",
    region="us-east",
)
def update_default_sampler():
    update_thompson_sampler(DEFAULT_MODEL)