import datetime
import time
import json
import uuid
import gradio as gr
from pathlib import Path
from .utils import *
from .log_utils import build_logger
from constants import IMAGE_DIR, EVALUATE_DIMS
from model.model_registry import model_info

t2s_logger = build_logger("gradio_web_server_text2shape", "gr_web_text2shape.log") # t2s = image generation, loggers for single model direct chat
t2s_multi_logger = build_logger("gradio_web_server_text2shape_multi", "gr_web_text2shape_multi.log") # t2s_multi = image generation multi, loggers for side-by-side and battle
i2s_logger = build_logger("gradio_web_server_image2shape", "gr_web_image2shape.log") # i2s = image editing, loggers for single model direct chat
i2s_multi_logger = build_logger("gradio_web_server_image2shape_multi", "gr_web_image2shape_multi.log") # i2s_multi = image editing multi, loggers for side-by-side and battle

def remove_suffix(name):
    if name.endswith("-t") or name.endswith("-i"):
        return name[:-2]
    else:
        return name

def vote_last_response_t2s(state, dim, vote_type, model_selector, request: gr.Request):
    with open(get_conv_log_filename(), "a") as fout:
        data = {
            "tstamp": round(time.time(), 4),
            "dim": dim,
            "type": vote_type,
            "model": model_selector,
            "state": state.dict(),
            "ip": get_ip(request),
        }
        fout.write(json.dumps(data) + "\n")
        append_json_item_on_log_server(data, get_conv_log_filename())
        
def vote_last_response_t2s_multi(states, dim, vote_type, is_anony: bool, model_selectors, request: gr.Request):
    with open(get_conv_log_filename(), "a") as fout:
        data = {
            "tstamp": round(time.time(), 4),
            "dim": dim,
            "type": vote_type,
            "anony": is_anony,
            "models": [x for x in model_selectors],
            "states": [x.dict() for x in states],
            "ip": get_ip(request),
        }
        fout.write(json.dumps(data) + "\n")
        append_json_item_on_log_server(data, get_conv_log_filename())
    # for state in states:
    #     output_file = f'{IMAGE_DIR}/text2shape/{state.conv_id}.png'
    #     with open(output_file, 'w') as f:
    #         state.output.save(f, 'PNG')
    #     save_image_file_on_log_server(output_file)

def vote_last_response_i2s(state, dim, vote_type, model_selector, request: gr.Request):
    with open(get_conv_log_filename(), "a") as fout:
        data = {
            "tstamp": round(time.time(), 4),
            "dim": dim,
            "type": vote_type,
            "model": model_selector,
            "state": state.dict(),
            "ip": get_ip(request),
        }
        fout.write(json.dumps(data) + "\n")
        append_json_item_on_log_server(data, get_conv_log_filename())
    # output_file = f'{IMAGE_DIR}/image2shape/{state.conv_id}.png'
    # source_file = f'{IMAGE_DIR}/image2shape/{state.conv_id}_source.png'
    # with open(output_file, 'w') as f:
    #     state.output.save(f, 'PNG')
    # with open(source_file, 'w') as sf:
    #     state.source_image.save(sf, 'PNG')
    # save_image_file_on_log_server(output_file)
    # save_image_file_on_log_server(source_file)
        
def vote_last_response_i2s_multi(states, dim, vote_type, is_anony, model_selectors, request: gr.Request):
    with open(get_conv_log_filename(), "a") as fout:
        data = {
            "tstamp": round(time.time(), 4),
            "dim": dim,
            "type": vote_type,
            "anony": is_anony,
            "models": [x for x in model_selectors],
            "states": [x.dict() for x in states],
            "ip": get_ip(request),
        }
        fout.write(json.dumps(data) + "\n")
        append_json_item_on_log_server(data, get_conv_log_filename())
    # for state in states:
    #     output_file = f'{IMAGE_DIR}/image2shape/{state.conv_id}.png'
    #     source_file = f'{IMAGE_DIR}/image2shape/{state.conv_id}_source.png'
    #     with open(output_file, 'w') as f:
    #         state.output.save(f, 'PNG')
    #     with open(source_file, 'w') as sf:
    #         state.source_image.save(sf, 'PNG')
    #     save_image_file_on_log_server(output_file)
    #     save_image_file_on_log_server(source_file)
            

