import os
import re
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel

# ==========================================
# 1. 모델 및 토크나이저 로드 (LoRa Inference Engine)
# ==========================================
LORA_DIR = "./lora_climate_misinfo"
BASE_MODEL_NAME = "klue/roberta-base"  # 기본 Base LLM

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

try:
    tokenizer = AutoTokenizer.from_pretrained(LORA_DIR)
    base_model = AutoModelForSequenceClassification.from_pretrained(
        BASE_MODEL_NAME, 
        num_labels=2
    )
    model = PeftModel.from_pretrained(base_model, LORA_DIR)
    model.to(device)
    model.eval()
    MODEL_LOADED = True
except Exception as e:
    print(f"[경고] 모델 로드 중 오류 발생 (시뮬레이션 모드로 전환): {e}")
    MODEL_LOADED = False

# ==========================================
# 2. 알고리즘 파이프라인 함수
# ==========================================
def preprocess_text(text: str) -> str:
    """텍스트 전처리 (Preprocessing)"""
    text = re.sub(r'\s+', ' ', text)
    return text.strip()

def extract_xai_attention(text: str):
    """
    XAI 추출 및 어텐션 맵 분석 (Attention Analysis & Highlight Word)
    UI 스케치 양식(WHO -> 소규모 집단 등) 표현
    """
    words = text.split()
    if len(words) >= 2:
        src_word = words[0]
        target_word = words[1] if len(words) > 1 else "소규모 집단"
    else:
        src_word = "WHO"
        target_word = "소규모 집단"
    
    return src_word, target_word

def run_pipeline(input_text: str):
    """전체 추론 및 결과 생성 파이프라인"""
    if not input_text.strip():
        return "⚠️ 텍스트를 입력해주세요.", "", "", "", "내용을 입력하세요."

    # Step 1: 텍스트 전처리
    cleaned_text = preprocess_text(input_text)
    
    # Step 2 & 3: 토큰화 및 LoRa 추론
    if MODEL_LOADED:
        inputs = tokenizer(cleaned_text, return_tensors="pt", truncation=True, max_length=512).to(device)
        with torch.no_grad():
            outputs = model(**inputs)
            probs = torch.softmax(outputs.logits, dim=-1).squeeze().cpu().numpy()
            fake_prob = float(probs[1]) if len(probs) > 1 else float(probs[0])
    else:
        # 가상 시뮬레이션 결과 (모델 파일 미로드 시)
        fake_prob = 0.98

    # Step 4: XAI 및 어텐션 분석
    src_w, tgt_w = extract_xai_attention(cleaned_text)
    
    # Step 5: 결과 구성 (UI 스케치 반영)
    is_fake = fake_prob >= 0.5
    verdict_badge = f"{'🔴 가짜뉴스' if is_fake else '🟢 진짜뉴스'} {fake_prob:.2f}"
    model_tag = "Active Dolphin"
    highlight_html = f"""
    <div style="border: 2px solid #333; border-radius: 15px; padding: 15px; text-align: center; margin: 10px 0;">
        <h4 style="margin-top:0;">Highlight Word</h4>
        <div style="display: flex; justify-content: center; align-items: center; gap: 15px;">
            <span style="border: 1px solid #666; border-radius: 10px; padding: 8px 15px; font-weight: bold;">{src_w}</span>
            <span>➔</span>
            <span style="border: 1px solid #666; border-radius: 10px; padding: 8px 15px; font-weight: bold;">{tgt_w}</span>
        </div>
    </div>
    """
    
    risk_score_text = "위험도 점수 +0.0007%"
    disclaimer_text = "⚠️ 풍자-반어법 텍스트 정확도 낮음!"

    return verdict_badge, model_tag, highlight_html, risk_score_text, disclaimer_text

