import streamlit as st
import leafmap.foliumap as leafmap
import folium
import requests
from urllib.parse import quote, quote_plus

st.set_page_config(
    layout="wide",
    initial_sidebar_state="expanded",
    page_title="COG Viewer",
    page_icon="🗺️"

)

# 隐藏右上角的菜单按钮
# st.markdown("""
#     <style>
#         #MainMenu {visibility: hidden;}
#         header {visibility: hidden;}
#         footer {visibility: hidden;}
#         [data-testid="stHeader"] {visibility: hidden;}
#         [data-testid="stToolbar"] {visibility: hidden;}
#     </style>
#     """, unsafe_allow_html=True)

# 侧边栏配置
st.sidebar.title("🗺️ COG 图层配置")

# COG 数据源配置
st.sidebar.subheader("数据源")
cog_url = st.sidebar.text_input(
    "COG URL",
    value="http://minio.minio-dev:9000/staticfiles/s2_mosaic_cog.tif",
    help="COG 文件的 URL 地址"
)

# 波段和表达式配置
st.sidebar.subheader("波段配置")
indexes = st.sidebar.text_input(
    "波段索引 (indexes)",
    value="4,8",
    help="要使用的波段索引，用逗号分隔，例如: 4,8"
)

expression = st.sidebar.text_input(
    "表达式 (expression)",
    value="(b2-b1)/(b2+b1)",
    help="波段计算表达式，例如: (b2-b1)/(b2+b1) 表示 NDVI"
)

# 数值范围配置
st.sidebar.subheader("数值范围")
rescale_min = st.sidebar.number_input(
    "最小值 (rescale min)",
    value=-1.0,
    step=0.1,
    format="%.2f"
)
rescale_max = st.sidebar.number_input(
    "最大值 (rescale max)",
    value=1.0,
    step=0.1,
    format="%.2f"
)

# 颜色映射配置
st.sidebar.subheader("颜色映射")
# titiler 支持的颜色映射列表（小写，使用下划线）
colormap_options = [
    "rdylgn", "viridis", "plasma", "inferno", "magma", "cividis",
    "spectral", "rdylbu", "rdgy", "rdbu", "piyg", "prgn", "brbg",
    "puor", "oranges", "greens", "blues", "reds", "greys",
    "turbo", "jet", "rainbow", "coolwarm", "seismic", "terrain",
    "hot", "cool", "spring", "summer", "autumn", "winter",
    "bone", "copper", "pink", "gray", "binary", "gist_earth",
    "gist_rainbow", "gist_stern", "gist_heat", "gist_ncar",
    "nipy_spectral", "tab10", "tab20", "set1", "set2", "set3",
    "pastel1", "pastel2", "paired", "accent", "dark2", "flag",
    "prism", "ocean", "gist_gray", "gist_yarg", "afmhot",
    "gist_rainbow_r", "hsv", "cubehelix", "brg", "hsv_r"
]

use_custom_colormap = st.sidebar.checkbox(
    "使用自定义颜色映射",
    value=False,
    help="勾选后可以输入自定义的颜色映射名称"
)

if use_custom_colormap:
    colormap_name = st.sidebar.text_input(
        "自定义颜色映射名称",
        value="rdylgn",
        help="输入 titiler 支持的颜色映射名称（小写，使用下划线）"
    )
else:
    colormap_name = st.sidebar.selectbox(
        "颜色映射 (colormap)",
        options=colormap_options,
        index=0,  # 默认使用 rdylgn
        help="选择颜色映射方案（titiler 支持的格式）"
    )

# 图层显示配置
st.sidebar.subheader("显示设置")
layer_opacity = st.sidebar.slider(
    "透明度 (opacity)",
    min_value=0.0,
    max_value=1.0,
    value=0.75,
    step=0.1
)

layer_name = st.sidebar.text_input(
    "图层名称",
    value="S2 Mosaic COG",
    help="在地图图层控制中显示的图层名称"
)

# 在侧边栏底部添加使用说明
with st.sidebar.expander("ℹ️ 使用说明", expanded=False):
    st.markdown("""
    **快速开始：**
    1. 输入 COG URL
    2. 配置波段和表达式
    3. 调整颜色映射和透明度
    4. 地图会自动定位
    
    **提示：** 修改 COG URL 后会自动获取信息并重新定位
    """)