## Text-to-Shape Generation (t2s) Single Model Direct Chat
def upvote_last_response_t2s(state, model_selector, dim_md, request: gr.Request):
    ip = get_ip(request)
    state.evaluted_dims += 1
    t2s_logger.info(f"upvote [{dim_md}]. ip: {ip}")
    vote_last_response_t2s(state, dim_md, "upvote", model_selector, request)
    return (state,) + (disable_btn,) * 3

def downvote_last_response_t2s(state, model_selector, dim_md, request: gr.Request):
    ip = get_ip(request)
    state.evaluted_dims += 1
    t2s_logger.info(f"downvote [{dim_md}]. ip: {ip}")
    vote_last_response_t2s(state, dim_md, "downvote", model_selector, request)
    return (state,) + (disable_btn,) * 3

def flag_last_response_t2s(state, model_selector, dim_md, request: gr.Request):
    ip = get_ip(request)
    state.evaluted_dims += 1
    t2s_logger.info(f"flag [{dim_md}]. ip: {ip}")
    vote_last_response_t2s(state, dim_md, "flag", model_selector, request)
    return (state,) + (disable_btn,) * 3


## Text-to-Shape Generation Multi (t2s_multi) Side-by-Side and Battle
def leftvote_last_response_t2s_named(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    t2s_multi_logger.info(f"leftvote [{dim_md}] (named). ip: {get_ip(request)}")
    vote_last_response_t2s_multi(
        [state0, state1], dim_md, "leftvote", False, [model_selector0, model_selector1], request
    )
    return (state0, state1) + (disable_btn,) * 4

def rightvote_last_response_t2s_named(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    t2s_multi_logger.info(f"rightvote [{dim_md}] (named). ip: {get_ip(request)}")
    vote_last_response_t2s_multi(
        [state0, state1], dim_md, "rightvote", False, [model_selector0, model_selector1], request
    )
    return (state0, state1) + (disable_btn,) * 4

def tievote_last_response_t2s_named(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    t2s_multi_logger.info(f"tievote [{dim_md}] (named). ip: {get_ip(request)}")
    vote_last_response_t2s_multi(
        [state0, state1], dim_md, "tievote", False, [model_selector0, model_selector1], request
    )
    return (state0, state1) + (disable_btn,) * 4

def bothbad_vote_last_response_t2s_named(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    t2s_multi_logger.info(f"bothbad_vote [{dim_md}] (named). ip: {get_ip(request)}")
    vote_last_response_t2s_multi(
        [state0, state1], dim_md, "bothbad_vote", False, [model_selector0, model_selector1], request
    )
    return (state0, state1) + (disable_btn,) * 4


def leftvote_last_response_t2s_anony(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    t2s_multi_logger.info(f"leftvote [{dim_md}] (anony). ip: {get_ip(request)}")
    vote_last_response_t2s_multi(
        [state0, state1], dim_md, "leftvote", True, [model_selector0, model_selector1], request
    )

    # if state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS:
    #     names = (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
    #     return (state0, state1) + (disable_btn,) * 4 + names
    # else:
    #     names = (gr.Markdown(f"### Model A: {state0.model_name}", visible=False), gr.Markdown(f"### Model B: {state1.model_name}", visible=False))
    #     return (state0, state1) + (disable_btn,) * 4 + names
    is_visible = (state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS)
    return (state0, state1) \
        + (disable_btn,) * 4 \
        + (gr.Markdown(f"### Model A: {model_info[state0.model_name].simple_name}", visible=is_visible), 
           gr.Markdown(f"### Model B: {model_info[state1.model_name].simple_name}", visible=is_visible))

def rightvote_last_response_t2s_anony(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    t2s_multi_logger.info(f"rightvote [{dim_md}] (anony). ip: {get_ip(request)}")
    vote_last_response_t2s_multi(
        [state0, state1], dim_md, "rightvote", True, [model_selector0, model_selector1], request
    )

    # if state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS:
    #     names = (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
    #     return (state0, state1) + (disable_btn,) * 4 + names
    # else:
    #     return (state0, state1) + (disable_btn,) * 4 + ("", "")
    is_visible = (state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS)
    return (state0, state1) \
        + (disable_btn,) * 4 \
        + (gr.Markdown(f"### Model A: {model_info[state0.model_name].simple_name}", visible=is_visible), 
           gr.Markdown(f"### Model B: {model_info[state1.model_name].simple_name}", visible=is_visible))

def tievote_last_response_t2s_anony(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    t2s_multi_logger.info(f"tievote [{dim_md}] (anony). ip: {get_ip(request)}")
    vote_last_response_t2s_multi(
        [state0, state1], dim_md, "tievote", True, [model_selector0, model_selector1], request
    )

    # if state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS:
    #     names = (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
    #     return (state0, state1) + (disable_btn,) * 4 + names
    # else:
    #     return (state0, state1) + (disable_btn,) * 4 + ("", "")
    is_visible = (state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS)
    return (state0, state1) \
        + (disable_btn,) * 4 \
        + (gr.Markdown(f"### Model A: {model_info[state0.model_name].simple_name}", visible=is_visible), 
           gr.Markdown(f"### Model B: {model_info[state1.model_name].simple_name}", visible=is_visible))
    
def bothbad_vote_last_response_t2s_anony(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    t2s_multi_logger.info(f"bothbad_vote [{dim_md}] (anony). ip: {get_ip(request)}")
    vote_last_response_t2s_multi(
        [state0, state1], dim_md, "bothbad_vote", True, [model_selector0, model_selector1], request
    )
    
    # if state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS:
    #     names = (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
    #     return (state0, state1) + (disable_btn,) * 4 + names
    # else:
    #     return (state0, state1) + (disable_btn,) * 4 + ("", "")
    is_visible = (state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS)
    return (state0, state1) \
        + (disable_btn,) * 4 \
        + (gr.Markdown(f"### Model A: {model_info[state0.model_name].simple_name}", visible=is_visible), 
           gr.Markdown(f"### Model B: {model_info[state1.model_name].simple_name}", visible=is_visible))

## Image-to-Shape (i2s) Single Model Direct Chat
def upvote_last_response_i2s(state, model_selector, dim_md, request: gr.Request):
    ip = get_ip(request)
    state.evaluted_dims += 1
    i2s_logger.info(f"upvote [{dim_md}]. ip: {ip}")
    vote_last_response_i2s(state, dim_md, "upvote", model_selector, request)
    return (state,) + (disable_btn,) * 3

def downvote_last_response_i2s(state, model_selector, dim_md, request: gr.Request):
    ip = get_ip(request)
    state.evaluted_dims += 1
    i2s_logger.info(f"downvote [{dim_md}]. ip: {ip}")
    vote_last_response_i2s(state, dim_md, "downvote", model_selector, request)
    return (state,) + (disable_btn,) * 3

def flag_last_response_i2s(state, model_selector, dim_md, request: gr.Request):
    ip = get_ip(request)
    state.evaluted_dims += 1
    i2s_logger.info(f"flag [{dim_md}]. ip: {ip}")
    vote_last_response_i2s(state, dim_md, "flag", model_selector, request)
    return (state,) + (disable_btn,) * 3


## Image-to-Shape Multi (i2s_multi) Side-by-Side and Battle
def leftvote_last_response_i2s_named(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    i2s_multi_logger.info(f"leftvote [{dim_md}] (named). ip: {get_ip(request)}")
    vote_last_response_i2s_multi(
        [state0, state1], dim_md, "leftvote", False, [model_selector0, model_selector1], request
    )
    return (state0, state1) + (disable_btn,) * 4
        
def rightvote_last_response_i2s_named(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    i2s_multi_logger.info(f"rightvote [{dim_md}] (named). ip: {get_ip(request)}")
    vote_last_response_i2s_multi(
        [state0, state1], dim_md, "rightvote", False, [model_selector0, model_selector1], request
    )
    return (state0, state1) + (disable_btn,) * 4

def tievote_last_response_i2s_named(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    i2s_multi_logger.info(f"tievote [{dim_md}] (named). ip: {get_ip(request)}")
    vote_last_response_i2s_multi(
        [state0, state1], dim_md, "tievote", False, [model_selector0, model_selector1], request
    )
    return (state0, state1) + (disable_btn,) * 4

def bothbad_vote_last_response_i2s_named(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    i2s_multi_logger.info(f"bothbad_vote [{dim_md}] (named). ip: {get_ip(request)}")
    vote_last_response_i2s_multi(
        [state0, state1], dim_md, "bothbad_vote", False, [model_selector0, model_selector1], request
    )
    return (state0, state1) + (disable_btn,) * 4


def leftvote_last_response_i2s_anony(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    i2s_multi_logger.info(f"leftvote [{dim_md}] (anony). ip: {get_ip(request)}")
    vote_last_response_i2s_multi(
        [state0, state1], dim_md, "leftvote", True, [model_selector0, model_selector1], request
    )

    # if state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS:
    #     names = (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
    #     return (state0, state1) + (disable_btn,) * 4 + names
    # else:
    #     return (state0, state1) + (disable_btn,) * 4 + ("", "")
    is_visible = (state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS)
    return (state0, state1) \
        + (disable_btn,) * 4 \
        + (gr.Markdown(f"### Model A: {model_info[state0.model_name].simple_name}", visible=is_visible), 
           gr.Markdown(f"### Model B: {model_info[state1.model_name].simple_name}", visible=is_visible))
        

def rightvote_last_response_i2s_anony(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    i2s_multi_logger.info(f"rightvote [{dim_md}] (anony). ip: {get_ip(request)}")
    vote_last_response_i2s_multi(
        [state0, state1], dim_md, "rightvote", True, [model_selector0, model_selector1], request
    )

    # if state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS:
    #     names = (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
    #     return (state0, state1) + (disable_btn,) * 4 + names
    # else:
    #     return (state0, state1) + (disable_btn,) * 4 + ("", "")
    is_visible = (state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS)
    return (state0, state1) \
        + (disable_btn,) * 4 \
        + (gr.Markdown(f"### Model A: {model_info[state0.model_name].simple_name}", visible=is_visible), 
           gr.Markdown(f"### Model B: {model_info[state1.model_name].simple_name}", visible=is_visible))
        

def tievote_last_response_i2s_anony(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    i2s_multi_logger.info(f"tievote [{dim_md}] (anony). ip: {get_ip(request)}")
    vote_last_response_i2s_multi(
        [state0, state1], dim_md, "tievote", True, [model_selector0, model_selector1], request
    )
    
    # if state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS:
    #     names = (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
    #     return (state0, state1) + (disable_btn,) * 4 + names
    # else:
    #     return (state0, state1) + (disable_btn,) * 4 + ("", "")
    is_visible = (state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS)
    return (state0, state1) \
        + (disable_btn,) * 4 \
        + (gr.Markdown(f"### Model A: {model_info[state0.model_name].simple_name}", visible=is_visible), 
           gr.Markdown(f"### Model B: {model_info[state1.model_name].simple_name}", visible=is_visible))
        

def bothbad_vote_last_response_i2s_anony(
    state0, state1, model_selector0, model_selector1, dim_md, request: gr.Request
):
    state0.evaluted_dims += 1
    state1.evaluted_dims += 1
    i2s_multi_logger.info(f"bothbad_vote [{dim_md}] (anony). ip: {get_ip(request)}")
    vote_last_response_i2s_multi(
        [state0, state1], dim_md, "bothbad_vote", True, [model_selector0, model_selector1], request
    )
    
    # if state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS:
    #     names = (gr.Markdown(f"### Model A: {state0.model_name}", visible=True), gr.Markdown(f"### Model B: {state1.model_name}", visible=True))
    #     return (state0, state1) + (disable_btn,) * 4 + names
    # else:
    #     return (state0, state1) + (disable_btn,) * 4 + ("", "")
    is_visible = (state0.evaluted_dims == state1.evaluted_dims == EVALUATE_DIMS)
    return (state0, state1) \
        + (disable_btn,) * 4 \
        + (gr.Markdown(f"### Model A: {model_info[state0.model_name].simple_name}", visible=is_visible), 
           gr.Markdown(f"### Model B: {model_info[state1.model_name].simple_name}", visible=is_visible))
        


share_js = """
function (a, b, c, d) {
    const captureElement = document.querySelector('#share-region-named');
    html2canvas(captureElement)
        .then(canvas => {
            canvas.style.display = 'none'
            document.body.appendChild(canvas)
            return canvas
        })
        .then(canvas => {
            const image = canvas.toDataURL('image/png')
            const a = document.createElement('a')
            a.setAttribute('download', 'chatbot-arena.png')
            a.setAttribute('href', image)
            a.click()
            canvas.remove()
        });
    return [a, b, c, d];
}
"""
def share_click_t2s_multi(state0, state1, model_selector0, model_selector1, request: gr.Request):
    t2s_multi_logger.info(f"share (anony). ip: {get_ip(request)}")
    if state0 is not None and state1 is not None:
        vote_last_response_t2s_multi(
            [state0, state1], "", "share", True, [model_selector0, model_selector1], request
        )

def share_click_i2s_multi(state0, state1, model_selector0, model_selector1, request: gr.Request):
    i2s_multi_logger.info(f"share (anony). ip: {get_ip(request)}")
    if state0 is not None and state1 is not None:
        vote_last_response_i2s_multi(
            [state0, state1], "", "share", True, [model_selector0, model_selector1], request
        )