import os
import torch
import torch.nn.functional as F
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel

# ---------------------------------------------------------------------------
# 1. 모델 및 토크나이저 로드 (LoRA Inference Engine)
# ---------------------------------------------------------------------------
LORA_PATH = "./lora_adapter"
BASE_MODEL_NAME = "klue/roberta-base"

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

tokenizer = AutoTokenizer.from_pretrained(
    LORA_PATH if os.path.exists(os.path.join(LORA_PATH, "tokenizer.json")) else BASE_MODEL_NAME
)

try:
    base_model = AutoModelForSequenceClassification.from_pretrained(
        BASE_MODEL_NAME, 
        num_labels=2, 
        output_attentions=True
    )
    model = PeftModel.from_pretrained(base_model, LORA_PATH)
    model.to(device)
    model.eval()
    MODEL_LOADED = True
except Exception as e:
    print(f"모델 로딩 중 경고/오류 (기본 데모 모드로 전환): {e}")
    MODEL_LOADED = False

# ---------------------------------------------------------------------------
# 2. 알고리즘 흐름도 기반 파이프라인 함수
# ---------------------------------------------------------------------------
def process_and_infer(user_text):
    if not user_text.strip():
        return "텍스트를 입력해주세요.", "0%", [], "입력값이 없습니다."

    # Step 1 & 2: 텍스트 전처리 및 토크나이징
    inputs = tokenizer(
        user_text, 
        return_tensors="pt", 
        truncation=True, 
        max_length=512, 
        padding=True
    ).to(device)

    # Step 3: LoRA 모델 추론 및 어텐션 맵 extraction (XAI)
    if MODEL_LOADED:
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            probs = F.softmax(logits, dim=-1)[0]
            
            # 어텐션 분석 (마지막 레이어의 어텐션 평균 활용)
            attentions = outputs.attentions
            last_layer_attn = attentions[-1][0].mean(dim=0)
            token_importance = last_layer_attn.sum(dim=0)
            
            pred_idx = torch.argmax(probs).item()
            confidence = probs[pred_idx].item() * 100
    else:
        # 모델 미로드 시 시뮬레이션 동작
        pred_idx = 1 if "미세플라스틱" in user_text or "무해" in user_text else 0
        confidence = 92.4
        tokens = tokenizer.tokenize(user_text)
        token_importance = torch.rand(len(tokens) + 2)

    # Label 매핑
    label_map = {0: "진실/신뢰할 수 있음 (True)", 1: "기후/환경 허위 정보 (Misinformation)"}
    result_label = label_map.get(pred_idx, "알 수 없음")
    
    # Step 4: 영향 단어 및 기여도 분석 (XAI Extraction)
    input_tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    scores = token_importance.cpu().tolist()
    
    xai_data = []
    total_score = sum(scores) if sum(scores) > 0 else 1.0
    for tok, sc in zip(input_tokens, scores):
        if tok not in [tokenizer.cls_token, tokenizer.sep_token, tokenizer.pad_token, "<s>", "</s>", "<pad>"]:
            clean_tok = tok.replace(" ", "")
            contrib = round((sc / total_score) * 100, 2)
            if clean_tok:
                xai_data.append((clean_tok, contrib))

    xai_data = sorted(xai_data, key=lambda x: x[1], reverse=True)[:5]
    xai_formatted = [[word, f"{score}%"] for word, score in xai_data]

    return result_label, f"{confidence:.1f}%", xai_formatted

def handle_appeal(appeal_text, user_input):
    if not appeal_text.strip():
        return "⚠️ 이의 제기 내용을 입력한 후 제출해주세요."
    return "✅ 이의 제기가 성공적으로 접수되었습니다. 모델 재학습 및 피드백 검토에 반영됩니다."

# ---------------------------------------------------------------------------
# 3. Gradio UI 스케치 레이아웃 구축
# ---------------------------------------------------------------------------
custom_css = """
.container { max-width: 900px; margin: auto; }
.result-box { background-color: #f0f7ff; border-radius: 8px; padding: 15px; border-left: 5px solid #2b6cb0; }
.disclaimer-box { background-color: #fff5f5; border-radius: 8px; padding: 15px; border-left: 5px solid #e53e3e; margin-top: 20px; }
"""

with gr.Blocks(css=custom_css, title="기후변화 허위정보 판별 AI") as demo:
    gr.Markdown("# 🌿 LoRA 기반 기후변화/미세플라스틱 허위정보 검증 시스템")
    gr.Markdown("알고리즘 및 XAI 기여도 분석을 통해 기후변화 관련 문장의 허위성을 검증합니다.")
    
    with gr.Row():
        input_text = gr.Textbox(
            label="검증할 문장 입력", 
            placeholder="예: 미세플라스틱은 체내에 전혀 흡수되지 않고 모두 배출되므로 무해하다.",
            lines=3
        )
    submit_btn = gr.Button("판정 및 분석 실행", variant="primary")
    
    gr.Divider()
    
    # [구역 1: 상단] 판정 결과 (점수/분류)
    gr.Markdown("### 1. 모델 판정 결과")
    with gr.Row(elem_classes=["result-box"]):
        out_result = gr.Textbox(label="판정 결과 (분류)", interactive=False)
        out_confidence = gr.Textbox(label="신뢰도 (점수)", interactive=False)
        
    # [구역 2: 중앙] 영향 단어 + 기여도 % (XAI Extraction)
    gr.Markdown("### 2. 판단 근거 추출 (XAI - 영향 단어 및 기여도 %)")
    out_xai_table = gr.Dataframe(
        headers=["영향 단어 (Tokens)", "기여도 (%)"],
        datatype=["str", "str"],
        interactive=False,
        row_count=5
    )

    # [구역 3: 하단] 한계 고지 + [이의 제기] 버튼
    with gr.Column(elem_classes=["disclaimer-box"]):
        gr.Markdown("### 3. Model Card 한계 및 사전 고지 (Model Disclaimer)")
        gr.Markdown(
            """
            ⚠️ **모델의 한계 및 주의사항:**
            1. **비꼬는 표현/반어법 제한**: 칭찬이나 우회적 비꼬기 표현이 포함된 기후 관련 문장은 잘못 판정할 위험이 높습니다 (위험 수준: **높음**).
            2. **전문 학술 용어 오반응**: 최신 과학 전문 용어나 문맥이 복잡한 경우 오반응할 수 있습니다 (위험 수준: **중간**).
            3. **책임 선언**: 본 모델은 보조 검증용이며, 최종 판단 및 법적/의학적 책임은 모델 개발자 및 시스템에 있지 않음을 사전에 인지하여 주시기 바랍니다.
            """
        )
        
        with gr.Accordion("📢 판정 결과에 동의하지 않으신가요? [이의 제기하기]", open=False):
            appeal_input = gr.Textbox(label="이의 제기 사유 입력", placeholder="판정이 잘못되었다고 생각하는 이유나 근거를 적어주세요.")
            appeal_btn = gr.Button("이의 제기 제출", variant="secondary")
            appeal_status = gr.Markdown()

    # Event Handlers
    submit_btn.click(
        fn=process_and_infer,
        inputs=[input_text],
        outputs=[out_result, out_confidence, out_xai_table]
    )
    
    appeal_btn.click(
        fn=handle_appeal,
        inputs=[appeal_input, input_text],
        outputs=[appeal_status]
    )

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