import gradio as gr
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from peft import PeftModel

base_model_name = "facebook/nllb-200-distilled-600M"
adapter_model_name = "entropy25/mt_en_no_oil"

tokenizer = AutoTokenizer.from_pretrained(base_model_name)
base_model = AutoModelForSeq2SeqLM.from_pretrained(
    base_model_name,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    device_map="auto"
)

model = PeftModel.from_pretrained(base_model, adapter_model_name)

def translate(text, source_lang, target_lang):
    if not text.strip():
        return ""
    
    if source_lang == target_lang:
        return text
    
    lang_map = {
        "English": "eng_Latn",
        "Norwegian": "nob_Latn"
    }
    
    sentences = text.split('\n')
    translated_sentences = []
    
    for sentence in sentences:
        if not sentence.strip():
            translated_sentences.append("")
            continue
        
        inputs = tokenizer(
            sentence,
            return_tensors="pt",
            truncation=True,
            max_length=512
        )
        
        if hasattr(model, 'device'):
            inputs = {k: v.to(model.device) for k, v in inputs.items()}
        
        outputs = model.generate(
            **inputs,
            forced_bos_token_id=tokenizer.convert_tokens_to_ids(lang_map[target_lang]),
            max_length=512,
            num_beams=5
        )
        
        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
        translated_sentences.append(result)
    
    return '\n'.join(translated_sentences)

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:
        with open(file.name, 'r', encoding='utf-8') as f:
            return f.read()
    except:
        try:
            with open(file.name, 'r', encoding='latin-1') as f:
                return f.read()
        except Exception as e:
            return f"Error reading file: {str(e)}"

custom_css = """
.gradio-container {
    max-width: 1100px !important;
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
}
.main-container {
    background: #f6f7f8 !important;
    padding: 0 !important;
    border-radius: 0 !important;
}
.translate-box {
    background: white !important;
    border-radius: 5px !important;
    padding: 0 !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;
}
.lang-selector {
    border: none !important;
    background: transparent !important;
    font-size: 15px !important;
    font-weight: 500 !important;
    color: #333 !important;
}
.text-area textarea {
    border: none !important;
    font-size: 17px !important;
    line-height: 1.7 !important;
    padding: 20px !important;
    min-height: 200px !important;
}
.swap-container {
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    padding: 20px 0 !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;
    cursor: pointer !important;
}
.swap-btn:hover {
    background: #f8f9fa !important;
    border-color: #0f6fff !important;
}
.footer-info {
    text-align: center !important;
    color: #999 !important;
    font-size: 13px !important;
    padding: 20px !important;
}
"""

EXAMPLES = {
    "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": "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": "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": "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 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."
}

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,
                        elem_classes="lang-selector",
                        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):
            with gr.Row(elem_classes="swap-container"):
                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,
                        elem_classes="lang-selector",
                        scale=1
                    )
                
                output_text = gr.Textbox(
                    placeholder="Translation",
                    show_label=False,
                    lines=8,
                    max_lines=20,
                    container=False,
                    elem_classes="text-area",
                    interactive=False
                )
    
    gr.HTML(
        "<div class='footer-info'>"
        "Oil & Gas Translation • English ↔ Norwegian"
        "</div>"
    )
    
    with gr.Accordion("Example Sentences", open=True):
        with gr.Row():
            example_text = gr.Textbox(
                value=EXAMPLES["drilling_short"],
                label="",
                lines=3,
                max_lines=5,
                show_copy_button=True
            )
            use_example_btn = gr.Button("Use This Example →", variant="primary", size="sm")
        
        with gr.Row():
            gr.Button("Drilling (Short)", size="sm").click(
                lambda: EXAMPLES["drilling_short"],
                outputs=example_text
            )
            gr.Button("Drilling (Long)", size="sm").click(
                lambda: EXAMPLES["drilling_long"],
                outputs=example_text
            )
            gr.Button("Reservoir (Short)", size="sm").click(
                lambda: EXAMPLES["reservoir_short"],
                outputs=example_text
            )
            gr.Button("Reservoir (Long)", size="sm").click(
                lambda: EXAMPLES["reservoir_long"],
                outputs=example_text
            )
            gr.Button("Subsea (Short)", size="sm").click(
                lambda: EXAMPLES["subsea_short"],
                outputs=example_text
            )
        
        with gr.Row():
            gr.Button("Subsea (Long)", size="sm").click(
                lambda: EXAMPLES["subsea_long"],
                outputs=example_text
            )
            gr.Button("Seismic (Short)", size="sm").click(
                lambda: EXAMPLES["seismic_short"],
                outputs=example_text
            )
            gr.Button("Seismic (Long)", size="sm").click(
                lambda: EXAMPLES["seismic_long"],
                outputs=example_text
            )
            gr.Button("Safety (Short)", size="sm").click(
                lambda: EXAMPLES["safety_short"],
                outputs=example_text
            )
            gr.Button("Safety (Long)", size="sm").click(
                lambda: EXAMPLES["safety_long"],
                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 to translate",
            file_types=[".txt"],
            type="filepath"
        )
    
    input_text.change(
        fn=translate,
        inputs=[input_text, source_lang, target_lang],
        outputs=output_text
    )
    
    source_lang.change(
        fn=translate,
        inputs=[input_text, source_lang, target_lang],
        outputs=output_text
    )
    
    target_lang.change(
        fn=translate,
        inputs=[input_text, source_lang, target_lang],
        outputs=output_text
    )
    
    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
    )

demo.launch()