import re
import copy
import random
import gradio as gr
from gradio_client import Client
from pathlib import Path

from modules import (
	ImageMaker, MusicMaker, palmchat, merge_video
)
from interfaces import utils

from pingpong import PingPong
from pingpong.context import CtxLastWindowStrategy

# TODO: Replace checkpoint filename to Huggingface URL
img_maker = ImageMaker('landscapeAnimePro_v20Inspiration.safetensors') # without VAE
#img_maker = ImageMaker('landscapeAnimePro_v20Inspiration.safetensors', vae="stabilityai/sd-vae-ft-mse")
#img_maker = ImageMaker('fantasyworldFp16.safetensors', vae="cute20vae.safetensors")
#img_maker = ImageMaker('forgesagalandscapemi.safetensors', vae="anythingFp16.safetensors")
bgm_maker = MusicMaker(model_size='small', output_format='mp3')

video_gen_client_url = "https://0447df3cf5f7c49c46.gradio.live"

async def update_story_gen(
	cursors, cur_cursor_idx,
	genre, place, mood,
	main_char_name, main_char_age, main_char_mbti, main_char_personality, main_char_job,
	side_char_enable1, side_char_name1, side_char_age1, side_char_mbti1, side_char_personality1, side_char_job1,
	side_char_enable2, side_char_name2, side_char_age2, side_char_mbti2, side_char_personality2, side_char_job2,
	side_char_enable3, side_char_name3, side_char_age3, side_char_mbti3, side_char_personality3, side_char_job3,
):
    if len(cursors) == 1:
        return await first_story_gen(
			cursors,
			genre, place, mood,
			main_char_name, main_char_age, main_char_mbti, main_char_personality, main_char_job,
			side_char_enable1, side_char_name1, side_char_age1, side_char_mbti1, side_char_personality1, side_char_job1,
			side_char_enable2, side_char_name2, side_char_age2, side_char_mbti2, side_char_personality2, side_char_job2,
			side_char_enable3, side_char_name3, side_char_age3, side_char_mbti3, side_char_personality3, side_char_job3,
			cur_cursor_idx=cur_cursor_idx
		)
    else:
        return await next_story_gen(
			cursors,
			None,
			genre, place, mood,
			main_char_name, main_char_age, main_char_mbti, main_char_personality, main_char_job,
			side_char_enable1, side_char_name1, side_char_age1, side_char_mbti1, side_char_personality1, side_char_job1,
			side_char_enable2, side_char_name2, side_char_age2, side_char_mbti2, side_char_personality2, side_char_job2,
			side_char_enable3, side_char_name3, side_char_age3, side_char_mbti3, side_char_personality3, side_char_job3,
			cur_cursor_idx=cur_cursor_idx
		)

