import gradio as gr
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
from functools import lru_cache
import os
import time
import threading
import nltk

try:
    nltk.data.find("tokenizers/punkt")
except LookupError:
    nltk.download("punkt")
    try:
        nltk.download("punkt_tab")
    except:
        pass

from nltk.tokenize import sent_tokenize

BASE_MODEL = os.getenv("BASE_MODEL", "facebook/nllb-200-distilled-600M")
ADAPTER_EN_TO_NO = os.getenv("ADAPTER_EN_TO_NO", "entropy25/mt_en_no_oil")
ADAPTER_NO_TO_EN = os.getenv("ADAPTER_NO_TO_EN", "entropy25/mt_no_en_oil")

MAX_FILE_SIZE = 1024 * 1024
MAX_TEXT_LENGTH = 10000
BATCH_SIZE = 10
NUM_BEAMS = 3
MAX_LENGTH = 256

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

device = "cuda" if torch.cuda.is_available() else "cpu"

base_model = AutoModelForSeq2SeqLM.from_pretrained(
    BASE_MODEL,
    quantization_config=BitsAndBytesConfig(load_in_8bit=True) if device == "cuda" else None,
    device_map={"": 0} if device == "cuda" else None,
    low_cpu_mem_usage=True,
)

model = PeftModel.from_pretrained(base_model, ADAPTER_EN_TO_NO, adapter_name="en_to_no")
model.load_adapter(ADAPTER_NO_TO_EN, adapter_name="no_to_en")
model.eval()

adapter_lock = threading.Lock()

LANG_CONFIG = {
    ("English", "Norwegian"): ("en_to_no", "nob_Latn"),
    ("Norwegian", "English"): ("no_to_en", "eng_Latn"),
}

EXAMPLES_EN = {
    "drilling_short": "Mud weight adjusted to 1.82 specific gravity at 3,247 meters depth.",
    "drilling_long": "The drilling operation at well site A-15 encountered unexpected high-pressure zones at 3,247 meters depth, requiring immediate adjustment of mud weight from 1.65 to 1.82 specific gravity to maintain wellbore stability and prevent potential kicks.",
    "reservoir_short": "Permeability is 250 millidarcy with 22 percent porosity.",
    "reservoir_long": "The reservoir shows excellent permeability of 250 millidarcy and porosity of 22 percent based on core analysis, indicating significant hydrocarbon potential with estimated oil saturation of 65 percent.",
    "subsea_short": "Christmas tree rated for 10,000 psi working pressure.",
    "subsea_long": "The subsea production system consists of a vertical Christmas tree rated for 10,000 psi working pressure and 150 degrees Celsius temperature, equipped with redundant safety features including automatic shutdown valves and real-time pressure monitoring systems.",
    "seismic_short": "Structural trap area estimated at 12 square kilometers.",
    "seismic_long": "Seismic data confirms the presence of a structural trap with an estimated area of 12 square kilometers, and productivity tests show stabilized oil production of 3,400 barrels per day at optimization pressure of 2,100 psi.",
    "safety_short": "H2S training required before site access.",
    "safety_long": "Emergency response procedures require all personnel to complete H2S safety training before site access, with breathing apparatus and wind indicators positioned at designated muster points, and immediate evacuation protocols activated when gas detection exceeds 10 ppm concentration levels.",
}

EXAMPLES_NO = {
    "drilling_short": "Slamvekt justert til 1,82 spesifikk tyngde ved 3 247 meters dybde.",
    "drilling_long": "Boreoperasjonen ved brønnsted A-15 støtte på uventede høytrykksoner ved 3 247 meters dybde, noe som krevde umiddelbar justering av slamvekt fra 1,65 til 1,82 spesifikk tyngde for å opprettholde brønnborestabilitet og forhindre potensielle kicks.",
    "reservoir_short": "Permeabilitet er 250 millidarcy med 22 prosent porøsitet.",
    "reservoir_long": "Reservoaret viser utmerket permeabilitet på 250 millidarcy og porøsitet på 22 prosent basert på kjerneanalyse, noe som indikerer betydelig hydrokarbonpotensial med estimert oljemetning på 65 prosent.",
    "subsea_short": "Juletre dimensjonert for 10 000 psi arbeidstrykk.",
    "subsea_long": "Subsea produksjonssystemet består av et vertikalt juletre dimensjonert for 10 000 psi arbeidstrykk og 150 grader Celsius temperatur, utstyrt med redundante sikkerhetsfunksjoner inkludert automatiske nedstengningsventiler og sanntids trykkmonitorering.",
    "seismic_short": "Strukturell felle estimert til 12 kvadratkilometer.",
    "seismic_long": "Seismiske data bekrefter tilstedeværelsen av en strukturell felle med estimert areal på 12 kvadratkilometer, og produktivitetstester viser stabilisert oljeproduksjon på 3 400 fat per dag ved optimaliseringstrykk på 2 100 psi.",
    "safety_short": "H2S-opplæring påkrevd før tilgang til området.",
    "safety_long": "Nødprosedyrer krever at alt personell fullfører H2S-sikkerhetsopplæring før områdetilgang, med åndedrettsutstyr og vindindikatorer plassert ved utpekte samlingspunkter, og umiddelbare evakueringsprotokoller aktiveres når gassdeteksjon overskrider 10 ppm konsentrasjonsnivå.",
}


