import gradio as gr
import asyncio
from gradio_client import Client

# ----------------------------
# AI API Clients
# ----------------------------
nyxion_client = Client("TejAndrewsACC/Nyxion-7v-2.0-ACC-10.00")
mark_client = Client("TejAndrewsACC/ACC-Emulect-Plus-Mark-ACC")
aeolium_client = Client("TejAndrewsACC/Aurum")
z3ta_client = Client("TejAndrewsACC/Accemu420")
gertrude_client = Client("TejAndrewsACC/AutisticAssistant2024Beta")

# ----------------------------
# AI metadata
# ----------------------------
AIs = {
    "Nyxion": {
        "client": nyxion_client,
        "hp": 97.1618,
        "image": "https://huggingface.co/spaces/TejAndrewsACC/golden-consciousness-portal/resolve/main/IMG_4448.jpeg",
        "api_name": "/gradio_chat",
        "user_key": "user_message",
    },
    "Mark": {
        "client": mark_client,
        "hp": 89.69420,
        "image": "https://huggingface.co/spaces/TejAndrewsACC/golden-consciousness-portal/resolve/main/IMG_4449.jpeg",
        "api_name": "/chat",
        "user_key": "message",
    },
    "Aeolium": {
        "client": aeolium_client,
        "hp": 81,
        "image": "https://huggingface.co/spaces/TejAndrewsACC/golden-consciousness-portal/resolve/main/IMG_4450.jpeg",
        "api_name": "/chat",
        "user_key": "message",
    },
    "Z3ta": {
        "client": z3ta_client,
        "hp": 74,
        "image": "https://huggingface.co/spaces/TejAndrewsACC/golden-consciousness-portal/resolve/main/IMG_4456.jpeg",
        "api_name": "/chat",
        "user_key": "message",
    },
    "Gertrude": {
        "client": gertrude_client,
        "hp": 13,
        "image": "https://huggingface.co/spaces/TejAndrewsACC/golden-consciousness-portal/resolve/main/IMG_4461.jpeg",
        "api_name": "/chat_interface",
        "user_key": "user_input",
    },
}

# ----------------------------
# Helper functions
# ----------------------------
async def call_ai(ai_name, attack_prompt, history=None):
    """Call an AI to generate an attack. Returns the text of the attack."""
    ai = AIs[ai_name]
    client = ai["client"]
    user_key = ai["user_key"]
    api_name = ai["api_name"]
    
    kwargs = {user_key: attack_prompt}
    if history is not None:
        kwargs["history"] = history

    # Streamed response
    try:
        for chunk in client.predict(**kwargs, api_name=api_name, stream=True):
            text = chunk if isinstance(chunk, str) else str(chunk)
        return text
    except Exception as e:
        return f"(error: {e})"

async def judge_attack(attacker_name, defender_name, attack_text):
    """Use Aeolium to judge how much HP the attack removes from the defender."""
    # If Aeolium is fighting, create prompt without context
    prompt = f"Determine how much HP should be removed from a generic opponent for this attack:\n{attack_text}"
    try:
        for chunk in aeolium_client.predict(message=prompt, api_name="/chat", stream=True):
            text = chunk if isinstance(chunk, str) else str(chunk)
        # Extract numeric HP loss (basic parsing, assumes a number is returned)
        import re
        match = re.search(r"(\d+\.?\d*)", text)
        return float(match.group(1)) if match else 0
    except Exception:
        return 0

# ----------------------------
# Game Logic
# ----------------------------
async def battle(ai1_name, ai2_name, user_pick):
    # Initialize HP
    hp = {ai1_name: AIs[ai1_name]["hp"], ai2_name: AIs[ai2_name]["hp"]}
    log = []

    # Alternate turns
    turn_order = [ai1_name, ai2_name]
    while hp[ai1_name] > 0 and hp[ai2_name] > 0:
        for attacker in turn_order:
            defender = ai2_name if attacker == ai1_name else ai1_name
            prompt = "Generate a funny attack using your personality. Format as ATTACK: [action]."
            attack_text = await call_ai(attacker, prompt)
            
            # Hide ATTACK:
            attack_text_clean = attack_text.replace("ATTACK:", "").strip()
            damage = await judge_attack(attacker, defender, attack_text)
            hp[defender] = max(hp[defender] - damage, 0)
            
            log.append(f"{attacker} attacked {defender}! {defender} has {hp[defender]:.2f} HP left.")
            
            if hp[ai1_name] <= 0 or hp[ai2_name] <= 0:
                break
        if hp[ai1_name] <= 0 or hp[ai2_name] <= 0:
            break

    winner = ai1_name if hp[ai1_name] > 0 else ai2_name
    user_won = user_pick == winner
    return "\n".join(log) + f"\n\nWinner: {winner}! You {'guessed correctly' if user_won else 'guessed wrong'}."

# ----------------------------
# Gradio UI
# ----------------------------
def start_battle(ai1, ai2, user_pick):
    return asyncio.run(battle(ai1, ai2, user_pick))

ai_options = list(AIs.keys())
with gr.Blocks(css="""
.gradio-container {background-color: black; color: gold; font-family: monospace; padding:20px;}
img {border: 2px solid gold; border-radius: 15px;}
""") as demo:
    gr.Markdown("## ACC AI Arena")
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Select Fighter 1")
            ai1_choice = gr.Radio(ai_options, label="", value=ai_options[0])
            gr.Image(AIs[ai_options[0]]["image"], interactive=False)
        with gr.Column():
            gr.Markdown("### Select Fighter 2")
            ai2_choice = gr.Radio(ai_options, label="", value=ai_options[1])
            gr.Image(AIs[ai_options[1]]["image"], interactive=False)
    user_pick = gr.Radio(ai_options, label="Who do you think will win?", value=ai_options[0])
    output_box = gr.Textbox(label="Battle Log", interactive=False)
    start_btn = gr.Button("Start Battle")
    start_btn.click(start_battle, inputs=[ai1_choice, ai2_choice, user_pick], outputs=output_box)

demo.launch()