import tensorflow as tf  # type: ignore
from keras_preprocessing.text import Tokenizer  # type: ignore
from keras_preprocessing.sequence import pad_sequences  # type: ignore
import numpy as np  # type: ignore
import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

# Optional: suppress ABSEIL warnings if TensorFlow still leaks something
try:
    import absl.logging  # type: ignore
    absl.logging.set_verbosity(absl.logging.ERROR)
except ImportError:
    pass


def load_dialogues(datapath="C:/Users/90372/ai/data/", num_words=10000, max_len=40): # noqa
    lines_file = os.path.join(datapath, "movie_lines.txt")
    convos_file = os.path.join(datapath, "movie_conversations.txt")

    id2line = {}
    with open(lines_file, encoding="utf-8", errors='ignore') as f:
        for line in f:
            parts = line.strip().split(" +++$+++ ")
            if len(parts) == 5:
                line_id, text = parts[0], parts[4]
                id2line[line_id] = text

    inputs, responses = [], []
    with open(convos_file, encoding="utf-8", errors='ignore') as f:
        for line in f:
            parts = line.strip().split(" +++$+++ ")
            if len(parts) == 4:
                line_ids = eval(parts[3])
                for i in range(len(line_ids) - 1):
                    input_line = id2line.get(line_ids[i], "")
                    response_line = id2line.get(line_ids[i + 1], "")
                    if input_line and response_line:
                        inputs.append(input_line)
                        responses.append(response_line)

    catchAllError = "I'm sorry, I cant seem to remember anything right now."
    tokenizer = Tokenizer(num_words=num_words, oov_token=catchAllError)
    tokenizer.fit_on_texts(inputs + responses)

    input_sequences = tokenizer.texts_to_sequences(inputs)
    response_sequences = tokenizer.texts_to_sequences(responses)

    inputs_padded = pad_sequences(input_sequences, maxlen=max_len, padding="post", truncating="post") # noqa
    responses_padded = pad_sequences(response_sequences, maxlen=max_len, padding="post", truncating="post") # noqa

    return inputs_padded, responses_padded, tokenizer

def build_seq2seq_model(vocab_size, max_len):
    inputs = tf.keras.Input(shape=(max_len,))
    embedding = tf.keras.layers.Embedding(vocab_size, 256)(inputs)
    lstm = tf.keras.layers.LSTM(256, return_sequences=True)(embedding)
    attention = tf.keras.layers.Attention()([lstm, lstm])
    decoder_lstm = tf.keras.layers.LSTM(256)(attention)
    outputs = tf.keras.layers.Dense(vocab_size)(decoder_lstm)  # Final output layer

    model = tf.keras.Model(inputs, outputs)
    loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)  # Removed label_smoothing
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4), loss=loss)  # Compile model
    return model

def train_conversational_model():
    inputs, responses, tokenizer = load_dialogues()

    target_sequences = responses[:, 1:]  # Shift the target by one timestep
    target_sequences = np.pad(target_sequences, ((0, 0), (0, 1)), mode='constant')  # Pad the target sequences

    model = build_seq2seq_model(vocab_size=10000, max_len=inputs.shape[1])
    model.fit(inputs, target_sequences, epochs=5, batch_size=32)  # Train with shifted target sequences

    return model, tokenizer


def generate_response(model, tokenizer, prompt, max_len=40):
    sequence = tokenizer.texts_to_sequences([prompt])
    sequence = pad_sequences(sequence, maxlen=max_len)
    pred = model.predict(sequence)
    next_word_id = np.argmax(pred[0])
    return tokenizer.index_word.get(next_word_id, "<UNK>")


if __name__ == "__main__":
    model, tokenizer = train_conversational_model()
    while True:
        prompt = input("User: ")
        if prompt.lower() == "exit":
            break
        print("AI:", generate_response(model, tokenizer, prompt))