import os
import gradio as gr
from transformers import AutoConfig

MODEL_ID = "PKUDigitalHealth/ECG-R1-8B-RL"

try:
    config = AutoConfig.from_pretrained(MODEL_ID)
    CONFIG_STATUS = "✅ Конфигурация загружена"
except Exception as e:
    config = None
    CONFIG_STATUS = f"❌ Ошибка загрузки конфигурации:\n{e}"


def get_info():
    result = []

    result.append("ECG-R1 DIAGNOSTICS")
    result.append("=" * 50)
    result.append("")
    result.append(f"Статус: {CONFIG_STATUS}")
    result.append("")

    if config is not None:
        result.append(f"Архитектура: {config.architectures}")
        result.append(f"Тип модели: {config.model_type}")
        result.append(f"Hidden Size: {config.hidden_size}")

        if hasattr(config, "text_config"):
            result.append("")
            result.append("TEXT CONFIG")
            result.append(f"Layers: {config.text_config.num_hidden_layers}")
            result.append(f"Heads: {config.text_config.num_attention_heads}")
            result.append(f"Intermediate Size: {config.text_config.intermediate_size}")

        result.append("")
        result.append("ECG CONFIG")
        result.append(f"ECG Hidden Size: {config.ecg_hidden_size}")
        result.append(f"ECG Model Config: {config.ecg_model_config}")
        result.append(f"ECG Projector: {config.ecg_projector_type}")
        result.append(f"ECG Tower Path: {config.ecg_tower_path}")

        result.append("")
        result.append("FILE CHECKS")

        ecg_file = "ecg_coca/checkpoint/cpt_wfep_epoch_20.pt"

        result.append(
            f"ECG Encoder Exists: {os.path.exists(ecg_file)}"
        )

        if os.path.exists("ecg_coca"):
            result.append("✅ ecg_coca folder found")
        else:
            result.append("❌ ecg_coca folder missing")

        if os.path.exists("ecg_r1"):
            result.append("✅ ecg_r1 folder found")
        else:
            result.append("❌ ecg_r1 folder missing")

    return "\n".join(result)


with gr.Blocks(title="ECG-R1 Russian") as demo:
    gr.Markdown("# ECG-R1 Russian")
    gr.Markdown("Диагностика модели ECG-R1")

    btn = gr.Button("Проверить систему")

    output = gr.Textbox(
        label="Результат",
        lines=30
    )

    btn.click(
        fn=get_info,
        outputs=output
    )

demo.launch()