import re
import gradio as gr
import spaces
import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

print("Initializing...")

########################################
# MODEL LOADING
########################################
model_id = "LatitudeGames/Wayfarer-12B"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_id)

print("Loading model in 4-bit precision...")
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    load_in_4bit=True
)

print("Model loaded. Creating pipeline...")
pipe = pipeline(
    task="text-generation",
    model=model,
    tokenizer=tokenizer,
    do_sample=True,
    temperature=0.8,
    min_p=0.025,
    repetition_penalty=1.05,
    max_new_tokens=450,
    pad_token_id=tokenizer.eos_token_id,
    eos_token_id=tokenizer.eos_token_id,
)

########################################
# ChatML prompts
########################################
SYSTEM_PROMPT = (
    "<|im_start|>system\n"
    "You're a masterful storyteller and gamemaster. Write in second person present tense (You are), "
    "crafting vivid, engaging narratives with authority and confidence.<|im_end|>"
)

DETAILED_SCENARIO = (
    "Welcome, adventurous soul, to a land of pastoral beauty and hidden perils. "
    "You find yourself upon timeworn cobblestone streets in a quaint fantasy village, "
    "nestled among rolling hills veiled by a gentle morning mist. The villagers whisper of "
    "a terrifying beast said to roam the outskirts at dusk, devouring livestock and frightening "
    "all who wander too far.\n\n"
    "Merchants hawk their wares in a bustling marketplace, while children chase after chickens "
    "in muddy lanes. Every so often, you catch sight of heavily-armed travellers passing through, "
    "eyes vigilant for any lurking danger. Rumours abound—some claim an ancient curse befalls "
    "this land, others say the beast is merely a savage wolf grown unnaturally large. Whatever "
    "the truth, it falls upon you to unravel the mystery and save these troubled folk—or perish "
    "in the attempt."
)

STARTING_PROMPT = (
    "\n<|im_start|>assistant\n"
    + DETAILED_SCENARIO + "<|im_end|>"
)

def build_displayed_text(chat_state):
    lines = []
    scenario_included = False
    for m in chat_state:
        if m["role"] == "assistant":
            if (not scenario_included) and (m["content"] == DETAILED_SCENARIO):
                lines.append(m["content"])
                scenario_included = True
            else:
                lines.append(m["content"])
        elif m["role"] == "user":
            lines.append(f"> {m['content']}")
        # We skip the system entirely

    return "\n\n".join(lines).strip()

########################################
punct_pattern = re.compile(r"[.!?\"']$")
leading_you = re.compile(r"^\\s*you\\s+", re.IGNORECASE)
leading_i   = re.compile(r"^\\s*i\\s+", re.IGNORECASE)
assistant_pat = re.compile(r"d?assistant", re.IGNORECASE)
system_pat_start = re.compile(r"^\\s*[A-Za-z]+system\\s*", re.IGNORECASE)
single_word_pat = re.compile(r"^[A-Za-z]+$")
remove_sys_block_pat = re.compile(r"<\\|im_start\\|>system.*?<\\|im_end\\|>", re.DOTALL)
remove_chatml_block_pat = re.compile(r"<\\|im_start\\|>.*?<\\|im_end\\|>", re.DOTALL)

########################################
# Helper functions
########################################

def maybe_add_period(text: str) -> str:
    text = text.strip()
    if not text:
        return text
    if punct_pattern.search(text[-1]):
        return text
    return text + "."

def first_to_second_person(text: str) -> str:
    text = text.strip()
    text = leading_you.sub("", text)
    text = leading_i.sub("", text)

    replacements = [
        (re.compile(r"\\bI\\b", re.IGNORECASE), "you"),
        (re.compile(r"\\bme\\b", re.IGNORECASE), "you"),
        (re.compile(r"\\bmy\\b", re.IGNORECASE), "your"),
        (re.compile(r"\\bmine\\b", re.IGNORECASE), "yours"),
        (re.compile(r"\\bmyself\\b", re.IGNORECASE), "yourself"),
        (re.compile(r"I\\'m", re.IGNORECASE), "you're"),
        (re.compile(r"I am", re.IGNORECASE), "you are"),
    ]
    for (pat, repl) in replacements:
        text = pat.sub(repl, text)
    return text.strip()

def strip_all_system_blocks(text: str) -> str:
    return remove_sys_block_pat.sub("", text)

def strip_all_chatml_blocks(text: str) -> str:
    return remove_chatml_block_pat.sub("", text)

def strip_spurious_tokens(text: str) -> str:
    cleaned = assistant_pat.sub("", text)
    cleaned = system_pat_start.sub("", cleaned)

    lines = cleaned.split("\n")
    if lines:
        first_line = lines[0].strip()
        if single_word_pat.match(first_line) and len(first_line) < 20:
            lines.pop(0)
    cleaned = "\n".join(lines)
    return cleaned.strip()

