import streamlit as st
import os
import pytesseract
import easyocr
from PIL import Image
import numpy as np

# Optional: Try PaddleOCR if installed
try:
    from paddleocr import PaddleOCR
    paddle_available = True
except ImportError:
    paddle_available = False

# Set page configuration
st.set_page_config(page_title="KAIRO.ai - Hugging Face Demo", layout="wide")
st.title("🧠 KAIRO.ai - AI 기반 언어 검수 플랫폼 (HF Demo)")

# Sidebar: OCR Engine
st.sidebar.header("📂 OCR 엔진 선택")
available_engines = ["EasyOCR", "Tesseract"]
if paddle_available:
    available_engines.append("PaddleOCR")

ocr_engines = st.sidebar.multiselect(
    "사용할 OCR 엔진을 선택하세요",
    available_engines,
    default=["EasyOCR"]
)

# Sidebar: 검수 항목
st.sidebar.markdown("### 🧠 검수 항목 선택")

checklist_korean = [
    "맞춤법 (Orthography)", "띄어쓰기 (Spacing Rules)", "문법 오류 (Grammatical Errors)",
    "표준어 규정 위반 (Standard Language Regulation)", "외래어 표기법 위반 (Loanword Orthography)",
    "문장 구조 오류 (Sentence Structure Error)", "어휘 적절성 (Lexical Appropriateness)",
    "중복어/군더더기 제거 (Redundancy Elimination)", "오탈자 (Typographical Errors)",
    "문장 부호 사용 (Punctuation Usage)", "어법 오류 (Usage Error)",
    "논리적 일관성 (Logical Coherence)", "형식적 오류 (Formatting Consistency)"
]

checklist_foreign = [
    "Spelling (철자 오류)", "Grammar (문법 오류)", "Capitalization (대소문자 오류)",
    "Subject-Verb Agreement (주어-동사 일치)", "Article Usage (관사의 적절성)",
    "Tense Consistency (시제 일관성)", "Word Choice (어휘 선택)",
    "Redundancy (중복 표현)", "Tone (문체 및 일관성)"
]

selected_korean_checks = st.sidebar.multiselect("✅ 한글 검수 항목 선택", checklist_korean)
selected_foreign_checks = st.sidebar.multiselect("✅ 외국어 검수 항목 선택", checklist_foreign)

# Cache EasyOCR for speed
@st.cache_resource
def load_easyocr():
    return easyocr.Reader(["ko", "en"])

# Main Section
log_messages = []
progress = st.empty()
uploaded_file = st.file_uploader("PNG 이미지 업로드", type=["png"])

if uploaded_file is not None:
    with st.spinner("📤 이미지 처리 중..."):
        try:
            progress.progress(10, text="이미지 열기")
            image = Image.open(uploaded_file)
            st.image(image, caption="업로드된 이미지", use_column_width=True)
            log_messages.append("✅ 이미지 업로드 성공")

            text_output = ""

            if "EasyOCR" in ocr_engines:
                try:
                    progress.progress(30, text="EasyOCR 인식 중...")
                    reader = load_easyocr()
                    result = reader.readtext(np.array(image), detail=0)
                    text_output += "\n".join(result) + "\n"
                    log_messages.append("✅ EasyOCR 처리 완료")
                except Exception as e:
                    log_messages.append(f"❌ EasyOCR 오류: {str(e)}")

            if "Tesseract" in ocr_engines:
                try:
                    progress.progress(60, text="Tesseract 인식 중...")
                    tess_path = shutil.which("tesseract")
                    if tess_path is None:
                        raise EnvironmentError("Tesseract is not installed or not in PATH.")
                    tess_result = pytesseract.image_to_string(image, lang="kor+eng")
                    text_output += tess_result + "\n"
                    log_messages.append("✅ Tesseract 처리 완료")
                except Exception as e:
                    log_messages.append(f"❌ Tesseract 오류: {str(e)}")

            if "PaddleOCR" in ocr_engines and paddle_available:
                try:
                    progress.progress(90, text="PaddleOCR 인식 중...")
                    paddle_ocr = PaddleOCR(use_angle_cls=True, lang='korean')
                    result = paddle_ocr.ocr(np.array(image), cls=True)
                    paddle_text = "\n".join([line[1][0] for line in result[0]])
                    text_output += paddle_text + "\n"
                    log_messages.append("✅ PaddleOCR 처리 완료")
                except Exception as e:
                    log_messages.append(f"❌ PaddleOCR 오류: {str(e)}")

            st.markdown("### 🔍 추출된 텍스트")
            st.text_area("OCR 결과", text_output, height=300)
            progress.progress(100, text="✅ OCR 완료")

            if selected_korean_checks or selected_foreign_checks:
                st.markdown("### 🛠️ 선택된 검수 항목")
                if selected_korean_checks:
                    st.markdown("**📌 한글 검수 항목:**")
                    for check in selected_korean_checks:
                        st.write(f"• {check}")
                if selected_foreign_checks:
                    st.markdown("**🌐 외국어 검수 항목:**")
                    for check in selected_foreign_checks:
                        st.write(f"• {check}")
                log_messages.append("✅ 검수 항목 적용 완료")
        except Exception as e:
            log_messages.append(f"❌ 전체 오류 발생: {str(e)}")

# Footer Log
st.markdown("---")
st.markdown("### 📋 시스템 로그")
for msg in log_messages:
    if "✅" in msg:
        st.success(msg)
    elif "❌" in msg:
        st.error(msg)
    else:
        st.info(msg)