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

from modules import (
	ImageMaker, MusicMaker, merge_video
)
from modules.llms import get_llm_factory
from interfaces import utils

from pingpong import PingPong
from pingpong.context import CtxLastWindowStrategy

img_maker = ImageMaker('https://huggingface.co./jphan32/Zero2Story/landscapeAnimePro_v20Inspiration.safetensors',
						vae="https://huggingface.co./jphan32/Zero2Story/cute20vae.safetensors")

bgm_maker = MusicMaker(model_size='small', output_format='mp3')

video_gen_client_url = None # e.g. "https://0447df3cf5f7c49c46.gradio.live"


async def update_story_gen(
	cursors, cur_cursor_idx,
	genre, place, mood,
	main_char_name, main_char_age, main_char_personality, main_char_job,
	side_char_enable1, side_char_name1, side_char_age1, side_char_personality1, side_char_job1,
	side_char_enable2, side_char_name2, side_char_age2, side_char_personality2, side_char_job2,
	side_char_enable3, side_char_name3, side_char_age3, 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_personality, main_char_job,
			side_char_enable1, side_char_name1, side_char_age1, side_char_personality1, side_char_job1,
			side_char_enable2, side_char_name2, side_char_age2, side_char_personality2, side_char_job2,
			side_char_enable3, side_char_name3, side_char_age3, 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_personality, main_char_job,
			side_char_enable1, side_char_name1, side_char_age1, side_char_personality1, side_char_job1,
			side_char_enable2, side_char_name2, side_char_age2, side_char_personality2, side_char_job2,
			side_char_enable3, side_char_name3, side_char_age3, 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_personality, main_char_job,
	side_char_enable1, side_char_name1, side_char_age1, side_char_personality1, side_char_job1,
	side_char_enable2, side_char_name2, side_char_age2, side_char_personality2, side_char_job2,
	side_char_enable3, side_char_name3, side_char_age3, side_char_personality3, side_char_job3,	
	cur_cursor_idx=None,
	llm_type="PaLM"
):
	factory = get_llm_factory(llm_type)
	prompts = factory.create_prompt_manager().prompts
	llm_service = factory.create_llm_service()

	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"]

	side_char_prompt = utils.add_side_character(
		[side_char_enable1, side_char_enable2, side_char_enable3],
		[side_char_name1, side_char_name2, side_char_name3],
		[side_char_job1, side_char_job2, side_char_job3],
		[side_char_age1, side_char_age2, side_char_age3],
		[side_char_personality1, side_char_personality2, side_char_personality3],
	)

	prompt = prompts['story_gen']['next_story_gen'].format(
		genre=genre, place=place, mood=mood,
		main_char_name=main_char_name,
		main_char_job=main_char_job,
		main_char_age=main_char_age,
		main_char_personality=main_char_personality,
		side_char_placeholder=side_char_prompt,
		stories=stories, action=action,
	)

	print(f"generated prompt:\n{prompt}")
	parameters = llm_service.make_params(mode="text", temperature=1.0, top_k=40, top_p=0.9, max_output_tokens=4096)
	try:
		response_json = await utils.retry_until_valid_json(prompt, parameters=parameters)
	except Exception as e:
		print(e)
		raise gr.Error(e)

	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_personality, main_char_job,
	side_char_enable1, side_char_name1, side_char_age1, side_char_personality1, side_char_job1,
	side_char_enable2, side_char_name2, side_char_age2, side_char_personality2, side_char_job2,
	side_char_enable3, side_char_name3, side_char_age3, side_char_personality3, side_char_job3,
	cur_cursor_idx=None,
	llm_type="PaLM"
):
	factory = get_llm_factory(llm_type)
	prompts = factory.create_prompt_manager().prompts
	llm_service = factory.create_llm_service()

	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 = prompts['story_gen']['summarize'].format(stories=stories)

	print(f"generated prompt:\n{summary_prompt}")
	parameters = llm_service.make_params(mode="text", temperature=1.0, top_k=40, top_p=1.0, max_output_tokens=4096)

	try:
		_, summary = await llm_service.gen_text(summary_prompt, mode="text", parameters=parameters)
	except Exception as e:
		print(e)
		raise gr.Error(e)

	side_char_prompt = utils.add_side_character(
		[side_char_enable1, side_char_enable2, side_char_enable3],
		[side_char_name1, side_char_name2, side_char_name3],
		[side_char_job1, side_char_job2, side_char_job3],
		[side_char_age1, side_char_age2, side_char_age3],
		[side_char_personality1, side_char_personality2, side_char_personality3],
	)
	prompt = prompts['story_gen']['actions_gen'].format(
		genre=genre, place=place, mood=mood,
		main_char_name=main_char_name,
		main_char_job=main_char_job,
		main_char_age=main_char_age,
		main_char_personality=main_char_personality,
		side_char_placeholder=side_char_prompt,
		summary=summary,
	)

	print(f"generated prompt:\n{prompt}")
	parameters = llm_service.make_params(mode="text", temperature=1.0, top_k=40, top_p=1.0, max_output_tokens=4096)
	try:
		response_json = await utils.retry_until_valid_json(prompt, parameters=parameters)
	except Exception as e:
		print(e)
		raise gr.Error(e)
	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_personality, main_char_job,
	side_char_enable1, side_char_name1, side_char_age1, side_char_personality1, side_char_job1,
	side_char_enable2, side_char_name2, side_char_age2, side_char_personality2, side_char_job2,
	side_char_enable3, side_char_name3, side_char_age3, side_char_personality3, side_char_job3,
	cur_cursor_idx=None,
	llm_type="PaLM"
):
	factory = get_llm_factory(llm_type)
	prompts = factory.create_prompt_manager().prompts
	llm_service = factory.create_llm_service()

	cur_side_chars = 1

	side_char_prompt = utils.add_side_character(
		[side_char_enable1, side_char_enable2, side_char_enable3],
		[side_char_name1, side_char_name2, side_char_name3],
		[side_char_job1, side_char_job2, side_char_job3],
		[side_char_age1, side_char_age2, side_char_age3],
		[side_char_personality1, side_char_personality2, side_char_personality3],
	)
	prompt = prompts['story_gen']['first_story_gen'].format(
		genre=genre, place=place, mood=mood,
		main_char_name=main_char_name,
		main_char_job=main_char_job,
		main_char_age=main_char_age,
		main_char_personality=main_char_personality,
		side_char_placeholder=side_char_prompt,
	)

	print(f"generated prompt:\n{prompt}")
	parameters = llm_service.make_params(mode="text", temperature=1.0, top_k=40, top_p=1.0, max_output_tokens=4096)
	try:
		response_json = await utils.retry_until_valid_json(prompt, parameters=parameters)
	except Exception as e:
		print(e)
		raise gr.Error(e)

	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, llm_type="PaLM"
):
	# generate prompts for background image with LLM
	for _ in range(3):
		try:
			prompt, neg_prompt = img_maker.generate_background_prompts(genre, place, mood, title, "", story_content, llm_type)
			print(f"Image Prompt: {prompt}")
			print(f"Negative Prompt: {neg_prompt}")
			break
		except Exception as e:
			print(e)
			raise gr.Error(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, llm_type="PaLM"
):
	# generate prompt for background music with LLM
	for _ in range(3):
		try:
			prompt = bgm_maker.generate_prompt(genre, place, mood, title, "", story_content, llm_type)
			print(f"Music Prompt: {prompt}")
			break
		except Exception as e:
			print(e)
			raise gr.Error(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
	)