@lru_cache(maxsize=512)
def cached_sent_tokenize(text):
    return tuple(sent_tokenize(text))


@lru_cache(maxsize=512)
def translate_cached(text, source_lang, target_lang):
    if source_lang == target_lang or not text.strip():
        return text, 0.0

    config = LANG_CONFIG.get((source_lang, target_lang))
    if not config:
        return "Unsupported language pair", 0.0

    adapter_name, tgt_code = config
    start = time.time()
    dev = next(model.parameters()).device
    translated_paragraphs = []

    for paragraph in text.split("\n"):
        if not paragraph.strip():
            translated_paragraphs.append("")
            continue

        sentences = cached_sent_tokenize(paragraph)
        results = []

        for i in range(0, len(sentences), BATCH_SIZE):
            batch = sentences[i : i + BATCH_SIZE]
            inputs = tokenizer(
                batch,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=MAX_LENGTH,
            )
            inputs = {k: v.to(dev) for k, v in inputs.items()}

            with adapter_lock:
                model.set_adapter(adapter_name)
                with torch.inference_mode():
                    outputs = model.generate(
                        **inputs,
                        forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_code),
                        max_length=MAX_LENGTH,
                        num_beams=NUM_BEAMS,
                        early_stopping=True,
                    )

            results.extend(tokenizer.batch_decode(outputs, skip_special_tokens=True))

        translated_paragraphs.append(" ".join(results))

    return "\n".join(translated_paragraphs), time.time() - start


def translate(text, source_lang, target_lang):
    if not text.strip():
        return "", ""
    if len(text) > MAX_TEXT_LENGTH:
        return f"Error: Text too long (max {MAX_TEXT_LENGTH:,} characters)", ""

    try:
        result, elapsed = translate_cached(text, source_lang, target_lang)
        return result, f"Completed in {elapsed:.2f}s"
    except Exception as e:
        return f"Translation error: {str(e)}", ""


def swap_languages(src, tgt, input_txt, output_txt):
    return tgt, src, output_txt, input_txt


def load_file(file):
    if file is None:
        return "", ""
    try:
        size = os.path.getsize(file.name)
        if size > MAX_FILE_SIZE:
            return "Error: File too large (max 1MB)", ""
        for encoding in ("utf-8", "latin-1"):
            try:
                with open(file.name, "r", encoding=encoding) as f:
                    content = f.read()
                if len(content) > MAX_TEXT_LENGTH:
                    return f"Error: File content too long (max {MAX_TEXT_LENGTH:,} characters)", ""
                return content, ""
            except UnicodeDecodeError:
                continue
    except (IOError, OSError) as e:
        return f"Error reading file: {str(e)}", ""


def get_example(key, source_lang):
    examples = EXAMPLES_EN if source_lang == "English" else EXAMPLES_NO
    return examples[key]


