import io
import os
import glob
import base64
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import geopandas as gpd
from flask import Blueprint, jsonify, request, current_app

bp = Blueprint('drought', __name__, url_prefix='/api/drought')

# SHP 렌더링 임시 JPG 경로
_SHP_JPG_DIR = None

def _get_shp_jpg_dir():
    global _SHP_JPG_DIR
    if _SHP_JPG_DIR is None:
        _SHP_JPG_DIR = os.path.join(current_app.static_folder, 'drought', 'tmp')
    os.makedirs(_SHP_JPG_DIR, exist_ok=True)
    return _SHP_JPG_DIR


def _cleanup_old_jpg():
    """이전 렌더링된 JPG 파일 삭제"""
    jpg_dir = _get_shp_jpg_dir()
    for f in glob.glob(os.path.join(jpg_dir, 'drought_shp_*.jpg')):
        try:
            os.remove(f)
        except OSError:
            pass


def _point_to_grade(point):
    """point 값(0~1)을 4등급으로 변환"""
    if point <= 0.25:
        return '관심'
    elif point <= 0.50:
        return '주의'
    elif point <= 0.75:
        return '경계'
    else:
        return '심각'


# 등급별 색상 (관심→심각: 초록→빨강)
GRADE_COLORS = {
    '관심': '#2ecc71',
    '주의': '#f1c40f',
    '경계': '#e67e22',
    '심각': '#e74c3c',
}


def render_shp_to_base64(grid_size=500):
    """SHP + XLSX 매칭하여 point 등급별 색상으로 격자 렌더링. 이전 JPG는 삭제."""
    import time
    import pandas as pd
    from matplotlib.patches import Patch

    _cleanup_old_jpg()

    base_dir = os.path.join(current_app.static_folder, 'drought', str(grid_size))
    shp_path = os.path.join(base_dir, f'goheung_{grid_size}m_merge.shp')
    xlsx_path = os.path.join(base_dir, f'goheung_{grid_size}m_center_lon_lat.xlsx')

    if not os.path.exists(shp_path):
        return None

    gdf = gpd.read_file(shp_path)

    # XLSX에서 point 값 읽어서 GRID_CD로 매칭
    if os.path.exists(xlsx_path):
        df = pd.read_excel(xlsx_path)
        point_map = dict(zip(df['GRID_CD'], df['point']))
        gdf['point'] = gdf['GRID_CD'].map(point_map).fillna(0.0)
    else:
        gdf['point'] = 0.0

    # point → 등급 → 색상
    gdf['grade'] = gdf['point'].apply(_point_to_grade)
    gdf['color'] = gdf['grade'].map(GRADE_COLORS)

    fig, ax = plt.subplots(figsize=(10, 10))

    # 등급별로 그려서 범례 생성
    for grade, color in GRADE_COLORS.items():
        subset = gdf[gdf['grade'] == grade]
        if not subset.empty:
            subset.plot(ax=ax, color=color, edgecolor='#555', linewidth=0.3, alpha=0.8)

    # 범례
    legend_handles = [Patch(facecolor=c, edgecolor='#555', label=g) for g, c in GRADE_COLORS.items()]
    ax.legend(handles=legend_handles, loc='lower right', fontsize=11, title='가뭄 등급', title_fontsize=12)

    ax.set_axis_off()
    plt.tight_layout(pad=0)

    # JPG 파일로 저장
    jpg_dir = _get_shp_jpg_dir()
    jpg_name = f'drought_shp_{int(time.time() * 1000)}.jpg'
    jpg_path = os.path.join(jpg_dir, jpg_name)
    fig.savefig(jpg_path, format='jpg', dpi=150, bbox_inches='tight', pad_inches=0)

    # base64 반환
    buf = io.BytesIO()
    fig.savefig(buf, format='jpg', dpi=150, bbox_inches='tight', pad_inches=0)
    buf.seek(0)
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close(fig)

    return image_base64


@bp.route('/shp-render', methods=['GET'])
def shp_render():
    """SHP 파일을 JPG로 렌더링하여 반환"""
    grid_size = request.args.get('grid', 500, type=int)
    try:
        image = render_shp_to_base64(grid_size)
        if image is None:
            return jsonify({'status': 'error', 'message': 'SHP 파일을 찾을 수 없습니다.'})
        return jsonify({'status': 'success', 'image': image})
    except Exception as e:
        return jsonify({'status': 'error', 'message': str(e)})