def submit_appeal(user_reason: str):
    """이의 제기 처리 함수"""
    if not user_reason.strip():
        return "이의 제기 사유를 입력해주세요."
    return f"✅ 이의 제기가 접수되었습니다. (담당자 검토 예정 - dev-team3@example.com)"

# ==========================================
# 3. Gradio UI 인터페이스 (UI 스케치 구현)
# ==========================================
custom_css = """
.verdict-box {
    border: 2px solid #333;
    border-radius: 12px;
    padding: 8px 15px;
    display: inline-block;
    font-size: 1.2rem;
    font-weight: bold;
}
.risk-box {
    border: 2px solid #333;
    border-radius: 12px;
    padding: 10px;
    text-align: center;
    font-size: 1.1rem;
    margin-top: 10px;
}
.disclaimer-box {
    border: 2px solid #ff4d4d;
    background-color: #fff2f2;
    color: #cc0000;
    border-radius: 10px;
    padding: 10px;
    font-weight: bold;
    text-align: center;
}
"""

with gr.Blocks(css=custom_css, title="기후변화 가짜뉴스 판별기") as demo:
    gr.Markdown("## 🔍 기후변화 가짜뉴스 AI 판별 및 XAI 분석 시스템")
    
    with gr.Row():
        with gr.Column(scale=1):
            input_area = gr.Textbox(
                label="뉴스 기사 입력 (Preprocess & Tokenize)",
                placeholder="검증할 기후변화/과학 관련 뉴스 텍스트를 입력하세요...",
                lines=8
            )
            btn_analyze = gr.Button("🔍 분석 및 판정 실행", variant="primary")
            
            with gr.Accordion("📋 Model Card 및 한계 고지 사항 확인", open=False):
                gr.Markdown("""
                - **학습 데이터:** 2024~2026년 기후/환경 관련 뉴스 200건
                - **주요 한계:**
                  1. 공신력 있는 기관 위장 시 오판 가능성
                  2. **풍자/반어법 텍스트 정확도 낮음**
                  3. 거절 정보 복합 문맥의 기계 오판 가능성
                - **개발자 책임 선언:** 본 모델의 판정 결과는 보조 지표이며 최종 결정 근거로 사용할 수 없습니다.
                """)

        with gr.Column(scale=1):
            # 1. 판정 결과 및 모델 태그
            with gr.Row():
                out_verdict = gr.Textbox(label="판정결과", elem_classes=["verdict-box"], interactive=False)
                out_tag = gr.Textbox(label="Model Tag", value="Active Dolphin", interactive=False)

            # 2. Highlight Word (어텐션 맵 visual)
            out_highlight = gr.HTML(label="Highlight Word")

            # 3. 위험도 점수 결과
            out_risk = gr.Textbox(label="결과", elem_classes=["risk-box"], interactive=False)

            # 4. 한계 고지 경고 문구 (UI 스케치 하단)
            out_disclaimer = gr.Textbox(
                label="한계 고지 알림", 
                elem_classes=["disclaimer-box"], 
                interactive=False
            )

            # 5. 이의 제기 버튼 및 모달 레이아웃
            with gr.Row():
                btn_appeal_open = gr.Button("이의 제기", variant="secondary")
            
            with gr.Group(visible=False) as appeal_group:
                appeal_input = gr.Textbox(label="이의 제기 사유 입력", placeholder="오판이라 생각하는 이유를 작성해주세요.")
                btn_appeal_submit = gr.Button("제출하기")
                appeal_status = gr.Markdown()

    # 이벤트 리스너 연결
    btn_analyze.click(
        fn=run_pipeline,
        inputs=[input_area],
        outputs=[out_verdict, out_tag, out_highlight, out_risk, out_disclaimer]
    )

    btn_appeal_open.click(
        fn=lambda: gr.update(visible=True),
        inputs=None,
        outputs=[appeal_group]
    )

    btn_appeal_submit.click(
        fn=submit_appeal,
        inputs=[appeal_input],
        outputs=[appeal_status]
    )

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