def parse_generated_for_assistant_content(full_generation: str, prompt_text: str) -> str:
    idx = full_generation.find(prompt_text)
    if idx != -1:
        full_generation = full_generation[idx + len(prompt_text):]

    cleaned = strip_all_system_blocks(full_generation)

    pattern = r"<\\|im_start\\|>assistant(.*?)<\\|im_end\\|>"
    matches = re.findall(pattern, cleaned, flags=re.DOTALL)
    if matches:
        final_chunk = matches[-1]
    else:
        final_chunk = cleaned

    final_chunk = strip_all_chatml_blocks(final_chunk)
    final_chunk = strip_spurious_tokens(final_chunk)
    return final_chunk.strip()

########################################
# We'll store messages in an array.
# Format: [ {"role": "system"|"user"|"assistant", "content": ...}, ... ]
# We'll convert them to ChatML when we call the pipeline.
########################################
def messages_to_chatml(messages):
    chatml = ""
    for m in messages:
        if m["role"] == "system":
            chatml += f"\n<|im_start|>system\n{m['content']}<|im_end|>"
        elif m["role"] == "user":
            chatml += f"\n<|im_start|>user\n> {m['content']}<|im_end|>"
        elif m["role"] == "assistant":
            chatml += f"\n<|im_start|>assistant\n{m['content']}<|im_end|>"
        else:
            pass
    return chatml

initial_messages = [
    {"role": "system", "content": "You're a masterful storyteller and gamemaster. Write in second person present tense (You are), crafting vivid, engaging narratives with authority and confidence."},
    {"role": "assistant", "content": DETAILED_SCENARIO},
]

DEFAULT_CHAT_STATE = initial_messages

def build_user_line(raw_user: str, action: str) -> str:
    if action == "Story":
        final_text = raw_user.strip()
    else:
        final_text = first_to_second_person(raw_user)

    final_text = maybe_add_period(final_text)

    def starts_with_you(t: str) -> bool:
        return t.lower().startswith("you ") or t.lower().startswith("you'")

    if action == "Do":
        if starts_with_you(final_text):
            return final_text
        return f"You {final_text}"
    elif action == "Say":
        if final_text.lower().startswith("you say"):
            return final_text
        return f"You say \"{final_text}\""
    else:
        return final_text

def remove_last_user_line_in_assistant(chat_state, assistant_text):
    last_user = None
    for m in reversed(chat_state):
        if m["role"] == "user":
            last_user = m["content"]
            break
    if last_user and last_user in assistant_text:
        assistant_text = assistant_text.replace(last_user, "")
        assistant_text = re.sub(r"\n\n+", "\n\n", assistant_text.strip())
    return assistant_text.strip()

@spaces.GPU(duration=60)
def handle_send(
    chat_state,
    displayed_text,
    user_text,
    action,
    last_user_text,
    last_displayed_before_ai
):
    #print("=== handle_send called ===")
    if not user_text.strip():
        raise gr.Error("You must type something before sending!")

    final_user_line = build_user_line(user_text, action)
    chat_state.append({"role": "user", "content": final_user_line})

    chatml_input = messages_to_chatml(chat_state)
    #print("=== Chat State BEFORE generation ===")
    #print(chatml_input)

    response = pipe(chatml_input)
    raw_generation = response[0]["generated_text"]

    #print("=== Raw Generation ===")
    #print(raw_generation)

    parsed_ai = parse_generated_for_assistant_content(raw_generation, chatml_input)
    parsed_ai = remove_last_user_line_in_assistant(chat_state, parsed_ai)
    chat_state.append({"role": "assistant", "content": parsed_ai})

    new_displayed = build_displayed_text(chat_state)

    #print("=== Chat State AFTER generation ===")
    final_chatml = messages_to_chatml(chat_state)
    #print(final_chatml)
    #print("=== Displayed Text ===")
    #print(new_displayed)

    return (
        chat_state,
        new_displayed,
        "",
        user_text,
        displayed_text
    )

@spaces.GPU(duration=60)
def handle_retry(
    chat_state,
    displayed_text,
    last_user_text,
    last_displayed_before_ai,
    action
):
    #print("=== handle_retry called ===")

    if not last_user_text.strip():
        raise gr.Error("No user prompt to resend!")

    roles_to_remove = ["assistant", "user"]
    for r in roles_to_remove:
        if len(chat_state) > 0 and chat_state[-1]["role"] == r:
            chat_state.pop()

    #print("=== After removing last user/assistant from chat_state===")
    final_chatml = messages_to_chatml(chat_state)
    #print(final_chatml)

    final_user_line = build_user_line(last_user_text, action)
    chat_state.append({"role": "user", "content": final_user_line})

    chatml_input = messages_to_chatml(chat_state)
    #print("=== Chat State BEFORE generation (retry) ===")
    #print(chatml_input)

    response = pipe(chatml_input)
    raw_generation = response[0]["generated_text"]

    #print("=== Raw Generation (retry) ===")
    #print(raw_generation)

    parsed_ai = parse_generated_for_assistant_content(raw_generation, chatml_input)
    parsed_ai = remove_last_user_line_in_assistant(chat_state, parsed_ai)
    chat_state.append({"role": "assistant", "content": parsed_ai})

    new_displayed = build_displayed_text(chat_state)

    #print("=== Chat State AFTER generation (retry) ===")
    final_chatml = messages_to_chatml(chat_state)
    #print(final_chatml)
    #print("=== Displayed Text (retry) ===")
    #print(new_displayed)

    return (
        chat_state,
        new_displayed,
        "",
        last_user_text,
        new_displayed
    )
    
