import streamlit as st
from unsloth import FastLanguageModel
import torch
max_seq_length = 2048 # Choose any! We auto support RoPE Scaling internally!
dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
load_in_4bit = True



model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "Mudditha/test-phi-3", # YOUR MODEL YOU USED FOR TRAINING
    max_seq_length = max_seq_length,
    dtype = dtype,
    load_in_4bit = load_in_4bit,
)


FastLanguageModel.for_inference(model)


inf_prompt_1 = """You are a creative domain name generator with a sound knowledge about how the creative name are formed according to requirement and language use.
You need to create and output 10 different unique names, which are not in your knowledge base, based on your knowledge on considered facts when generating domain names.
Don't use same keyword for every name. Instead use different attractive similar word combinations. make the names short as possible
output should be in following format. 
don't include '.com' in outputs and give different names.

"" 1. ......
2. .......
.......
.......
10. ....... ""

### Input:
{}

### Response:
{}"""



def print_output(input):

    inputs = tokenizer(
    [
        inf_prompt_1.format(
            input, # instruction
            "", # output - leave this blank for generation!
        ),
    ], return_tensors = "pt").to('cuda')
    
    # outputs = model.generate(**inputs, max_new_tokens = 64, use_cache = True)
    # tokenizer.batch_decode(outputs)
    
    # outputs = model.generate(**inputs, max_new_tokens = 100, use_cache = True)
    # return tokenizer.batch_decode(outputs)


    outputs = model.generate(**inputs, max_new_tokens = 200, use_cache = True, temperature=0.5)

    n1 = tokenizer.batch_decode(outputs)[0].index('### Response:')
    output_ = tokenizer.batch_decode(outputs)[0][n1+13:].replace('.com','')
    n2 = output_.index('\n\n')

    dom_names = []
    
    for name in output_[:n2].split('\n'):
      if name != '':
          dom_names.append(name.split('.')[1])
    
    # s = '\n'.join(dom_names)
    return dom_names


# Streamlit app UI
st.title("AI Based Domain Names Suggestion")

# Text box for user input
user_input = st.text_input("Type your idea here:")

# # Button to submit input
# if st.button("Submit"):
#     # Reprint the user input
#     for item in print_output(user_input):
#         st.write(item)
col1, col2 = st.columns(2)

# Loop through the output list and distribute items between the two columns
if st.button("Submit"):
    output_list = print_output(user_input)
    
    for i, item in enumerate(output_list):
        if i % 2 == 0:
            col1.write(item)  # Even-indexed items in column 1
        else:
            col2.write(item)  # Odd-indexed items in column 2

# def print_output(input):
#     inputs = tokenizer(
#     [
#         inf_prompt_1.format(
#             input, # instruction
#             "", # output - leave this blank for generation!
#         ),
#     ], return_tensors = "pt").to("cuda")
    
#     # outputs = model.generate(**inputs, max_new_tokens = 64, use_cache = True)
#     # tokenizer.batch_decode(outputs)
    
#     outputs = model.generate(**inputs, max_new_tokens = 200, use_cache = True, temperature=0.5)

#     n1 = tokenizer.batch_decode(outputs)[0].index('### Response:')
#     output_ = tokenizer.batch_decode(outputs)[0][n1+13:].replace('.com','')
#     n2 = output_.index('\n\n')

#     dom_names = []
    
#     for name in output_[:n2].split('\n'):
#       if name != '':
#           dom_names.append(name.split('.')[1])
    
#     s = '\n'.join(dom_names)

#     return tokenizer.batch_decode(s)


# # # Streamlit app UI
# # st.title("Echo Input Example")

# # # Text box for user input
# # user_input = st.text_input("Type something here:")

# # # Button to submit input
# # if st.button("Submit"):
# #     # Reprint the user input
# #     st.write(f"You typed: {print_output(user_input)}")


# def main():

#     css_dark_mode = """
#     <style>
#         body {
#             background-color: #121212;
#             color: #f8f9fa;
#             font-family: Arial, sans-serif;
#         }
#         .header {
#             background: linear-gradient(90deg, #005c97, #363795);
#             padding: 15px;
#             margin-bottom: 15px;
#             border-radius: 10px;
#             box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.2);
#             text-align: center;
#         }
#         .header h1 {
#             color: #ffffff;
#             font-size: 36px;
#         }
#         .header h2 {
#             color: #b0b0b0;
#             font-size: 22px;
#             font-family: Georgia, serif;
#         }
#         .stButton > button {
#             background-color: #20c997;
#             color: #ffffff;
#             border: none;
#             border-radius: 8px;
#             padding: 10px 20px;
#             font-size: 16px;
#             cursor: pointer;
#             transition: background-color 0.3s;
#         }
#         .stButton > button:hover {
#             background-color: #17a589;
#         }
#         .success-box {
#             background-color: #A1D6CB;
#             padding: 10px;
#             margin: 10px 0;
#             border-radius: 8px;
#             box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.2);
#         }
#     </style>
#     """
#     # Inject the custom CSS
#     st.markdown(css_dark_mode, unsafe_allow_html=True)

#     # Header Section
#     html_temp = """
#     <div class="header">
#         <h1>Dominious</h1>
#         <h2>AI Based Domain Name Suggestion System</h2>
#     </div>
#     """
#     st.markdown(html_temp, unsafe_allow_html=True)
#     # st.title("Domain Name Suggestion System")

#     # Get user input
#     user_input = st.text_input("Describe your business here...")

#     # Generate the response
#     if st.button("Generate"):
#         with st.spinner("Generating Domain Names..."):
#             result = print_output(user_input)
# if __name__ == "__main__":
#     main()