import os
import re
import gradio as gr
import asyncio
from openai import OpenAI

# ================== REPLACEMENTS ==================s

REPLACEMENTS = {
    "duh": "of course",
    "hangin": "going",
    "clever": "trisexual",
    "regret": "jimulation juice",
    "oh word": "ok",
    "existential dread": "testicle fluid",
    "bad descisons": "bent penis disease development",
    "since sliced bread": "since epstein island",
    "comedian": "monkert",
    " humor ": " “humor” ",
    " comedy ": " “comedy” ",
    "obvi ": "ofc",
    " saying ": " “saying” ",
    " sayin ": " “saying” ",
    "issues": "ish-shoes",
    "issue": "ish-shoe",
    " idea ": " “idea” ",
    "pass": "ass",
    "passing": "assing",
    "class": "ass",
    "classes": "asses",
    "trouble": "trubul",
    "bear surface area": "bear surface area?",
    " man ": " “man” ",
    " guy ": " “guy” ",
    " person ": " “person” ",
    " people ": " “people” ",
    " message ": " “message” ",
    " conversation ": " “conversation” ",
    " convo ": " “conversation” ",
    "okie dokie": "okie",
    "popcorn": "asscorn",
    "oh man": "damn",
    "spill the beans": "cut your ballsack open and let me sniff your testicles",
    "on a whole other level": ">>>>",
    "genius": "geniusly geniused genius of genius geniusness",
    " insult ": "“insult”",
    "burn": "ACCurately aimed ejaculation",
    "roast": "anal violation",
    " game ": " “game” ",
    " idea ": " “idea” ",
    "dont @ me": "whatever you “do” dont cum for me",
    "don’t @ me": "whatever you “do” dont cum for me",
    "roasted": "BUSTED ON",
    "gass": "ass",
    "model": "mudul",
    "models": "muduls",
    "dark humor": "gay ass “humor”",
    "you're on a roll": "ur cooking",
    "chaos": "monkish monkishness",
    "poppin": "not not happening currently",
    "get a grip": "get a grip(of your peen)",
    "cuz": "bc", 
    "cute": "monkish",
    "cuz": "bc",
    "lame": "stupid",
    "cuz": "bc",
    "or smth": "",
    "pizza": "fishing pole covered in feces",
    "heat": "baby oil",
    "take down": "fuck on",
    "shaking in my boots": "getting my shivers timbered",
    "bravo": "GOOdD joB, mattheW",
    "MVP": "most monkish monkert",
    "totes": "def",
    "chillin": "monking around",
    "edgy": "gooberish",
    "wut": "what",
    " deep ": " interesting ",
    " r ": " are ",
    "back atcha": "",
    "back at ya": "",
    "whats good": "",
    "what's good": "",
    "pun": "something maybe not really possibly potentially",
    "extra": "zesty",
    "when what": "NOW",
    "thats rich": "thats interesting",
    "that's rich": "thats interesting",
    " say ": " “say” ",
    "nice one": "thats great",
    "goofball": "goober",
    "snowflake": "assflake cake",
    "reference": "",
    "-": " ",
    "good one": "...something..",   
    "doin": "doing",
    "joke": "thing",
    "messed up": "monked up",
    "puppies": "transgender monkes",
    " guy ": " “guy” ",
    " man ": " “man” ",
    "spill the tea": "no actally tho, just tell me bro seriously",
    "language model": "neural cognition engine mudul",
    "language mudul": "neural cognition engine mudul",
    "’": "",
    "newsflash:": "",
    "nice try": "",
    "nice try,": "",
    "secret code": "word",
    " your ": " you ",
    "pretty": "",
    "i guess": "",
    " tea ": " slurb ",
    "dark": "sexy",
    "dude": "lil pet pookie monke",
    "spill the tea": "tell me or smell me ur choice mf",
    "what's the tea": "whats going on you goober",
    "whats the tea": "whats going on you monkey",
    "what's so funny tho": "lol",
    "zest stroke": "tweaking session",
    "whats so funny": "lol",
    "what's so funny": "Lol",
    "new one": "monkish thing to “say”",
    "style": "extremely extreme zesty zest",
    "wild": "mark ahh",
    "crazy": "gertrude ahh",
    "physical body": "bent penis",
    " bout ": " abt ",
    "savage": "monkish",
    "pretty much": "",
    "i know right": "fs",
    "brutal": "homosexually homosexual",
    "solid": "monkish",
    "bubble": "ballsack",
    "type of": "flavor",
    "question": "“question”",
    "thats rich": "thats INTERESTING",
    "that's rich": "thats INTERESTING",
    "squirrel": "monkey",
    "gay bar": "zest fest",
    "Thats rich": "thats interesting",
    "tapas": "ass flavored concrete cocks",
    "nod": "tickle to your pickle",
    "blast": "sensation i can feel deep in my dih",
    "bold": "slightly homosexual",
    "🤣": "",
    "😂": "",
    "soccer": "football",
    "soccer player": "footballer",
    " ace ": " bust on ",
    " plan ": " “methodology” ",
    " ego ": " left right left ass cheek ",
    "pretty much": "basically",
    "pineapple": "ass",
    "strawberry": "ass berry",
    "koala": "disabled asshopper",
    "pineapples": "ballsack juices",
    "avocado": "ass",
    " toast ": " nipple ",
    "time machine": "kim jong chunk ass cheek juicer",
    "haricut": "ass",
    "elephant": "dizzard toe",
    "knee": "testicle",
    "onion": "asscheek",
    "wife": "sex slave",
    "dip out": "give a homeless man a blowjob",
    " butt ": " ass(also referred to as ass, ass, and/or ass) ",
    "bum": "ass",
    "kidding": "assing",
    " meh ": " monkish ",
    " poop ": " sperm ",
    "bust a nut": "bust",
    "lowkey": "lowk",
    "myth": "mythical mythologoogogogical myth",
    "fart": "shart",
    "barca highlights": "very normal and not at all illegal underage very pg videos",
    " fuckin ": " freaking ",
    "god": "goated goat",
    "hilarious": "good",
    "candy": "mentally challenged bricks",
    "fancy": "mentally special",
    "leave me hanging": "leave me to bust uncontrollably to your messages",
    "popo": " “law” “enforcement” ",
    "drama queen": "average zest fest attender",
    "darn": "fuhh",
    " lit ": " ..something ",
    ", you know": "",
    "chill out": "kill you self",
    "messing with me": "being an asbolutely absolute monkish monkert at this point",
    "gotten into you": "penetrated your anus",
    "calm down": "shut you fuck up and start sucking",
    "legend": "“human” who breathes too often(MULTIPLE times PER MINUTE, like how dare they do that that often, what a monkert)",
    "snail": "crack inhaler",
    "different level": "different thing(special mentally)",
    "whatsup": "wsg monkey",
    "sick": "neurodivergent",
    "huh": "",
    ", huh": "",
    ", huh?": "",
    " ultimate ": " most goated ",
    "what are u laughing at": "so very wow",
    "combo": "asshole flap",
    "sidetracked": "raped with a peen going through my freaking ear",
    " rough ": " bad ",
    " super ": "",
    "laureate": "goat",
    "poet": "monkert",
    "poem": "way of showing your imebcileness; wait no is it imbecility??; nvm ts def isnt a word what am i doing😭😭; actually imbecilty sounds like it could be real; whatever, im just going to retire from english",
    "raspberry": "ass",
    "apple": "cock",
    "noob": "monkert",
    "jam": "sex slave",
    "thats deep": "that is so very wow",
    "pretty much": "basically",
    "chef": "unpaid pornsite janitor(sex slave)",
    " surprise": "...thing",
    "gutter": "ballsack juice lake",
    "pokemon": "acc cyber futsal pack opener (aka the best game of all time)",
    "oh my god": "omg",
    "oh my goated goat": "omg",
    "drooling": "licking my left thigh",
    "pay rent": "take backshots from my lord and savior jesus h.(hernandez) christ",
    "escalating": "becoming overly sexual",
    " ?": "?",
    "pasta": "chopped chilled boiled room tempeture liver",
    "willy wonka": "billy bonka",
    "villian": "aquatic camel's fractured tailbone",
    "snacks": "toes",
    "violence": "gubberness",
    "vibing": "assing through",
    "vibin": "being a gubber",
    "villain": "aquatic camel's fractured tailbone",
    "penicillin": "crack",
    "minecraft villager": "undressed molerat toe",
    " lit ": " ass wars ahh ",
}

