import gradio as gr
import subprocess
import os
import sys
from datetime import datetime
import shutil

# --- CONFIGURATION UPDATED FOR HYBRID MODEL ---
TRAINING_SCRIPT = "train_hybrid.py"
MODEL_OUTPUT_DIR = "checkpoints"
MODEL_FILE_NAME = "layoutlmv3_bilstm_crf_hybrid.pth"
MODEL_FILE_PATH = os.path.join(MODEL_OUTPUT_DIR, MODEL_FILE_NAME)

# ----------------------------------------------------------------

def retrieve_model():
    """
    Checks for the final model file and prepares it for download.
    Useful for when the training job finishes server-side but the
    client connection has timed out.
    """
    if os.path.exists(MODEL_FILE_PATH):
        file_size = os.path.getsize(MODEL_FILE_PATH) / (1024 * 1024)  # Size in MB

        # Copy to a simple location that Gradio can reliably serve
        import tempfile
        temp_dir = tempfile.gettempdir()
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        temp_model_path = os.path.join(temp_dir, f"hybrid_model_recovered_{timestamp}.pth")

        try:
            shutil.copy2(MODEL_FILE_PATH, temp_model_path)
            download_path = temp_model_path

            log_output = (
                f"--- Model Status Check: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ---\n"
                f"🎉 SUCCESS! The Hybrid LayoutLMv3+BiLSTM+CRF model was found.\n"
                f"📦 Model file: {MODEL_FILE_PATH}\n"
                f"📊 Model size: {file_size:.2f} MB\n"
                f"🔗 Download path prepared: {download_path}\n\n"
                f"⬇️ Click the '📥 Download Model' button below to save your model."
            )
            return log_output, download_path, gr.Button(visible=True)

        except Exception as e:
            log_output = (
                f"--- Model Status Check FAILED ---\n"
                f"⚠️ Trained model found, but could not prepare for download: {e}\n"
                f"📁 Original Path: {MODEL_FILE_PATH}. Try again or check Space logs."
            )
            return log_output, None, gr.Button(visible=False)

    else:
        log_output = (
            f"--- Model Status Check: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ---\n"
            f"❌ Model file not found at {MODEL_FILE_PATH}.\n"
            f"Training may still be running or it failed. Check back later."
        )
        return log_output, None, gr.Button(visible=False)


def clear_memory(dataset_file: gr.File):
    """
    Deletes the model output directory and the uploaded dataset file.
    """
    log_output = f"--- Memory Clear Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ---\n"

    # 1. Clear Model Checkpoints Directory
    if os.path.exists(MODEL_OUTPUT_DIR):
        try:
            shutil.rmtree(MODEL_OUTPUT_DIR)
            log_output += f"✅ Successfully deleted model directory: {MODEL_OUTPUT_DIR}\n"
        except Exception as e:
            log_output += f"❌ ERROR deleting model directory {MODEL_OUTPUT_DIR}: {e}\n"
    else:
        log_output += f"ℹ️ Model directory not found: {MODEL_OUTPUT_DIR} (Nothing to delete)\n"

    # 2. Clear Uploaded Dataset File (Temporary file cleanup)
    if dataset_file is not None:
        input_path = dataset_file.name if hasattr(dataset_file, 'name') else str(dataset_file)
        if os.path.exists(input_path):
            try:
                os.remove(input_path)
                log_output += f"✅ Successfully deleted uploaded dataset file: {input_path}\n"
            except Exception as e:
                log_output += f"❌ ERROR deleting dataset file {input_path}: {e}\n"
        else:
            log_output += f"ℹ️ Uploaded dataset file not found at {input_path}.\n"
    else:
        log_output += f"ℹ️ No dataset file currently tracked for deletion.\n"

    log_output += f"--- Memory Clear Complete: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ---\n"
    log_output += "✨ Files and checkpoints have been removed. You can now start a fresh training run."

    return log_output, None, gr.Button(visible=False), None


