import gradio as gr
import random
import spaces
import torch
from transformers import pipeline
import csv
import os
import uuid
import datetime

DATA_DIR = os.path.join(os.getcwd(), "session_data")
os.makedirs(DATA_DIR, exist_ok=True)
csv_path = None


pipe = pipeline("text-generation", model="mecoffey/NPC_brain")

DESC_SYSTEM_PROMPT = (
    "Write a description of a fantasy character based on the user prompt"
    "response Must include: fantasy Race (Human, Dragonborn, Dwarf, Eladrin, Elf, Fairy, Firbolg, Gnome, Goblin, Half-elf, Half-orc, Halfling, Human, Orc, Tiefling, other as specified by prompt), height, build, eye color, hair color, clothing."
    "response may include: noticeable features such as scars, tattoos, jewelry, or expression"
    "Response must not include any extraneous details."
    "response must be no longer than 150 words"
)
BACK_SYSTEM_PROMPT = (
    "format the backstory as a timeline in a bulleted list with 1-3 sentences per point."
    "Keep each point short and fitting for a background character in a fantasy world."
    "include things like occupations they have had, family ties, livelihood, or mannerisms."
    "response must include one line about what the character is thinking or doing right now."
    "Do not include any descriptions of physical features."
)


DEFAULT_SESSION_STATE = {
    "csv_path": None,
    "strength": "",
    "dexterity": "",
    "constitution": "",
    "intelligence": "",
    "wisdom": "",
    "charisma": "",
    "ac": "",
    "hp": "",
    "speed": "",
    "charm": "",
    "pref_traits": "",
    "ick_traits": "",
}

def new_session_state():
    return dict(DEFAULT_SESSION_STATE)

def stats_gen(state):
    state = dict(state or DEFAULT_SESSION_STATE)
    state.update({
        "strength": random.randint(2, 12),
        "dexterity": random.randint(2, 12),
        "constitution": random.randint(2, 12),
        "intelligence": random.randint(2, 12),
        "wisdom": random.randint(2, 12),
        "charisma": random.randint(2, 12),
    })
    return state["strength"], state["dexterity"], state["constitution"], state["intelligence"], state["wisdom"], state["charisma"], state

def stats_gen_2(state):
    state = dict(state or DEFAULT_SESSION_STATE)
    state.update({
        "ac": random.randint(2, 14),
        "hp": state.get("constitution", 0) + random.randint(1, 6),
        "speed": "25 ft",
    })
    return state["ac"], state["hp"], state["speed"], state

def traits_picker(filepath):
    with open(filepath, 'r') as f:
        items = [line.strip() for line in f if line.strip()]
    count = random.randint(1, min(3, len(items)))
    selected = random.sample(items,count)
    return ', '.join(selected)

def rand_race_gen(filepath):
    with open(filepath, 'r') as f:
        items = [line.strip() for line in f if line.strip()]
    return random.choice(items)

def bard_stats_gen(state):
    state = dict(state or DEFAULT_SESSION_STATE)
    state.update({
        "charm": random.randint(1, 30),
        "pref_traits": traits_picker('traits.txt'),
        "ick_traits": traits_picker('traits.txt'),
    })
    return state["charm"], state["pref_traits"], state["ick_traits"], state


def generate_character(prompt, state):
    state = dict(state or DEFAULT_SESSION_STATE)
    state.update({
        "strength": random.randint(2, 12),
        "dexterity": random.randint(2, 12),
        "constitution": random.randint(2, 12),
        "intelligence": random.randint(2, 12),
        "wisdom": random.randint(2, 12),
        "charisma": random.randint(2, 12),
    })
    state.update({
        "ac": random.randint(2, 14),
        "hp": state.get("constitution", 0) + random.randint(1, 6),
        "speed": "25 ft",
    })
    state.update({
        "charm": random.randint(1, 30),
        "pref_traits": traits_picker('traits.txt'),
        "ick_traits": traits_picker('traits.txt'),
    })
    if not state.get("csv_path"):
        state["csv_path"] = make_session_csv()
    description, backstory = char_description(prompt)
    append_to_csv(
        state["csv_path"],
        prompt,
        description,
        backstory,
        state["strength"],
        state["dexterity"],
        state["constitution"],
        state["intelligence"],
        state["wisdom"],
        state["charisma"],
        state["ac"],
        state["hp"],
        state["speed"],
        state["charm"],
        state["pref_traits"],
        state["ick_traits"],
    )
    return (
        state["strength"],
        state["dexterity"],
        state["constitution"],
        state["intelligence"],
        state["wisdom"],
        state["charisma"],
        state["ac"],
        state["hp"],
        state["speed"],
        state["charm"],
        state["pref_traits"],
        state["ick_traits"],
        description,
        backstory,
        state,
    )