# 获取 COG 信息
def get_cog_info(cog_url):
    """从 titiler 获取 COG 的完整信息"""
    try:
        encoded_url = quote_plus(cog_url)
        info_url = f"https://titiler.dev.maimaiag.com/cog/info?url={encoded_url}"
        response = requests.get(info_url, timeout=15)
        if response.status_code == 200:
            info = response.json()
            return info
        else:
            st.error(f"获取 COG 信息失败: HTTP {response.status_code}")
            return None
    except requests.exceptions.Timeout:
        st.error("请求超时，请检查网络连接或 COG URL")
        return None
    except Exception as e:
        st.error(f"无法获取 COG 信息: {e}")
        return None

# 在侧边栏添加获取信息按钮
st.sidebar.subheader("COG 信息")
auto_fit = st.sidebar.checkbox(
    "自动定位到 COG 范围",
    value=True,
    help="勾选后会自动根据 COG 边界定位地图"
)

if st.sidebar.button("🔄 获取 COG 信息", help="点击获取 COG 的详细信息"):
    with st.spinner("正在获取 COG 信息..."):
        cog_info = get_cog_info(cog_url)
        if cog_info:
            st.session_state['cog_info'] = cog_info
            # titiler 返回的是 bounds 字段，格式为 [minx, miny, maxx, maxy]
            st.session_state['cog_bbox'] = cog_info.get('bounds')
            st.sidebar.success("✓ COG 信息获取成功！")
        else:
            st.session_state['cog_info'] = None
            st.session_state['cog_bbox'] = None

# 检查 COG URL 是否改变，如果改变则清除旧信息并重新获取
if 'last_cog_url' not in st.session_state:
    st.session_state['last_cog_url'] = cog_url
    st.session_state['cog_info'] = None
    st.session_state['cog_bbox'] = None
elif st.session_state['last_cog_url'] != cog_url:
    # URL 改变了，清除旧信息
    st.session_state['last_cog_url'] = cog_url
    st.session_state['cog_info'] = None
    st.session_state['cog_bbox'] = None
    # URL 改变时，如果启用了自动定位，立即获取新信息
    if auto_fit and cog_url.strip():
        with st.spinner("检测到 COG URL 变化，正在获取新信息..."):
            cog_info = get_cog_info(cog_url)
            if cog_info:
                st.session_state['cog_info'] = cog_info
                # titiler 返回的是 bounds 字段，格式为 [minx, miny, maxx, maxy]
                st.session_state['cog_bbox'] = cog_info.get('bounds')
                # st.success("✓ 已自动定位到新的 COG 范围")
            else:
                st.session_state['cog_bbox'] = None

# 获取 COG 边界（从 session_state 或自动获取）
if st.session_state.get('cog_bbox') is None:
    if auto_fit and cog_url.strip():
        with st.spinner("正在获取 COG 边界信息..."):
            cog_info = get_cog_info(cog_url)
            if cog_info:
                st.session_state['cog_info'] = cog_info
                # titiler 返回的是 bounds 字段，格式为 [minx, miny, maxx, maxy]
                st.session_state['cog_bbox'] = cog_info.get('bounds')
            else:
                st.session_state['cog_bbox'] = None

cog_bbox = st.session_state.get('cog_bbox')
cog_info = st.session_state.get('cog_info')

# 显示 COG 信息
if cog_info:
    with st.sidebar.expander("📊 COG 详细信息", expanded=False):
        if 'bounds' in cog_info:
            bounds = cog_info['bounds']
            # bounds 格式: [minx, miny, maxx, maxy] = [最小经度, 最小纬度, 最大经度, 最大纬度]
            st.write(f"**边界框 (Bounds):**")
            st.write(f"- 最小经度: {bounds[0]:.6f}")
            st.write(f"- 最小纬度: {bounds[1]:.6f}")
            st.write(f"- 最大经度: {bounds[2]:.6f}")
            st.write(f"- 最大纬度: {bounds[3]:.6f}")
            # 计算中心点
            center_lon = (bounds[0] + bounds[2]) / 2
            center_lat = (bounds[1] + bounds[3]) / 2
            st.write(f"**中心点:** ({center_lat:.6f}, {center_lon:.6f})")
        
        if 'width' in cog_info and 'height' in cog_info:
            st.write(f"**尺寸:** {cog_info['width']} × {cog_info['height']} 像素")
        
        if 'count' in cog_info:
            st.write(f"**波段数:** {cog_info['count']}")
        
        if 'crs' in cog_info:
            st.write(f"**坐标系:** {cog_info['crs']}")
        
        if 'dtype' in cog_info:
            st.write(f"**数据类型:** {cog_info['dtype']}")