########################################
# UI
########################################
with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.orange)) as interface:
    gr.Markdown(
        """ 
        # **GradioDungeon**  
        
        Generate adventures using the Wayfarer-12B model by Latitude. The model may generate inappropriate or harmful content, and by interacting with it, you agree that you understand this risk.
        """
    )

    chat_state = gr.State(value=DEFAULT_CHAT_STATE)
    current_scenario = gr.State(value=DETAILED_SCENARIO)

    displayed_text = gr.Textbox(
        lines=20,
        label="Conversation",
        value=DETAILED_SCENARIO,
        interactive=False,
    )

    last_user_text = gr.State(value="")
    last_displayed_before_ai = gr.State(value="")

    with gr.Row():
        action_selector = gr.Radio([
            "Do", "Say", "Story"
        ], value="Do", label="Action")

    user_input = gr.Textbox(
        lines=2,
        label="Your Input",
        placeholder="(Type your action or dialogue here. Write in second person, example using 'you' instead of me.)"
    )

    with gr.Row():
        send_btn = gr.Button("Send", variant="primary")
        retry_btn = gr.Button("Retry", variant="secondary")

    send_btn.click(
        fn=handle_send,
        inputs=[chat_state, displayed_text, user_input, action_selector, last_user_text, last_displayed_before_ai],
        outputs=[chat_state, displayed_text, user_input, last_user_text, last_displayed_before_ai]
    )

    retry_btn.click(
        fn=handle_retry,
        inputs=[chat_state, displayed_text, last_user_text, last_displayed_before_ai, action_selector],
        outputs=[chat_state, displayed_text, user_input, last_user_text, last_displayed_before_ai]
    )

    with gr.Accordion("Edit Starting Scenario", open=False):
        prompt_edit_dropdown = gr.Dropdown(
            choices=["Custom Scenario", "Default Scenario", "Sci-Fi Setting", "Medieval Quest"],
            label="Select Scenario Template"
        )
        
        custom_prompt_box = gr.Textbox(
            lines=4, 
            label="Custom Scenario", 
            placeholder="Enter your own custom starting scenario..."
        )
        
        apply_scenario_btn = gr.Button("Apply New Scenario")

    def apply_new_scenario(dropdown_choice, custom_prompt):
        if dropdown_choice == "Custom Scenario":
            if not custom_prompt.strip():
                raise gr.Error("Custom scenario cannot be empty!")
            new_scenario = custom_prompt
        elif dropdown_choice == "Sci-Fi Setting":
            new_scenario = (
                "You find yourself aboard a massive interstellar cruiser, its metal corridors stretching endlessly. "
                "Alarms blare softly in the distance, and the ship's AI speaks in measured tones about an "
                "unidentified anomaly approaching from the starboard side. Your mission, still unclear, "
                "seems to hang in the balance of the next few moments.\n\n"
                "\"Initiating threat assessment protocol,\" the AI's synthetic voice announces. \"Crew member designation: "
                "you are required to report to the command deck immediately.\""
            )
        elif dropdown_choice == "Medieval Quest":
            new_scenario = (
                "In the shadow of a towering castle, you stand among a group of weary knights and adventurers. "
                "The king's messenger has just delivered grave news of a dark sorcerer threatening the realm. "
                "Whispers of an ancient prophecy circulate, and all eyes turn to you as a potential savior.\n\n"
                "\"The prophecy speaks of one marked by an unusual birthmark,\" an old wizard mutters, his gnarled finger "
                "tracing a pattern in the air near your shoulder. \"If the legends are true, you might be our last hope.\""
            )
        else:
            new_scenario = DETAILED_SCENARIO

        new_initial_messages = [
            {"role": "system", "content": "You're a masterful storyteller and gamemaster. Write in second person present tense (You are), crafting vivid, engaging narratives with authority and confidence."},
            {"role": "assistant", "content": new_scenario},
        ]
        
        return {
            chat_state: new_initial_messages,
            displayed_text: new_scenario,
            current_scenario: new_scenario,
        }

    apply_scenario_btn.click(
        fn=apply_new_scenario, 
        inputs=[prompt_edit_dropdown, custom_prompt_box], 
        outputs=[chat_state, displayed_text, current_scenario]
    )

    interface.launch()