def generate_plot():
    """가뭄 예측 샘플 플롯 생성"""
    fig, ax = plt.subplots(figsize=(10, 6))

    # 샘플 데이터 (실제 데이터로 교체 필요)
    np.random.seed(456)
    months = ['1월', '2월', '3월', '4월', '5월', '6월',
              '7월', '8월', '9월', '10월', '11월', '12월']
    drought_index = np.random.uniform(-2, 2, 12)

    colors = ['brown' if x < 0 else 'green' for x in drought_index]
    ax.bar(months, drought_index, color=colors, alpha=0.7)
    ax.axhline(y=0, color='black', linestyle='-', linewidth=0.5)
    ax.set_title('가뭄 예측 결과 (월별 가뭄 지수)', fontsize=14)
    ax.set_xlabel('월')
    ax.set_ylabel('가뭄 지수')
    ax.grid(True, alpha=0.3, axis='y')

    # 이미지를 base64로 변환
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
    buf.seek(0)
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close(fig)

    return image_base64


@bp.route('/predict', methods=['GET'])
def predict():
    """가뭄 예측 결과 반환"""
    try:
        image = generate_plot()
        return jsonify({
            'status': 'success',
            'message': '가뭄 예측 완료',
            'image': image
        })
    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e),
            'image': None
        })


def generate_simulation_plot():
    """가뭄 시뮬레이션 플롯 생성"""
    fig, ax = plt.subplots(figsize=(10, 6))

    np.random.seed(333)
    weeks = np.arange(1, 53)
    drought_severity = np.sin(weeks / 10) * 1.5 + np.random.normal(0, 0.2, 52)

    colors = ['brown' if x < 0 else 'forestgreen' for x in drought_severity]
    ax.bar(weeks, drought_severity, color=colors, alpha=0.7, width=0.8)
    ax.axhline(y=0, color='black', linestyle='-', linewidth=0.5)
    ax.axhline(y=-1, color='red', linestyle='--', label='심각 가뭄')
    ax.set_title('연간 가뭄 시뮬레이션', fontsize=14)
    ax.set_xlabel('주차')
    ax.set_ylabel('가뭄 심각도')
    ax.legend()
    ax.grid(True, alpha=0.3, axis='y')

    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
    buf.seek(0)
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close(fig)

    return image_base64


def generate_waterlevel_plot():
    """가뭄 수위 곡선 플롯 생성"""
    fig, ax = plt.subplots(figsize=(10, 6))

    np.random.seed(444)
    days = np.arange(1, 91)
    reservoir_level = 80 - np.cumsum(np.abs(np.random.randn(90)) * 0.3)
    reservoir_level = np.maximum(reservoir_level, 20)

    ax.plot(days, reservoir_level, 'b-', linewidth=2)
    ax.fill_between(days, 0, reservoir_level, alpha=0.3, color='blue')
    ax.axhline(y=40, color='orange', linestyle='--', label='경고 수위')
    ax.axhline(y=25, color='red', linestyle='--', label='위험 수위')
    ax.set_title('저수지 수위 변화', fontsize=14)
    ax.set_xlabel('일')
    ax.set_ylabel('저수율 (%)')
    ax.set_ylim(0, 100)
    ax.legend()
    ax.grid(True, alpha=0.3)

    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
    buf.seek(0)
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close(fig)

    return image_base64


@bp.route('/simulation', methods=['GET'])
def simulation():
    """가뭄 시뮬레이션 결과 반환"""
    try:
        image = generate_simulation_plot()
        return jsonify({
            'status': 'success',
            'message': '가뭄 시뮬레이션 완료',
            'image': image
        })
    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e),
            'image': None
        })


@bp.route('/waterlevel', methods=['GET'])
def waterlevel():
    """가뭄 수위 곡선 결과 반환"""
    try:
        image = generate_waterlevel_plot()
        return jsonify({
            'status': 'success',
            'message': '수위 곡선 조회 완료',
            'image': image
        })
    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e),
            'image': None
        })