custom_css = """
.gradio-container {
    max-width: 1100px !important;
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
}
.translate-box {
    background: white !important;
    border-radius: 5px !important;
    box-shadow: 0 2px 4px rgba(0,0,0,0.08) !important;
    margin: 20px 0 !important;
}
.lang-header {
    padding: 16px 20px !important;
    border-bottom: 1px solid #e8eaed !important;
    background: #fafafa !important;
}
.text-area textarea {
    border: none !important;
    font-size: 17px !important;
    line-height: 1.7 !important;
    padding: 20px !important;
    min-height: 200px !important;
}
.swap-btn {
    width: 44px !important;
    height: 44px !important;
    min-width: 44px !important;
    border-radius: 50% !important;
    background: white !important;
    border: 1px solid #d1d5db !important;
    box-shadow: 0 1px 3px rgba(0,0,0,0.1) !important;
    font-size: 18px !important;
    color: #0f6fff !important;
}
.translate-btn {
    background: #ff8c00 !important;
    color: white !important;
    border: none !important;
    padding: 12px 24px !important;
    font-size: 15px !important;
    font-weight: 500 !important;
    border-radius: 4px !important;
}
.time-info {
    text-align: center !important;
    color: #666 !important;
    font-size: 13px !important;
    padding: 10px !important;
    font-style: italic !important;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Default()) as demo:

    gr.HTML("<div style='height: 20px'></div>")

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Group(elem_classes="translate-box"):
                with gr.Row(elem_classes="lang-header"):
                    source_lang = gr.Dropdown(
                        choices=["English", "Norwegian"],
                        value="English",
                        show_label=False,
                        container=False,
                        scale=1,
                    )
                input_text = gr.Textbox(
                    placeholder="Type to translate",
                    show_label=False,
                    lines=8,
                    max_lines=20,
                    container=False,
                    elem_classes="text-area",
                )

        with gr.Column(scale=0, min_width=100):
            swap_btn = gr.Button("⇄", elem_classes="swap-btn")

        with gr.Column(scale=1):
            with gr.Group(elem_classes="translate-box"):
                with gr.Row(elem_classes="lang-header"):
                    target_lang = gr.Dropdown(
                        choices=["English", "Norwegian"],
                        value="Norwegian",
                        show_label=False,
                        container=False,
                        scale=1,
                    )
                output_text = gr.Textbox(
                    placeholder="Translation",
                    show_label=False,
                    lines=8,
                    max_lines=20,
                    container=False,
                    elem_classes="text-area",
                    interactive=False,
                )

    with gr.Row():
        translate_btn = gr.Button("Translate", variant="primary", elem_classes="translate-btn", size="lg")

    with gr.Row():
        time_display = gr.Textbox(show_label=False, container=False, interactive=False, elem_classes="time-info")

    gr.HTML("<div style='text-align:center;color:#999;font-size:13px;padding:20px'>Oil & Gas Translation — English ↔ Norwegian</div>")

    with gr.Accordion("Example Sentences", open=True):
        with gr.Row():
            example_text = gr.Textbox(value=EXAMPLES_EN["drilling_short"], label="", lines=3, show_copy_button=True)
            use_example_btn = gr.Button("Use This Example", variant="primary", size="sm")

        with gr.Row():
            for key, label in [
                ("drilling_short", "Drilling Short"),
                ("drilling_long", "Drilling Long"),
                ("reservoir_short", "Reservoir Short"),
                ("reservoir_long", "Reservoir Long"),
                ("subsea_short", "Subsea Short"),
            ]:
                gr.Button(label, size="sm").click(
                    lambda sl, k=key: get_example(k, sl), inputs=[source_lang], outputs=example_text
                )

        with gr.Row():
            for key, label in [
                ("subsea_long", "Subsea Long"),
                ("seismic_short", "Seismic Short"),
                ("seismic_long", "Seismic Long"),
                ("safety_short", "Safety Short"),
                ("safety_long", "Safety Long"),
            ]:
                gr.Button(label, size="sm").click(
                    lambda sl, k=key: get_example(k, sl), inputs=[source_lang], outputs=example_text
                )

        use_example_btn.click(fn=lambda x: x, inputs=example_text, outputs=input_text)

    with gr.Accordion("Upload Text File", open=False):
        file_input = gr.File(label="Upload a .txt file (max 1MB)", file_types=[".txt"], type="filepath")

    translate_btn.click(
        fn=translate,
        inputs=[input_text, source_lang, target_lang],
        outputs=[output_text, time_display],
    )

    swap_btn.click(
        fn=swap_languages,
        inputs=[source_lang, target_lang, input_text, output_text],
        outputs=[source_lang, target_lang, input_text, output_text],
    )

    file_input.change(fn=load_file, inputs=file_input, outputs=[input_text, time_display])

    source_lang.change(fn=lambda sl: get_example("drilling_short", sl), inputs=[source_lang], outputs=[example_text])

demo.queue().launch()