def train_model(dataset_file: gr.File, batch_size: int, epochs: int, lr: float, max_len: int, progress=gr.Progress()):
    """
    Handles the Gradio submission and executes the training script using subprocess.
    """
    # 1. Setup: Create output directory
    os.makedirs(MODEL_OUTPUT_DIR, exist_ok=True)

    # 2. File Handling
    if dataset_file is None:
        yield "❌ ERROR: Please upload a file.", None, gr.Button(visible=False)
        return

    input_path = dataset_file.name if hasattr(dataset_file, 'name') else str(dataset_file)

    if not os.path.exists(input_path):
        yield f"❌ ERROR: Uploaded file not found at {input_path}.", None, gr.Button(visible=False)
        return

    progress(0.1, desc="Initializing Hybrid Model Training...")

    log_output = f"--- Training Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ---\n"
    log_output += f"🤖 Architecture: LayoutLMv3 + BiLSTM + CRF\n"

    # 3. Construct the subprocess command
    command = [
        sys.executable,
        TRAINING_SCRIPT,
        "--mode", "train",
        "--input", input_path,
        "--batch_size", str(batch_size),
        "--epochs", str(epochs),
        "--lr", str(lr),
        "--max_len", str(max_len)
    ]

    log_output += f"Executing command: {' '.join(command)}\n\n"
    yield log_output, None, gr.Button(visible=False)

    try:
        # 4. Run the training script
        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1
        )

        # Stream logs
        for line in iter(process.stdout.readline, ""):
            log_output += line
            print(line, end='')
            yield log_output, None, gr.Button(visible=False)

        process.stdout.close()
        return_code = process.wait()

        # 5. Check completion
        if return_code == 0:
            log_output += "\n" + "=" * 60 + "\n"
            log_output += "✅ HYBRID TRAINING COMPLETE!\n"
            log_output += "=" * 60 + "\n"

            if os.path.exists(MODEL_FILE_PATH):
                file_size = os.path.getsize(MODEL_FILE_PATH) / (1024 * 1024)
                log_output += f"\n📦 Model file found: {MODEL_FILE_PATH} ({file_size:.2f} MB)"

                # Copy for download
                import tempfile
                temp_dir = tempfile.gettempdir()
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                temp_model_path = os.path.join(temp_dir, f"hybrid_model_{timestamp}.pth")

                try:
                    shutil.copy2(MODEL_FILE_PATH, temp_model_path)
                    download_path = temp_model_path
                except Exception as e:
                    log_output += f"\n⚠️ Copy failed: {e}, using original path"
                    download_path = MODEL_FILE_PATH

                log_output += f"\n\n⬇️ Click the '📥 Download Model' button below."
                yield log_output, download_path, gr.Button(visible=True)
                return
            else:
                log_output += f"\n❌ Error: Training finished but {MODEL_FILE_PATH} was not found."
                yield log_output, None, gr.Button(visible=False)
                return
        else:
            log_output += f"\n❌ TRAINING FAILED with return code {return_code}\n"
            yield log_output, None, gr.Button(visible=False)
            return

    except FileNotFoundError:
        yield log_output + f"\n❌ ERROR: '{TRAINING_SCRIPT}' not found.", None, gr.Button(visible=False)
    except Exception as e:
        yield log_output + f"\n❌ Unexpected Error: {e}", None, gr.Button(visible=False)


# --- Gradio Interface Setup ---
with gr.Blocks(title="Hybrid LayoutLM Training", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🧬 Hybrid LayoutLMv3 + BiLSTM + CRF Training")
    gr.Markdown(
        """
        **Architecture:** This app trains a state-of-the-art stack:
        1. **LayoutLMv3** (Visual & Textual Embeddings)
        2. **Bi-LSTM** (Sequence Context Modeling)
        3. **CRF** (Label Consistency Enforcement)
        
        **Instructions:** Upload your Label Studio JSON, set parameters, and train. 
        **Note:** This model is slower to train than standard LayoutLM but typically achieves higher accuracy on complex layouts.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📁 Dataset")
            file_input = gr.File(label="Upload Label Studio JSON", file_types=[".json"])

            gr.Markdown("### ⚙️ Hyperparameters")
            batch_size_input = gr.Slider(1, 16, value=4, step=1, label="Batch Size")
            epochs_input = gr.Slider(1, 10, value=5, step=1, label="Epochs")
            lr_input = gr.Number(value=2e-5, label="Learning Rate (Backbone)", info="LSTM/CRF head uses 1e-4")
            max_len_input = gr.Slider(128, 512, value=512, step=128, label="Max Seq Len")

            train_button = gr.Button("🔥 Start Hybrid Training", variant="primary", size="lg")
            check_button = gr.Button("🔍 Check Status / Recover Model", variant="secondary")
            clear_button = gr.Button("🧹 Clear Files", variant="stop")

        with gr.Column(scale=2):
            log_output = gr.Textbox(
                label="Training Logs", lines=25, autoscroll=True, show_copy_button=True,
                placeholder="Logs will appear here..."
            )
            
            download_btn = gr.Button("📥 Download Hybrid Model", variant="primary", size="lg", visible=False)
            
            # State and hidden download component
            model_path_state = gr.State(value=None)
            model_download = gr.File(label="Download", interactive=False, visible=True)

    # Actions
    train_button.click(
        fn=train_model,
        inputs=[file_input, batch_size_input, epochs_input, lr_input, max_len_input],
        outputs=[log_output, model_path_state, download_btn]
    )

    check_button.click(
        fn=retrieve_model,
        inputs=[],
        outputs=[log_output, model_path_state, download_btn]
    )

    clear_button.click(
        fn=clear_memory,
        inputs=[file_input],
        outputs=[log_output, model_path_state, download_btn, model_download]
    )

    download_btn.click(
        fn=lambda path: path,
        inputs=[model_path_state],
        outputs=[model_download]
    )

if __name__ == "__main__":
    demo.launch()