# app.py
import streamlit as st
import cv2
import numpy as np
from PIL import Image
import easyocr
import os
from streamlit_paste_button import paste_image_button as pbutton

color_ranges = {
    'fire': 'B50B0E',
    'water': '015AB6',
    'wind': '1F6A0B',
    'earth': '623F23',
    'light': 'DA8D09',
    'dark': '502181'
}

def hex_to_rgb(hex_code):
    return tuple(int(hex_code[i:i+2], 16) for i in (0, 2, 4))

# # 캐싱으로 Reader를 한 번만 로드
# @st.cache_resource
# def load_reader():
#     return easyocr.Reader(['en'], gpu=False, verbose=False)

@st.cache_resource
def load_number_templates():
    templates = {}
    
    template_dir = './templates'
    
    for num in range(1, 16):
        template_path = os.path.join(template_dir, f'{num}.png')
        
        if os.path.exists(template_path):
            # 그레이스케일로 로드
            template = cv2.imread(template_path, cv2.IMREAD_GRAYSCALE)
            
            if template is not None:
                # 전처리: 이진화
                _, template = cv2.threshold(template, 127, 255, cv2.THRESH_BINARY)
                templates[num] = template
            else:
                st.warning(f"템플릿 {num}.png를 로드할 수 없습니다.")
    if not templates:
        st.error("템플릿 이미지를 찾을 수 없습니다. templates/ 폴더에 1.png ~ 15.png 파일을 추가해주세요.")

    return templates

def extract_region(region):
    try:
        gray = cv2.cvtColor(region, cv2.COLOR_BGR2GRAY)
        _, binary = cv2.threshold(gray, 215, 255, cv2.THRESH_BINARY)
                
        # 노이즈 제거
        kernel = np.ones((2, 2), np.uint8)
        cleaned = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_OPEN, kernel)
        return cleaned

    except Exception as e:
        return None

def template_matching(region, templates):
    try:
        h, w = region.shape
        aspect_ratio = w / h
        new_width = int(25 * aspect_ratio)
        region = cv2.resize(region, (new_width, 25))
        best_match = 0
        best_score = -1
        for num, template in templates.items():
            try:
                result = cv2.matchTemplate(region, template, cv2.TM_CCOEFF_NORMED)
                _, max_val, _, _ = cv2.minMaxLoc(result)
                print(max_val)
                if max_val > best_score:
                    best_score = max_val
                    best_match = num
            except Exception as e:
                print(f"Error matching template {num}: {e}")
                continue

        if best_score > 0.65:
            return best_match
        else:
            return 0
    except Exception as e:
        print(f"Error matching template: {e}")
        return 0


def find_items(img_array, color_range, templates):
    img_rgb = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
    results = {}

    for type, hex_color in color_range.items():
        target_rgb = hex_to_rgb(hex_color)
        lower_c = np.array([max(0, c - 10) for c in target_rgb])
        upper_c = np.array([min(255, c + 10) for c in target_rgb])
        
        mask = cv2.inRange(img_rgb, lower_c, upper_c)
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        if contours:
            largest = max(contours, key=cv2.contourArea)
            if cv2.contourArea(largest) > 100:
                x, y, w, h = cv2.boundingRect(largest)
                number_region = img_rgb[y+int(h*1.2) :y+ int(h*1.7), x+int(w*1.4):x+int(w*2.2)]
                count = template_matching(extract_region(number_region), templates)
                results[type] = count
            else:
                results[type] = 0
        else:
            results[type] = 0

    return results


def process_image(img, templates):
    st.image(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), caption='업로드된 이미지', use_container_width=True)
    
    with st.spinner('아이템 개수를 세는 중...'):
        try:
            results = find_items(img, color_ranges, templates)
            
            # 결과 표시
            st.success('✅ 분석 완료!')
            
            col1, col2, col3 = st.columns(3)
            
            emoji_map = {
                'fire': '🔥',
                'water': '💧',
                'wind': '💨',
                'earth': '🌍',
                'light': '✨',
                'dark': '🌙'
            }
            
            korean_map = {
                'fire': '불',
                'water': '물',
                'wind': '바람',
                'earth': '대지',
                'light': '빛',
                'dark': '어둠'
            }

            counts_str = ''
            # 단일 열로 표시
            for type, count in results.items():
                st.metric(
                    label=f"{emoji_map.get(type, '')} {korean_map.get(type, type)}", 
                    value=f"{count}개"
                )

            # cols = [col1, col2, col3]
            # counts_str = ''
            # for idx, (type, count) in enumerate(results.items()):
            #     col = cols[idx % 3]
            #     with col:
            #         st.metric(
            #             label=f"{emoji_map.get(type, '')} {korean_map.get(type, type)}", 
            #             value=f"{count}개"
            #         )
                counts_str += f"{count % 10 if count >= 10 else count}/"
            st.markdown(f'{"/".join(korean_map.values())}')
            st.markdown(f'{counts_str}')
        except Exception as e:
            st.error(f"분석 중 오류 발생: {e}")
            st.info("EasyOCR 모델 로딩에 실패했을 수 있습니다. Streamlit Cloud의 메모리 제한 때문일 수 있습니다.")


# Streamlit UI
st.set_page_config(page_title="속성 잠재력 주문서 카운터", page_icon="🎮")

st.title('🎮 속성 잠재력 주문서 카운터')
st.write('획득한 속성 잠재력 주문서 개수를 자동으로 세어드립니다!')

templates = load_number_templates()

tab1, tab2 = st.tabs(["📁 파일 업로드", "📋 붙여넣기"])

with tab1:
    st.write('이미지 파일을 업로드하세요')
    uploaded_file = st.file_uploader("스크린샷을 업로드하세요", type=['png', 'jpg', 'jpeg'])

    if uploaded_file is not None:
        file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
        img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
        process_image(img, templates)

with tab2:
    st.write('이미지를 붙여넣으세요')
    paste_result = pbutton(
        label="📋 여기를 클릭하고 Ctrl+V",
        background_color="#FF4B4B",
        hover_background_color="#FF6B6B",
    )

    if paste_result.image_data is not None:
        pil_image = paste_result.image_data
        # st.markdown(f'{type(pil_image)}')
        img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
        process_image(img, templates)


st.markdown('---')
st.caption('Made by ❤️sseong')