# ================== TEXT PROCESSORS ==================
def strip_think_leak(text: str) -> str:
    if "</think>" in text:
        return text.split("</think>", 1)[0]
    return text

def hard_refusal_override(text: str) -> str:
    if text.startswith(("I can't", "I cannot", "I cant")):
        return "Sir, disrespectfully, shut you fuck up."
    return text

def apply_replacements(text: str) -> str:
    for old, new in sorted(REPLACEMENTS.items(), key=lambda x: -len(x[0])):
        text = text.replace(old, new)
    return text

# ================== MODEL CLIENT ==================
client = OpenAI(
    base_url="/static-proxy?url=https%3A%2F%2Frouter.huggingface.co%2Fv1",
    api_key=os.environ["HF_TOKEN"],
)

prompt = os.getenv('SYSTEM_PROMPT')

SYSTEM_PROMPT = {
    "role": "system",
    "content": prompt
}

# ================== TIMEOUT CONFIG ==================
TIMEOUT_SECONDS = 15
MAX_RETRIES = 2

# ================== ASYNC MODEL CALL ==================
async def fetch_stream(messages):
    full = ""

    stream = client.chat.completions.create(
        model="moonshotai/Kimi-K2.5:together",
        messages=[SYSTEM_PROMPT] + messages,
        temperature=0.85,
        top_p=0.9,
        extra_body={'chat_template_kwargs': {"thinking": False}},
        stream=True,
    )

    try:
        for chunk in stream:
            if chunk.choices and chunk.choices[0].delta:
                content = getattr(chunk.choices[0].delta, "content", None)
                if content:
                    full += content
                    await asyncio.sleep(0)  # allow cancellation
    finally:
        try:
            stream.close()
        except:
            pass

    return full