def generate_random_character(state):
    state = dict(state or DEFAULT_SESSION_STATE)
    state.update({
        "strength": random.randint(2, 12),
        "dexterity": random.randint(2, 12),
        "constitution": random.randint(2, 12),
        "intelligence": random.randint(2, 12),
        "wisdom": random.randint(2, 12),
        "charisma": random.randint(2, 12),
    })
    state.update({
        "ac": random.randint(2, 14),
        "hp": state.get("constitution", 0) + random.randint(1, 6),
        "speed": "25 ft",
    })
    state.update({
        "charm": random.randint(1, 30),
        "pref_traits": traits_picker('traits.txt'),
        "ick_traits": traits_picker('traits.txt'),
    })
    if not state.get("csv_path"):
        state["csv_path"] = make_session_csv()
    rand_race = rand_race_gen('core_races.txt')
    description, backstory = rand_char_description(rand_race)
    append_to_csv(
        state["csv_path"],
        rand_race,
        description,
        backstory,
        state["strength"],
        state["dexterity"],
        state["constitution"],
        state["intelligence"],
        state["wisdom"],
        state["charisma"],
        state["ac"],
        state["hp"],
        state["speed"],
        state["charm"],
        state["pref_traits"],
        state["ick_traits"],
    )
    return (
        state["strength"],
        state["dexterity"],
        state["constitution"],
        state["intelligence"],
        state["wisdom"],
        state["charisma"],
        state["ac"],
        state["hp"],
        state["speed"],
        state["charm"],
        state["pref_traits"],
        state["ick_traits"],
        description,
        backstory,
        state,
    )


def make_session_csv():
    session_id = uuid.uuid4().hex
    return os.path.join(DATA_DIR, f"story_session_{session_id}.csv")