async def next_story_gen(
	cursors,
	action,
	genre, place, mood,
	main_char_name, main_char_age, main_char_mbti, main_char_personality, main_char_job,
	side_char_enable1, side_char_name1, side_char_age1, side_char_mbti1, side_char_personality1, side_char_job1,
	side_char_enable2, side_char_name2, side_char_age2, side_char_mbti2, side_char_personality2, side_char_job2,
	side_char_enable3, side_char_name3, side_char_age3, side_char_mbti3, side_char_personality3, side_char_job3,	
	cur_cursor_idx=None
):
	stories = ""
	cur_side_chars = 1

	action = cursors[cur_cursor_idx]["action"] if cur_cursor_idx is not None else action
	end_idx = len(cursors) if cur_cursor_idx is None else len(cursors)-1

	for cursor in cursors[:end_idx]:
		stories = stories + cursor["story"]

	prompt = f"""Write the next paragraphs. The next paragraphs should be determined by an option and well connected to the current stories. 

background information:
- genre: {genre}
- where: {place}
- mood: {mood}

main character
- name: {main_char_name}
- job: {main_char_job}
- age: {main_char_age}
- mbti: {main_char_mbti}
- personality: {main_char_personality}
"""

	prompt, cur_side_chars = utils.add_side_character(
		side_char_enable1, prompt, cur_side_chars,
		side_char_name1, side_char_job1, side_char_age1, side_char_mbti1, side_char_personality1
	)
	prompt, cur_side_chars = utils.add_side_character(
		side_char_enable2, prompt, cur_side_chars,
		side_char_name2, side_char_job2, side_char_age2, side_char_mbti2, side_char_personality2
	)
	prompt, cur_side_chars = utils.add_side_character(
		side_char_enable3, prompt, cur_side_chars,
		side_char_name3, side_char_job3, side_char_age3, side_char_mbti3, side_char_personality3
	)

	prompt = prompt + f"""
stories
{stories}

option to the next stories: {action}

Fill in the following JSON output format:
{{
	"paragraphs": "string"
}}

"""

	print(f"generated prompt:\n{prompt}")
	parameters = {
		'model': 'models/text-bison-001',
		'candidate_count': 1,
		'temperature': 1.0,
		'top_k': 40,
		'top_p': 1,
		'max_output_tokens': 4096,
	}    	
	response_json = await utils.retry_until_valid_json(prompt, parameters=parameters)

	story = response_json["paragraphs"]
	if isinstance(story, list):
		story = "\n\n".join(story)
  
	if cur_cursor_idx is None:
		cursors.append({
			"title": "",
			"story": story,
			"action": action
		})
	else:
		cursors[cur_cursor_idx]["story"] = story
		cursors[cur_cursor_idx]["action"] = action

	return (
		cursors, len(cursors)-1,
		story,
		gr.update(
			maximum=len(cursors), value=len(cursors),
			label=f"{len(cursors)} out of {len(cursors)} stories",
			visible=True, interactive=True
		),
		gr.update(interactive=True),
		gr.update(interactive=True),
		gr.update(value=None, visible=False, interactive=True),
		gr.update(value=None, visible=False, interactive=True),
		gr.update(value=None, visible=False, interactive=True),	
	)

async def actions_gen(
	cursors,
	genre, place, mood,
	main_char_name, main_char_age, main_char_mbti, main_char_personality, main_char_job,
	side_char_enable1, side_char_name1, side_char_age1, side_char_mbti1, side_char_personality1, side_char_job1,
	side_char_enable2, side_char_name2, side_char_age2, side_char_mbti2, side_char_personality2, side_char_job2,
	side_char_enable3, side_char_name3, side_char_age3, side_char_mbti3, side_char_personality3, side_char_job3,
	cur_cursor_idx=None
):
	stories = ""
	cur_side_chars = 1
	end_idx = len(cursors) if cur_cursor_idx is None else len(cursors)-1

	for cursor in cursors[:end_idx]:
		stories = stories + cursor["story"]

	summary_prompt = f"""Summarize the text below

{stories}

"""
	print(f"generated prompt:\n{summary_prompt}")
	parameters = {
		'model': 'models/text-bison-001',
		'candidate_count': 1,
		'temperature': 1.0,
		'top_k': 40,
		'top_p': 1,
		'max_output_tokens': 4096,
	}    	
	_, summary = await palmchat.gen_text(summary_prompt, mode="text", parameters=parameters)

	prompt = f"""Suggest the 30 options to drive the stories to the next based on the information below. 

background information:
- genre: {genre}
- where: {place}
- mood: {mood}

main character
- name: {main_char_name}
- job: {main_char_job}
- age: {main_char_age}
- mbti: {main_char_mbti}
- personality: {main_char_personality}
"""
	prompt, cur_side_chars = utils.add_side_character(
		side_char_enable1, prompt, cur_side_chars,
		side_char_name1, side_char_job1, side_char_age1, side_char_mbti1, side_char_personality1
	)
	prompt, cur_side_chars = utils.add_side_character(
		side_char_enable2, prompt, cur_side_chars,
		side_char_name2, side_char_job2, side_char_age2, side_char_mbti2, side_char_personality2
	)
	prompt, cur_side_chars = utils.add_side_character(
		side_char_enable3, prompt, cur_side_chars,
		side_char_name3, side_char_job3, side_char_age3, side_char_mbti3, side_char_personality3
	)

	prompt = prompt + f"""
summary of the story
{summary}

Fill in the following JSON output format:
{{
    "options": ["string", "string", "string", ...]
}}

"""

	print(f"generated prompt:\n{prompt}")
	parameters = {
		'model': 'models/text-bison-001',
		'candidate_count': 1,
		'temperature': 1.0,
		'top_k': 40,
		'top_p': 1,
		'max_output_tokens': 4096,
	}    	
	response_json = await utils.retry_until_valid_json(prompt, parameters=parameters)
	actions = response_json["options"]

	random_actions = random.sample(actions, 3)

	return (
		gr.update(value=random_actions[0], interactive=True),
		gr.update(value=random_actions[1], interactive=True),
		gr.update(value=random_actions[2], interactive=True),
		"   "
	)

