import os
import sys
import site
import torch
import transformers
import gradio as gr
import re
import json
from datetime import datetime

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

print("=" * 80)
print("Vidyaapati Translator")
print("=" * 80)

print("Python :", sys.version)
print("Transformers :", transformers.__version__)

# -----------------------------------------------------------------------------
# Fix IndicTransToolkit import
# -----------------------------------------------------------------------------

def fix_collator():
    try:
        collator_path = None

        for p in site.getsitepackages():
            test = os.path.join(
                p,
                "IndicTransToolkit",
                "collator.py"
            )

            if os.path.exists(test):
                collator_path = test
                break

        if collator_path:

            with open(collator_path, "r", encoding="utf-8") as f:
                content = f.read()

            content = content.replace(
                "from transformers.tokenization_utils import PreTrainedTokenizerBase",
                "from transformers import PreTrainedTokenizerBase",
            )

            with open(collator_path, "w", encoding="utf-8") as f:
                f.write(content)

            print("✓ IndicTransToolkit fixed")

    except Exception as e:
        print(e)


fix_collator()

from IndicTransToolkit import IndicProcessor

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

print("Using device:", DEVICE)

if DEVICE == "cuda":
    print(torch.cuda.get_device_name(0))

# -----------------------------------------------------------------------------
# Model IDs
# -----------------------------------------------------------------------------

HI_KO_MODEL = "shriyadiker/KonkaniTrans-Hi-Ko"
KO_HI_MODEL = "shriyadiker/KonkaniTrans-Ko-Hi"

print("\nLoading Hindi → Konkani model...")

hi_ko_tokenizer = AutoTokenizer.from_pretrained(
    HI_KO_MODEL,
    trust_remote_code=True,
)

hi_ko_model = AutoModelForSeq2SeqLM.from_pretrained(
    HI_KO_MODEL,
    trust_remote_code=True,
    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
).to(DEVICE)

print("✓ Hindi → Konkani loaded")

print("\nLoading Konkani → Hindi model...")

ko_hi_tokenizer = AutoTokenizer.from_pretrained(
    KO_HI_MODEL,
    trust_remote_code=True,
)

ko_hi_model = AutoModelForSeq2SeqLM.from_pretrained(
    KO_HI_MODEL,
    trust_remote_code=True,
    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
).to(DEVICE)

print("✓ Konkani → Hindi loaded")

ip = IndicProcessor(inference=True)

print("\nAll models loaded successfully.")

# -----------------------------------------------------------------------------
# File for storing corrections
# -----------------------------------------------------------------------------

CORRECTIONS_FILE = "translation_corrections.txt"
ENTRY_COUNTER_FILE = "entry_counter.json"

def get_next_entry_number():
    """Get the next entry number from counter file"""
    try:
        if os.path.exists(ENTRY_COUNTER_FILE):
            with open(ENTRY_COUNTER_FILE, 'r', encoding='utf-8') as f:
                data = json.load(f)
                return data.get('next_entry', 1)
    except:
        pass
    return 1

def update_entry_counter(entry_number):
    """Update the entry counter after saving"""
    try:
        with open(ENTRY_COUNTER_FILE, 'w', encoding='utf-8') as f:
            json.dump({'next_entry': entry_number + 1}, f)
    except:
        pass

def save_correction(direction, source, original_target, corrected_target, entry_number):
    """Save the correction to file with proper formatting"""
    
    # Ensure the directory exists
    os.makedirs(os.path.dirname(CORRECTIONS_FILE) if os.path.dirname(CORRECTIONS_FILE) else '.', exist_ok=True)
    
    # Format the entry
    entry = f"""entryno:{entry_number}
direction:{direction}
source:{source}
target:{corrected_target}
---------------------
"""
    
    # Append to file
    with open(CORRECTIONS_FILE, 'a', encoding='utf-8') as f:
        f.write(entry)
    
    return entry_number + 1

