import os
import pickle
import sys
import traceback
import numpy as np
import streamlit as st
import torch
import torch.nn.functional as F
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# --- [1] 파일 및 폴더 경로 설정 ---
BASE_DIR = os.path.dirname(__file__)
KC_DIR = os.path.join(BASE_DIR, "kcbert_web")
DEBERTA_DIR = os.path.join(BASE_DIR, "deberta_web")
CNN_PATH = os.path.join(BASE_DIR, "char_cnn_web.pt")
VOCAB_PATH = os.path.join(BASE_DIR, "vocab.pkl")

# --- [2] 모델 로드 함수 (서버 기동 시 딱 한 번 실행) ---
@st.cache_resource(show_spinner=False)
def load_verification_models():
    try:
        # 1. KcBERT 로드
        kc_tokenizer = AutoTokenizer.from_pretrained(KC_DIR)
        kc_model = AutoModelForSequenceClassification.from_pretrained(KC_DIR)
        kc_model.eval()
        
        # 2. DeBERTa 로드
        deberta_tokenizer = AutoTokenizer.from_pretrained(DEBERTA_DIR)
        deberta_model = AutoModelForSequenceClassification.from_pretrained(DEBERTA_DIR)
        deberta_model.eval()
        
        # 3. Char-CNN 및 보카 로드
        vocab = None
        if os.path.exists(VOCAB_PATH):
            with open(VOCAB_PATH, "rb") as f:
                vocab = pickle.load(f)
        cnn_model = None  # 추후 고유 CNN 구조 필요 시 연결
        
        return kc_tokenizer, kc_model, deberta_tokenizer, deberta_model, cnn_model
    except Exception as e:
        raise RuntimeError(
            "모델 로드 중 오류가 발생했습니다. 모델 폴더 경로 및 파일들을 확인해주세요.\n"
            f"상세 오류: {e}"
        )

# 모델 로드 구동
try:
    kc_tok, kc_mod, deb_tok, deb_mod, cnn_mod = load_verification_models()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    # 모델들을 적절한 디바이스로 이동 (CPU/GPU)
    kc_mod.to(device)
    deb_mod.to(device)
except Exception as e:
    st.error("⚠️ 시스템 초기화 실패 (모델 로드 에러)")
    st.code(str(e))
    st.stop()

# --- [3] 두 문장의 스타일 유사도를 계산하는 함수 ---
def calculate_similarity(text_a, text_b):
    with torch.no_grad():
        # --- 1) KcBERT 스타일 벡터 분석 ---
        inputs_a = kc_tok(text_a, return_tensors="pt", truncation=True, max_length=128).to(device)
        outputs_a = kc_mod(**inputs_a)
        prob_a_kc = F.softmax(outputs_a.logits, dim=-1).flatten()

        inputs_b = kc_tok(text_b, return_tensors="pt", truncation=True, max_length=128).to(device)
        outputs_b = kc_mod(**inputs_b)
        prob_b_kc = F.softmax(outputs_b.logits, dim=-1).flatten()

        kc_sim = F.cosine_similarity(prob_a_kc.unsqueeze(0), prob_b_kc.unsqueeze(0)).item() * 100

        # --- 2) DeBERTa 스타일 벡터 분석 ---
        deb_inputs_a = deb_tok(text_a, return_tensors="pt", truncation=True, max_length=128).to(device)
        deb_outputs_a = deb_mod(**deb_inputs_a)
        prob_a_deb = F.softmax(deb_outputs_a.logits, dim=-1).flatten()

        deb_inputs_b = deb_tok(text_b, return_tensors="pt", truncation=True, max_length=128).to(device)
        deb_outputs_b = deb_mod(**deb_inputs_b)
        prob_b_deb = F.softmax(deb_outputs_b.logits, dim=-1).flatten()

        deb_sim = F.cosine_similarity(prob_a_deb.unsqueeze(0), prob_b_deb.unsqueeze(0)).item() * 100

        # --- 3) 최종 앙상블 유사도 산출 ---
        final_similarity = (kc_sim + deb_sim) / 2.0
        
        return kc_sim, deb_sim, final_similarity