async def first_story_gen(
	cursors,
	genre, place, mood,
	main_char_name, main_char_age, main_char_mbti, main_char_personality, main_char_job,
	side_char_enable1, side_char_name1, side_char_age1, side_char_mbti1, side_char_personality1, side_char_job1,
	side_char_enable2, side_char_name2, side_char_age2, side_char_mbti2, side_char_personality2, side_char_job2,
	side_char_enable3, side_char_name3, side_char_age3, side_char_mbti3, side_char_personality3, side_char_job3,
	cur_cursor_idx=None
):
	cur_side_chars = 1

	prompt = f"""Write the first three paragraphs of a novel as much detailed as possible. They should be based on the background information. Blend 5W1H principle into the stories as a plain text. Don't let the paragraphs end the whole story.

background information:
- genre: {genre}
- where: {place}
- mood: {mood}

main character
- name: {main_char_name}
- job: {main_char_job}
- age: {main_char_age}
- mbti: {main_char_mbti}
- personality: {main_char_personality}
"""

	prompt, cur_side_chars = utils.add_side_character(
		side_char_enable1, prompt, cur_side_chars,
		side_char_name1, side_char_job1, side_char_age1, side_char_mbti1, side_char_personality1
	)
	prompt, cur_side_chars = utils.add_side_character(
		side_char_enable2, prompt, cur_side_chars,
		side_char_name2, side_char_job2, side_char_age2, side_char_mbti2, side_char_personality2
	)
	prompt, cur_side_chars = utils.add_side_character(
		side_char_enable3, prompt, cur_side_chars,
		side_char_name3, side_char_job3, side_char_age3, side_char_mbti3, side_char_personality3
	)

	prompt = prompt + f"""
Fill in the following JSON output format:
{{
	"paragraphs": "string"
}}

"""	

	print(f"generated prompt:\n{prompt}")
	parameters = {
		'model': 'models/text-bison-001',
		'candidate_count': 1,
		'temperature': 1.0,
		'top_k': 40,
		'top_p': 1,
		'max_output_tokens': 4096,
	}    	
	response_json = await utils.retry_until_valid_json(prompt, parameters=parameters)

	story = response_json["paragraphs"]
	if isinstance(story, list):
		story = "\n\n".join(story)
  
	if cur_cursor_idx is None:
		cursors.append({
			"title": "",
			"story": story
		})
	else:
		cursors[cur_cursor_idx]["story"] = story

	return (
		cursors, len(cursors)-1,
		story,
		gr.update(
			maximum=len(cursors), value=len(cursors),
			label=f"{len(cursors)} out of {len(cursors)} stories",
			visible=False if len(cursors) == 1 else True, interactive=True
		),
		gr.update(interactive=True),
		gr.update(interactive=True),
		gr.update(value=None, visible=False, interactive=True),
		gr.update(value=None, visible=False, interactive=True),	
  		gr.update(value=None, visible=False, interactive=True),
	)

def video_gen(
	image, audio, title, cursors, cur_cursor, use_ffmpeg=True
):
	if use_ffmpeg:
		output_filename = merge_video(image, audio, story_title="")
	
	if not use_ffmpeg or not output_filename:
		client = Client(video_gen_client_url)
		result = client.predict(
			"",
			audio,
			image,
			f"{utils.id_generator()}.mp4",
			api_name="/predict"
		)
		output_filename = result[0]

	cursors[cur_cursor]["video"] = output_filename

	return (
		gr.update(visible=False),
		gr.update(visible=False),
		gr.update(visible=True, value=output_filename),
		cursors,
		"   "
	)