# -----------------------------------------------------------------------------
# Translation
# -----------------------------------------------------------------------------

def translate(text, direction):

    if text is None or text.strip() == "":
        return ""

    try:

        if direction == "Hindi → Konkani":

            tokenizer = hi_ko_tokenizer
            model = hi_ko_model

            src_lang = "hin_Deva"
            tgt_lang = "gom_Deva"

        else:

            tokenizer = ko_hi_tokenizer
            model = ko_hi_model

            src_lang = "gom_Deva"
            tgt_lang = "hin_Deva"

        batch = ip.preprocess_batch(
            [text],
            src_lang=src_lang,
            tgt_lang=tgt_lang,
        )

        inputs = tokenizer(
            batch,
            truncation=True,
            padding=True,
            return_tensors="pt",
            return_attention_mask=True,
        ).to(DEVICE)

        with torch.no_grad():

            generated = model.generate(
                **inputs,
                max_length=256,
                num_beams=5,
                early_stopping=True,
            )

        decoded = tokenizer.batch_decode(
            generated,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=True,
        )

        output = ip.postprocess_batch(
            decoded,
            lang=tgt_lang,
        )[0]

        output = fix_escaped_unicode(output)

        return output

    except Exception as e:

        return str(e)


UNICODE_ESCAPE_PATTERN = re.compile(r'\\u([0-9a-fA-F]{4})')

def fix_escaped_unicode(text):

    if not isinstance(text, str):
        return text

    if "\\u" not in text:
        return text

    text = UNICODE_ESCAPE_PATTERN.sub(
        lambda m: chr(int(m.group(1), 16)),
        text
    )

    text = text.replace("ऱ", "ऱ")
    text = text.replace("ऴ", "ऴ")

    return text

# -----------------------------------------------------------------------------
# CSS
# -----------------------------------------------------------------------------

css = """
footer {
    display: none !important;
}

.gradio-container {
    max-width: 900px !important;
    margin: auto !important;
    padding: 30px 20px !important;
}

textarea {
    font-size: 18px !important;
    font-family: 'Segoe UI', 'Noto Sans Devanagari', Arial, sans-serif !important;
    line-height: 1.8 !important;
    padding: 15px !important;
    min-height: 180px !important;
}

button {
    padding: 12px 40px !important;
    font-size: 16px !important;
    font-weight: 600 !important;
    border-radius: 8px !important;
}

.title {
    text-align: center;
    margin-bottom: 30px;
}

.title h1 {
    font-size: 28px;
    font-weight: 600;
    color: #1a1a1a;
    margin-bottom: 5px;
}

.title p {
    color: #666;
    font-size: 14px;
}

.correction-section {
    margin-top: 20px;
    padding: 20px;
    background: #f8f9fa;
    border-radius: 10px;
    border: 1px solid #e9ecef;
}

.correction-section textarea {
    min-height: 100px !important;
}

.success-message {
    color: #28a745;
    font-weight: 600;
    padding: 10px;
    background: #d4edda;
    border-radius: 5px;
    margin-top: 10px;
}

.error-message {
    color: #dc3545;
    font-weight: 600;
    padding: 10px;
    background: #f8d7da;
    border-radius: 5px;
    margin-top: 10px;
}
"""

# -----------------------------------------------------------------------------
# Helper Functions
# -----------------------------------------------------------------------------

def clear_text():
    return "", "", "", ""

def submit_correction(direction, source, original_target, corrected_target):
    """Handle submission of corrected translation"""
    
    # Validate inputs
    if not source or source.strip() == "":
        return "", "Error: Source text is empty. Please translate something first."
    
    if not corrected_target or corrected_target.strip() == "":
        return "", "Error: Corrected translation is empty. Please provide the corrected text."
    
    if corrected_target == original_target:
        return "", "No changes detected. Please edit the translation before submitting."
    
    # Get next entry number
    entry_number = get_next_entry_number()
    
    # Save correction
    try:
        save_correction(direction, source, original_target, corrected_target, entry_number)
        update_entry_counter(entry_number)
        
        success_msg = f"✓ Correction saved successfully! (Entry #{entry_number})"
        return "", success_msg
        
    except Exception as e:
        return "", f"Error saving correction: {str(e)}"