def generate_heatmap_image(size=(400, 300)):
    """가뭄 히트맵 이미지 생성"""
    fig, ax = plt.subplots(figsize=(size[0]/100, size[1]/100))

    np.random.seed(72)

    # 히트맵 데이터 생성 (가뭄 위험도)
    x = np.linspace(0, 10, 100)
    y = np.linspace(0, 8, 80)
    X, Y = np.meshgrid(x, y)

    # 다중 가우시안으로 히트맵 생성
    Z = np.zeros_like(X)
    # 중심 핫스팟 (심각 지역)
    Z += 0.9 * np.exp(-((X-5)**2 + (Y-4)**2) / 3)
    # 주변 핫스팟들
    Z += 0.6 * np.exp(-((X-3)**2 + (Y-2)**2) / 2)
    Z += 0.7 * np.exp(-((X-7)**2 + (Y-5)**2) / 2.5)
    Z += 0.4 * np.exp(-((X-2)**2 + (Y-6)**2) / 1.5)

    # 노이즈 추가
    Z += np.random.rand(80, 100) * 0.1

    # 커스텀 컬러맵 (녹색 -> 노랑 -> 주황 -> 빨강)
    from matplotlib.colors import LinearSegmentedColormap
    colors = ['#27ae60', '#f1c40f', '#e67e22', '#e74c3c']
    cmap = LinearSegmentedColormap.from_list('drought', colors)

    ax.imshow(Z, cmap=cmap, aspect='auto', alpha=0.8)
    ax.axis('off')

    plt.tight_layout(pad=0)
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=100, bbox_inches='tight', pad_inches=0, transparent=True)
    buf.seek(0)
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close(fig)

    return image_base64


def generate_mini_heatmap_image():
    """미니 히트맵 이미지 생성"""
    fig, ax = plt.subplots(figsize=(2.6, 1.2))

    np.random.seed(72)

    x = np.linspace(0, 5, 50)
    y = np.linspace(0, 3, 30)
    X, Y = np.meshgrid(x, y)

    Z = 0.8 * np.exp(-((X-2.5)**2 + (Y-1.5)**2) / 1.5)
    Z += np.random.rand(30, 50) * 0.15

    from matplotlib.colors import LinearSegmentedColormap
    colors = ['#27ae60', '#f1c40f', '#e67e22', '#e74c3c']
    cmap = LinearSegmentedColormap.from_list('drought', colors)

    ax.imshow(Z, cmap=cmap, aspect='auto', alpha=0.85)
    ax.axis('off')

    plt.tight_layout(pad=0)
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=100, bbox_inches='tight', pad_inches=0, transparent=True)
    buf.seek(0)
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close(fig)

    return image_base64


def generate_monthly_drought_chart():
    """월별 가뭄 지수 차트 생성"""
    fig, ax = plt.subplots(figsize=(4, 1.2))

    np.random.seed(2025)
    months = np.arange(1, 21)

    # 2025년 데이터
    drought_2025 = 25 + np.sin(months / 3) * 8 + np.random.rand(20) * 5
    # 2024년 데이터
    drought_2024 = 30 + np.sin(months / 3 + 0.5) * 6 + np.random.rand(20) * 4

    ax.plot(months, drought_2025, 'b-', linewidth=1.5, label='2025')
    ax.plot(months, drought_2024, color='#e67e22', linewidth=1.5, linestyle='--', label='2024')

    ax.fill_between(months, drought_2025, alpha=0.2, color='blue')

    ax.set_xlim(1, 20)
    ax.set_ylim(15, 45)
    ax.set_xticks([1, 8, 9, 12, 15, 20])
    ax.set_xticklabels(['1', '8', '9', '12', '15', '20'], fontsize=7)
    ax.tick_params(axis='y', labelsize=7)
    ax.grid(True, alpha=0.3)

    plt.tight_layout(pad=0.5)
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=100, bbox_inches='tight', transparent=True)
    buf.seek(0)
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close(fig)

    return image_base64


@bp.route('/monitoring-heatmap', methods=['GET'])
def monitoring_heatmap():
    """가뭄 모니터링 히트맵 이미지 반환"""
    try:
        heatmap_image = generate_heatmap_image()
        mini_heatmap_image = generate_mini_heatmap_image()

        return jsonify({
            'status': 'success',
            'heatmap_image': heatmap_image,
            'mini_heatmap_image': mini_heatmap_image,
            'risk_score': 72,
            'risk_level': '심각'
        })
    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e)
        })


@bp.route('/monthly-chart', methods=['GET'])
def monthly_chart():
    """월별 가뭄 지수 차트 반환"""
    try:
        image = generate_monthly_drought_chart()
        return jsonify({
            'status': 'success',
            'image': image
        })
    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e),
            'image': None
        })


