import gradio as gr
import numpy as np
import tensorflow as tf
import pickle
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.sequence import pad_sequences
from huggingface_hub import hf_hub_download

from fastapi import FastAPI, Request
import uvicorn
import threading

app = FastAPI()

@app.post("/reload")
async def reload_model(request: Request):
    token = request.headers.get("x-hub-token")
    if token != "mi-secret-token":
        print("🚫 Unauthorized webhook request.")
        return {"status": "unauthorized"}

    load_latest_model()
    print("🔁 Model reloaded securely via webhook.")
    return {"status": "ok"}


np.random.seed(42)
tf.random.set_seed(42)

# Load tokenizers
en_tok_path = hf_hub_download(repo_id="Juna190825/github_jeffprosise_model", filename="tokenizers/en_tokenizer.pkl")
fr_tok_path = hf_hub_download(repo_id="Juna190825/github_jeffprosise_model", filename="tokenizers/fr_tokenizer.pkl")

with open(en_tok_path, "rb") as f:
    en_tokenizer = pickle.load(f)
with open(fr_tok_path, "rb") as f:
    fr_tokenizer = pickle.load(f)

fr_index_lookup = fr_tokenizer.index_word
sequence_len = 20  # adjust as needed

# Load model function
def load_latest_model():
    global model
    model_path = hf_hub_download(
        repo_id="Juna190825/github_jeffprosise_model", filename="model.keras", cache_dir="cache", force_download=True
    )
    en_tok_path = hf_hub_download(repo_id="Juna190825/github_jeffprosise_model", filename="tokenizers/en_tokenizer.pkl")
    fr_tok_path = hf_hub_download(repo_id="Juna190825/github_jeffprosise_model", filename="tokenizers/fr_tokenizer.pkl")
    model = load_model(model_path)
    return "✅ Model reloaded!"

# Initial load
load_latest_model()

# Translation function
def translate_text(text):
    input_sequence = en_tokenizer.texts_to_sequences([text])
    padded_input_sequence = pad_sequences(input_sequence, maxlen=sequence_len, padding='post')
    decoded_text = '[start]'

    for i in range(sequence_len):
        target_sequence = fr_tokenizer.texts_to_sequences([decoded_text])
        padded_target_sequence = pad_sequences(target_sequence, maxlen=sequence_len, padding='post')[:, :-1]

        prediction = model([padded_input_sequence, padded_target_sequence])
        idx = np.argmax(prediction[0, i, :])
        token = fr_tokenizer.index_word.get(idx)

        if token is None or token.strip() == '' or token == '[unk]':
            return f"Translation unavailable for '{text}'"
        decoded_text += ' ' + token
        if token == '[end]':
            break

    return decoded_text.replace('[start] ', '').replace(' [end]', '')

# Gradio Interface using Blocks
with gr.Blocks() as demo:
    gr.Markdown("# 🌍 English to French Translator")

    with gr.Row():
        input_text = gr.Textbox(label="English", placeholder="Enter text here…")
        output_text = gr.Textbox(label="French Translation")

    with gr.Row():
        translate_btn = gr.Button("Translate")
        reload_btn = gr.Button("🔄 Reload Model")

    translate_btn.click(translate_text, inputs=input_text, outputs=output_text)
    reload_btn.click(load_latest_model, outputs=output_text)

    gr.Examples(["What is your name?", "I love chocolate.", "Where is the nearest station?"], inputs=input_text)


def run_api():
    uvicorn.run(app, host="0.0.0.0", port=7862)


threading.Thread(target=run_api, daemon=True).start()


demo.launch(share=True)