# -----------------------------------------------------------------------------
# Gradio UI
# -----------------------------------------------------------------------------

with gr.Blocks(
    title="VIDYAAPATI (Hindi-Konkani) | Goa University",
    css=css,
) as demo:

    # Updated Heading with Project Description & Affiliation
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 25px;">
        <h2 style="margin-bottom: 6px; color: #1a1a1a; font-size: 24px;">VIDYAAPATI (Hindi-Konkani)</h2>
        <div style="font-weight: 600; color: #4a5568; margin-bottom: 10px;">Goa University</div>
        <p style="font-size: 13px; color: #666; line-height: 1.5; max-width: 750px; margin: 0 auto;">
            <b>Project:</b> "VIDYAAPATI: Bidirectional Machine Translation Involving Bengali, Konkani, Maithili, Marathi, and Hindi". 
            Under the project titled, <b>"National Language Translation Mission (NLTM): BHASHINI"</b> 
            funded by the Ministry of Electronics and Information Technology (MeitY), Government of India.
        </p>
        <div style="margin-top: 10px; padding: 8px; background: #f0f4ff; border-radius: 5px; font-size: 13px; color: #2563eb;">
            💡 Translate, edit if needed, and submit corrections to help improve the model!
        </div>
    </div>
    """)

    # Direction dropdown
    direction = gr.Dropdown(
        choices=[
            "Hindi → Konkani",
            "Konkani → Hindi",
        ],
        value="Hindi → Konkani",
        label="Translation Direction",
    )

    # Input textbox
    input_text = gr.Textbox(
        label="Input Text",
        placeholder="Type or paste text here...",
        lines=6,
        max_lines=12,
    )

    # Translate button
    translate_btn = gr.Button(
        "Translate",
        variant="primary",
    )

    # Output textbox (editable)
    output_text = gr.Textbox(
        label="Translation (edit if needed)",
        placeholder="Translation will appear here...",
        lines=6,
        max_lines=12,
    )

    # Hidden field to store original translation
    original_translation = gr.State("")

    # Correction section
    with gr.Column(elem_classes="correction-section"):
        gr.Markdown("### ✏️ Submit Correction")
        gr.Markdown("If you've edited the translation above, click the button below to submit your correction.")
        
        with gr.Row():
            submit_btn = gr.Button(
                "📤 Submit Correction",
                variant="secondary",
            )
            clear_btn = gr.Button(
                "🗑️ Clear All",
                variant="stop",
            )
        
        status_message = gr.Textbox(
            label="Status",
            lines=2,
            interactive=False,
            show_label=True,
        )

    # Event handlers
    def translate_and_store(text, direction):
        result = translate(text, direction)
        return result, result, "", ""  # output_text, original_translation, status
    
    translate_btn.click(
        fn=translate_and_store,
        inputs=[input_text, direction],
        outputs=[output_text, original_translation, status_message],
    )

    input_text.submit(
        fn=translate_and_store,
        inputs=[input_text, direction],
        outputs=[output_text, original_translation, status_message],
    )

    # Submit correction handler
    submit_btn.click(
        fn=submit_correction,
        inputs=[direction, input_text, original_translation, output_text],
        outputs=[status_message, status_message],  # The first output is unused, second is status
    )

    # Clear all handler
    clear_btn.click(
        fn=clear_text,
        inputs=[],
        outputs=[input_text, output_text, original_translation, status_message],
    )

# -----------------------------------------------------------------------------
# Launch
# -----------------------------------------------------------------------------

if __name__ == "__main__":
    demo.queue()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
    )