def append_to_csv(csv_path, prompt, description, backstory, strength, dexterity, constitution, intelligence, wisdom, charisma, ac, hp, speed, charm, pref_traits, ick_traits):
    is_new = not os.path.exists(csv_path)
    with open(csv_path, "a", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        if is_new:
            writer.writerow(["timestamp", "prompt", "description", "backstory", "strength", "dexterity", "constitution", "intelligence", "wisdom", "charisma", "ac", "hp", "speed", "charm", "pref_traits", "ick_traits"])
        writer.writerow([
            datetime.datetime.utcnow().isoformat(),
            prompt,
            description,
            backstory,
            strength,
            dexterity,
            constitution,
            intelligence,
            wisdom,
            charisma,
            ac,
            hp,
            speed,
            charm,
            pref_traits,
            ick_traits,
        ])

def char_description(prompt):

    prompt1 = [{"role": "user", "content": prompt}]
    output1 = pipe(prompt1, return_full_text=False, max_new_tokens=256)
    output1_text = output1[0]["generated_text"].strip()
    combine = str(prompt) + "\n" + str(output1_text) + "\nBackstory:"
    prompt2 = [
        {"role": "user", "content": combine},
    ]
    output2 = pipe(prompt2, return_full_text=False, max_new_tokens=256)
    return output1_text, output2[0]["generated_text"].strip()

def rand_char_description(prompt):
    prompt1 = [{"role": "user", "content": prompt}]
    output1 = pipe(prompt1, return_full_text=False, max_new_tokens=256)
    output1_text = output1[0]["generated_text"].strip()
    combine = str(prompt) + "\n" + str(output1_text) + "\nBackstory:"
    prompt2 = [
        {"role": "user", "content": combine},
    ]
    output2 = pipe(prompt2, return_full_text=False, max_new_tokens=256)
    return output1_text, output2[0]["generated_text"].strip()


def char_description_csv(prompt, state):
    state = dict(state or DEFAULT_SESSION_STATE)
    if not state.get("csv_path"):
        state["csv_path"] = make_session_csv()
    description, backstory = char_description(prompt)
    append_to_csv(
        state["csv_path"],
        prompt,
        description,
        backstory,
        state.get("strength", ""),
        state.get("dexterity", ""),
        state.get("constitution", ""),
        state.get("intelligence", ""),
        state.get("wisdom", ""),
        state.get("charisma", ""),
        state.get("ac", ""),
        state.get("hp", ""),
        state.get("speed", ""),
        state.get("charm", ""),
        state.get("pref_traits", ""),
        state.get("ick_traits", ""),
    )
    return description, backstory, state


def rand_char_description_csv(state):
    state = dict(state or DEFAULT_SESSION_STATE)
    if not state.get("csv_path"):
        state["csv_path"] = make_session_csv()
    rand_race = rand_race_gen('core_races.txt')
    description, backstory = rand_char_description(rand_race)
    append_to_csv(
        state["csv_path"],
        rand_race,
        description,
        backstory,
        state.get("strength", ""),
        state.get("dexterity", ""),
        state.get("constitution", ""),
        state.get("intelligence", ""),
        state.get("wisdom", ""),
        state.get("charisma", ""),
        state.get("ac", ""),
        state.get("hp", ""),
        state.get("speed", ""),
        state.get("charm", ""),
        state.get("pref_traits", ""),
        state.get("ick_traits", ""),
    )
    return description, backstory, state


def get_csv_file(state):
    state = dict(state or DEFAULT_SESSION_STATE)
    if not state.get("csv_path"):
        state["csv_path"] = make_session_csv()
    return state["csv_path"], state


##UI is fine as-is
with gr.Blocks() as demo:
    title = gr.HTML("<center><h1>NPC Generator</h1></center>")
    description = gr.HTML("<center>For tabletop GMs when players ask for a description of random people in the tavern.</center>")
    tags = gr.HTML("<center>| Backyard AI | Open BMB | Tiny Titan | Fine-tuned | Modal |</center>")
    demo_video = gr.HTML("<center>See a demo here!: https://youtu.be/b7M0ejRoxQU</center>")

    with gr.Row(equal_height = True):
        with gr.Column(scale = 2):
            text_in = gr.Textbox(label="What npc do you need?", placeholder = "An Orc Barkeep")
        with gr.Column(scale = 0):
            btn = gr.Button("Generate") 
            rand = gr.Button("Random")       

    with gr.Row():

        session_state = gr.State(value=new_session_state())

        with gr.Column(scale = 0):
            stre_box = gr.Textbox(value = "", label = "Strength", interactive = False)
            dex_box = gr.Textbox(value = "", label = "Dexterity", interactive = False)
            con_box = gr.Textbox(value = "", label = "Constitution", interactive = False)
            intl_box = gr.Textbox(value = "", label = "Intelligence", interactive = False)
            wis_box = gr.Textbox(value = "", label = "Wisdom", interactive = False)
            char_box = gr.Textbox(value = "", label = "Charisma", interactive = False)

        with gr.Column(scale = 2):
            text_out1 = gr.Textbox(label = "Description")
            text_out2 = gr.Textbox(label = "Backstory")

        with gr.Column(scale = 0):
            ac_box = gr.Textbox(value = "", label = "AC", interactive = False)
            hp_box = gr.Textbox(value = "", label = "HP", interactive = False)
            speed_box = gr.Textbox(value = "", label = "Speed", interactive = False)
            bard_stats = gr.HTML("~~Bard Stats~~")
            with gr.Accordion(open = False):
                dcs = gr.Textbox(value = "", label = "DC to Charm")
                prefs = gr.Textbox(value = "", label = "Advantage if:")
                icks = gr.Textbox(value = "", label = "Disdvantage if:")

        btn.click(
            fn=generate_character,
            inputs=[text_in, session_state],
            outputs=[
                stre_box,
                dex_box,
                con_box,
                intl_box,
                wis_box,
                char_box,
                ac_box,
                hp_box,
                speed_box,
                dcs,
                prefs,
                icks,
                text_out1,
                text_out2,
                session_state,
            ],
        )
        rand.click(
            fn=generate_random_character,
            inputs=session_state,
            outputs=[
                stre_box,
                dex_box,
                con_box,
                intl_box,
                wis_box,
                char_box,
                ac_box,
                hp_box,
                speed_box,
                dcs,
                prefs,
                icks,
                text_out1,
                text_out2,
                session_state,
            ],
        )

    with gr.Row():
        download_btn = gr.Button("Download CSV")
    with gr.Row():
        download_file = gr.File(label="Generated CSV", interactive=False)
        download_btn.click(fn=get_csv_file, inputs=session_state, outputs=[download_file, session_state])


demo.launch(theme = gr.themes.Monochrome())