# 创建地图
if cog_bbox and auto_fit:
    # bounds 格式: [minx, miny, maxx, maxy] = [最小经度, 最小纬度, 最大经度, 最大纬度]
    # 计算中心点
    center_lat = (cog_bbox[1] + cog_bbox[3]) / 2
    center_lon = (cog_bbox[0] + cog_bbox[2]) / 2
    
    # 计算合适的缩放级别（基于边界框大小）
    lat_diff = cog_bbox[3] - cog_bbox[1]
    lon_diff = cog_bbox[2] - cog_bbox[0]
    max_diff = max(lat_diff, lon_diff)
    
    # 根据范围大小估算合适的缩放级别
    if max_diff > 10:
        initial_zoom = 6
    elif max_diff > 5:
        initial_zoom = 7
    elif max_diff > 1:
        initial_zoom = 8
    elif max_diff > 0.5:
        initial_zoom = 9
    elif max_diff > 0.1:
        initial_zoom = 10
    else:
        initial_zoom = 11
    
    # 创建地图并设置初始视图
    m = leafmap.Map(
        center=[center_lat, center_lon],
        zoom=initial_zoom,
        minimap_control=True
    )
    
    # 使用边界框来适应视图（确保 COG 完全可见）
    # fit_bounds 需要 [[min_lat, min_lon], [max_lat, max_lon]] 格式
    m.fit_bounds([[cog_bbox[1], cog_bbox[0]], [cog_bbox[3], cog_bbox[2]]])
    
    # 添加边界框标记（可选）
    if st.sidebar.checkbox("显示 COG 边界框", value=False):
        folium.Rectangle(
            bounds=[[cog_bbox[1], cog_bbox[0]], [cog_bbox[3], cog_bbox[2]]],
            color="red",
            weight=2,
            fill=False,
            popup=f"COG 边界框\n经度范围: {cog_bbox[0]:.6f} ~ {cog_bbox[2]:.6f}\n纬度范围: {cog_bbox[1]:.6f} ~ {cog_bbox[3]:.6f}"
        ).add_to(m)
else:
    # 如果无法获取边界或未启用自动定位，使用默认视图
    m = leafmap.Map(
        center=[39.9042, 116.4074],  # 默认北京
        zoom=8,
        minimap_control=True
    )

# 添加底图作为参考
m.add_basemap("OpenStreetMap")

# 只有在有边界信息时才添加 COG 图层，避免无效请求
if cog_bbox and cog_url.strip():
    # 构建 COG 瓦片 URL
    # 对 URL 参数进行编码
    # 使用 quote_plus 确保 + 号被编码为 %2B
    encoded_cog_url = quote_plus(cog_url)
    encoded_expression = quote_plus(expression)

    # 拼接完整的瓦片 URL
    cog_tiles_url = (
        f"https://titiler.dev.maimaiag.com/cog/tiles/WebMercatorQuad/{{z}}/{{x}}/{{y}}.png"
        f"?url={encoded_cog_url}"
        f"&indexes={indexes}"
        f"&expression={encoded_expression}"
        f"&rescale={rescale_min},{rescale_max}"
        f"&colormap_name={colormap_name}"
    )

    # 在侧边栏显示生成的 URL（用于调试）
    if st.sidebar.checkbox("显示生成的 URL", value=False):
        st.sidebar.code(cog_tiles_url, language=None)

    # 添加自定义 COG 瓦片图层
    folium.TileLayer(
        tiles=cog_tiles_url,
        attr="COG Viewer - Titiler",
        name=layer_name,
        overlay=True,
        control=True,
        opacity=layer_opacity,
        show=True,  # 默认显示
        # error_tile_url="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
    ).add_to(m)
else:
    # 如果没有边界信息，显示提示
    if cog_url.strip() and not cog_bbox:
        st.info("⚠️ 无法获取 COG 边界信息，请检查 COG URL 是否正确或点击'获取 COG 信息'按钮手动获取。")
# 添加图层控制
m.add_layer_control()

# 使用更大的地图高度，占据更多屏幕空间
# 根据常见屏幕高度，使用 900px 可以让地图占据大部分可视区域
m.to_streamlit(height=800)