def image_gen(
	genre, place, mood, title, story_content, cursors, cur_cursor
):
	# generate prompts for background image with PaLM
	for _ in range(3):
		try:
			prompt, neg_prompt = img_maker.generate_background_prompts(genre, place, mood, title, "", story_content)
			neg_prompt
			print(f"Image Prompt: {prompt}")
			print(f"Negative Prompt: {neg_prompt}")
			break
		except Exception as e:
			print(e)
			
	if not prompt:
		raise ValueError("Failed to generate prompts for background image.")

	# generate image
	try:
		img_filename = img_maker.text2image(prompt, neg_prompt=neg_prompt, ratio='16:9', cfg=6.5)
	except ValueError as e:
		print(e)
		img_filename = str(Path('.') / 'assets' / 'nsfw_warning_wide.png')
	
	cursors[cur_cursor]["img"] = img_filename

	return  (
		gr.update(visible=True, value=img_filename),
		cursors,
		"  "
	)


def audio_gen(
	genre, place, mood, title, story_content, cursors, cur_cursor
):
	# generate prompt for background music with PaLM
	for _ in range(3):
		try:
			prompt = bgm_maker.generate_prompt(genre, place, mood, title, "", story_content)
			print(f"Music Prompt: {prompt}")
			break
		except Exception as e:
			print(e)

	if not prompt:
		raise ValueError("Failed to generate prompt for background music.")

	# generate music
	bgm_filename = bgm_maker.text2music(prompt, length=60)
	cursors[cur_cursor]["audio"] = bgm_filename

	return (
		gr.update(visible=True, value=bgm_filename),
		cursors,
		" "
	)

def move_story_cursor(moved_cursor, cursors):
	cursor_content = cursors[moved_cursor-1]
	max_cursor = len(cursors)

	action_btn = (
			gr.update(interactive=False),
			gr.update(interactive=False),
			gr.update(interactive=False)
	)

	if moved_cursor == max_cursor:
		action_btn = (
			gr.update(interactive=True),
			gr.update(interactive=True),
			gr.update(interactive=True)
		)

	if "video" in cursor_content:
		outputs = (
			moved_cursor-1,
			gr.update(label=f"{moved_cursor} out of {len(cursors)} chapters"),
			cursor_content["story"],
			gr.update(value=None, visible=False),
			gr.update(value=None, visible=False),
			gr.update(value=cursor_content["video"], visible=True),
		)

	else:
		image_container = gr.update(value=None, visible=False)
		audio_container = gr.update(value=None, visible=False)

		if "img" in cursor_content:
			image_container = gr.update(value=cursor_content["img"], visible=True)

		if "audio" in cursor_content:
			audio_container = gr.update(value=cursor_content["audio"], visible=True)

		outputs = (
			moved_cursor-1,
			gr.update(label=f"{moved_cursor} out of {len(cursors)} stories"),
			cursor_content["story"],
			image_container,
			audio_container,
			gr.update(value=None, visible=False),
		)        

	return outputs + action_btn

def update_story_content(story_content, cursors, cur_cursor):
	cursors[cur_cursor]["story"] = story_content
	return cursors

def disable_btns():
	return (
		gr.update(interactive=False), # image_gen_btn
		gr.update(interactive=False), # audio_gen_btn
		gr.update(interactive=False), # img_audio_combine_btn

		gr.update(interactive=False), # regen_actions_btn
		gr.update(interactive=False), # regen_story_btn
		gr.update(interactive=False), # custom_prompt_txt

		gr.update(interactive=False), # action_btn1
		gr.update(interactive=False), # action_btn2
		gr.update(interactive=False), # action_btn3

		gr.update(interactive=False), # custom_action_txt

		gr.update(interactive=False), # restart_from_story_generation_btn
		gr.update(interactive=False), # story_writing_done_btn
	)
 
def enable_btns(story_image, story_audio):
	video_gen_btn_state = gr.update(interactive=False)
    
	if story_image is not None and \
		story_audio is not None:
		video_gen_btn_state = gr.update(interactive=True)
    
	return (
		gr.update(interactive=True), # image_gen_btn
		gr.update(interactive=True), # audio_gen_btn
		video_gen_btn_state, # img_audio_combine_btn

		gr.update(interactive=True), # regen_actions_btn
		gr.update(interactive=True), # regen_story_btn
		gr.update(interactive=True), # custom_prompt_txt

		gr.update(interactive=True), # action_btn1
		gr.update(interactive=True), # action_btn2
		gr.update(interactive=True), # action_btn3

		gr.update(interactive=True), # custom_action_txt

		gr.update(interactive=True), # restart_from_story_generation_btn
		gr.update(interactive=True), # story_writing_done_btn
	)