# --- [4] Streamlit UI 디자인 (동일인 식별 전용) ---
st.set_page_config(page_title="저자 식별 시스템", layout="wide", page_icon="🕵️‍♂️")
st.title("🕵️‍♂️ 알파 프로젝트: 문체학 기반 저자 동일성 검증 시스템")
st.subheader("두 개의 글을 비교하여 동일 인물이 작성했는지 실시간으로 분석합니다.")
st.write("---")

# 화면을 왼쪽, 오른쪽 두 칸으로 분할
col1, col2 = st.columns(2)
with col1:
    st.markdown("### 📝 분석 대상 글 A")
    text_a = st.text_area(
        "첫 번째 글을 입력하세요:",
        placeholder="비교할 첫 번째 본문을 입력하세요.",
        height=250,
        key="text_a",
    )
with col2:
    st.markdown("### 📝 분석 대상 글 B")
    text_b = st.text_area(
        "두 번째 글을 입력하세요:",
        placeholder="비교할 두 번째 본문을 입력하세요.",
        height=250,
        key="text_b",
    )

st.write("---")

# 실행 버튼
if st.button("🔍 동일인 여부 정밀 검증 시작", use_container_width=True):
    if not text_a.strip() or not text_b.strip():
        st.error("⚠️ 글 A와 글 B 모두 텍스트를 입력해야 분석이 가능합니다!")
    else:
        try:
            with st.spinner("3대장 알고리즘이 두 문장의 문체 패턴을 대조하는 중..."):
                # 유사도 연산 실행
                kc_sim, deb_sim, final_sim = calculate_similarity(text_a, text_b)

            # --- [5] 판정 결과 시각화 ---
            st.success("🎉 문체 대조 분석 완료!")

            # 기준점(Threshold) 설정
            THRESHOLD = 85.0
            is_same_author = final_sim >= THRESHOLD

            # 대형 패널로 결과 노출
            rect_col1, rect_col2 = st.columns(2)
            with rect_col1:
                if is_same_author:
                    st.metric(label="최종 판정 결과", value="🟢 동일인 가능성 매우 높음")
                else:
                    st.metric(label="최종 판정 결과", value="🔴 다른 인물일 가능성 높음")
            with rect_col2:
                st.metric(label="최종 문체 유사도 점수", value=f"{final_sim:.2f} / 100점")

            st.write("---")
            
            # 진행 바 시각화 추가
            st.progress(min(max(final_sim / 100.0, 0.0), 1.0))

            # 레이아웃 분할: 왼쪽은 차트, 오른쪽은 상세 표
            result_col1, result_col2 = st.columns([4, 3])
            
            with result_col1:
                st.markdown("### 📊 모델별 문체 대조 스코어")
                scores = {
                    "KcBERT 대조 점수": kc_sim,
                    "DeBERTa 대조 점수": deb_sim,
                    "최종 앙상블 결론": final_sim,
                }
                st.bar_chart(scores)

            with result_col2:
                st.markdown("### 📋 모델별 결과 수치")
                model_table = {
                    "평가 항목": ["KcBERT", "DeBERTa", "최종 앙상블 결과"],
                    "유사도 스코어": [f"{kc_sim:.2f}점", f"{deb_sim:.2f}점", f"{final_sim:.2f}점"]
                }
                st.table(model_table)

            # 연구실 보고용 상세 텍스트 피드백
            st.info(
                f"💡 **분석 결과 요약**: 두 문장의 문장 구조, 단어 선택 패턴, 어조를 종합한 결과 "
                f"최종 **{final_sim:.1f}%**의 일치율을 보였습니다. (판정 기준선: {THRESHOLD}%)"
            )

            # 고급 확장 탭 정보 (기존 디버깅용 UI 계승)
            with st.expander("🛠️ 시스템 기술 정보"):
                st.write("실행 장치 (Device):", device)
                st.write("Python 실행 경로:")
                st.code(sys.executable)
                st.write("모델 기본 디렉토리:")
                st.code(BASE_DIR)

            st.caption("주의: 본 결과는 AI 모델 기반 추정값이며, 실제 동일인을 단정하는 법적 판단이 아닙니다.")

        except Exception as e:
            st.error("고급 모델 실행 중 오류가 발생했습니다.")
            st.code(str(e))
            with st.expander("상세 오류 보기 (Traceback)"):
                st.code(traceback.format_exc())