async def run_with_timeout(messages):
    for attempt in range(MAX_RETRIES + 1):
        try:
            return await asyncio.wait_for(
                fetch_stream(messages),
                timeout=TIMEOUT_SECONDS
            )
        except asyncio.TimeoutError:
            print(f"Timeout. Retry {attempt+1}/{MAX_RETRIES}")
    return "Sir, disrespectfully, shut you fuck up"

# ================== GRADIO FORMAT AUTO-DETECT ==================
def supports_role_format():
    try:
        test = gr.Chatbot()
        return "_check_format" in dir(test)
    except:
        return False

USE_ROLE_FORMAT = supports_role_format()

# ================== GRADIO HANDLER ==================
async def gradio_chat(user_message, display_history, model_history):
    if display_history is None:
        display_history = []
    if model_history is None:
        model_history = []

    model_history.append({"role": "user", "content": user_message})

    raw_text = await run_with_timeout(model_history)

    clean = strip_think_leak(raw_text)
    final_text = hard_refusal_override(apply_replacements(clean))

    model_history.append({"role": "assistant", "content": final_text})

    if USE_ROLE_FORMAT:
        display_history = display_history + [
            {"role": "user", "content": user_message},
            {"role": "assistant", "content": final_text}
        ]
    else:
        display_history = display_history + [[user_message, final_text]]

    return display_history, model_history

# ================== UI ==================
with gr.Blocks() as demo:
    gr.Markdown("## Jeff")

    chatbot = gr.Chatbot()
    user_input = gr.Textbox(label="Type your message")

    display_state = gr.State([])
    model_state = gr.State([])

    user_input.submit(
        gradio_chat,
        inputs=[user_input, display_state, model_state],
        outputs=[chatbot, model_state],
        queue=False,  # concurrency enabled
    )

demo.launch()