def generate_spi_chart(period='month'):
    """SPI (표준강수지수) 시계열 그래프 생성"""
    fig, ax = plt.subplots(figsize=(10, 5))

    np.random.seed(2025)

    if period == 'month':
        days = np.arange(1, 31)
        xlabel = '일'
        title = '최근 1개월 SPI 변화'
    elif period == 'quarter':
        days = np.arange(1, 91)
        xlabel = '일'
        title = '최근 3개월 SPI 변화'
    else:  # year
        days = np.arange(1, 366)
        xlabel = '일'
        title = '최근 1년 SPI 변화'

    # SPI 값 생성 (-3 ~ 3 범위)
    spi_values = np.sin(days / 15) * 1.2 + np.random.normal(0, 0.3, len(days))
    spi_values = np.clip(spi_values, -3, 3)

    # 색상별 영역 표시
    ax.axhspan(-3, -2, alpha=0.2, color='darkred', label='극심한 가뭄')
    ax.axhspan(-2, -1.5, alpha=0.2, color='red', label='심한 가뭄')
    ax.axhspan(-1.5, -1, alpha=0.2, color='orange', label='중간 가뭄')
    ax.axhspan(-1, 0, alpha=0.2, color='yellow', label='약한 가뭄')
    ax.axhspan(0, 3, alpha=0.1, color='green', label='정상/습윤')

    ax.plot(days, spi_values, 'b-', linewidth=1.5)
    ax.axhline(y=0, color='black', linestyle='-', linewidth=0.5)

    ax.set_title(title, fontsize=14)
    ax.set_xlabel(xlabel)
    ax.set_ylabel('SPI 값')
    ax.set_ylim(-3, 3)
    ax.legend(loc='upper right', fontsize=8)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
    buf.seek(0)
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close(fig)

    # 통계 계산
    current = round(spi_values[-1], 2)
    avg = round(np.mean(spi_values), 2)
    min_val = round(np.min(spi_values), 2)

    # 가뭄 등급 결정
    if current <= -2:
        grade = '극심한 가뭄'
    elif current <= -1.5:
        grade = '심한 가뭄'
    elif current <= -1:
        grade = '중간 가뭄'
    elif current <= 0:
        grade = '약한 가뭄'
    else:
        grade = '정상'

    return image_base64, {'current': current, 'avg': avg, 'min': min_val, 'grade': grade}


def generate_vhi_chart(period='month'):
    """VHI (식생건강지수) 시계열 그래프 생성"""
    fig, ax = plt.subplots(figsize=(10, 5))

    np.random.seed(2026)

    if period == 'month':
        days = np.arange(1, 31)
        xlabel = '일'
        title = '최근 1개월 VHI 변화'
    elif period == 'quarter':
        days = np.arange(1, 91)
        xlabel = '일'
        title = '최근 3개월 VHI 변화'
    else:  # year
        days = np.arange(1, 366)
        xlabel = '일'
        title = '최근 1년 VHI 변화'

    # VHI 값 생성 (0 ~ 100 범위)
    vhi_values = 50 + np.sin(days / 20) * 25 + np.random.normal(0, 5, len(days))
    vhi_values = np.clip(vhi_values, 0, 100)

    # 색상별 영역 표시
    ax.axhspan(0, 10, alpha=0.2, color='darkred', label='극심한 가뭄')
    ax.axhspan(10, 20, alpha=0.2, color='red', label='심한 가뭄')
    ax.axhspan(20, 30, alpha=0.2, color='orange', label='중간 가뭄')
    ax.axhspan(30, 40, alpha=0.2, color='yellow', label='약한 가뭄')
    ax.axhspan(40, 100, alpha=0.1, color='green', label='정상')

    ax.plot(days, vhi_values, 'g-', linewidth=1.5)

    ax.set_title(title, fontsize=14)
    ax.set_xlabel(xlabel)
    ax.set_ylabel('VHI 값')
    ax.set_ylim(0, 100)
    ax.legend(loc='upper right', fontsize=8)
    ax.grid(True, alpha=0.3)

    plt.tight_layout()
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=100, bbox_inches='tight')
    buf.seek(0)
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close(fig)

    # 통계 계산
    current = round(vhi_values[-1], 1)
    avg = round(np.mean(vhi_values), 1)
    min_val = round(np.min(vhi_values), 1)

    # 가뭄 등급 결정
    if current <= 10:
        grade = '극심한 가뭄'
    elif current <= 20:
        grade = '심한 가뭄'
    elif current <= 30:
        grade = '중간 가뭄'
    elif current <= 40:
        grade = '약한 가뭄'
    else:
        grade = '정상'

    return image_base64, {'current': current, 'avg': avg, 'min': min_val, 'grade': grade}


@bp.route('/graph', methods=['GET'])
def drought_graph():
    """가뭄 지표 그래프 반환"""
    from flask import request

    graph_type = request.args.get('type', 'spi')
    period = request.args.get('period', 'month')

    try:
        if graph_type == 'spi':
            image, stats = generate_spi_chart(period)
        else:  # vhi
            image, stats = generate_vhi_chart(period)

        return jsonify({
            'status': 'success',
            'image': image,
            'stats': stats
        })
    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e),
            'image': None
        })