import xarray as xr
import numpy as np
from datetime import datetime
import pandas as pd
from typing import Union

from config import WIND_JSON, SIM
from logger import get_logger

logger = get_logger("createWindJson")


def extract_wind_data_json(wind_ds_or_path: Union[str, xr.Dataset],
                           target_time, center_lon: float, center_lat: float) -> list:
    """
    NetCDF 파일 또는 이미 열린 Dataset에서 특정 시간의 바람 데이터를 추출하고 JSON 형식으로 반환

    Parameters
    ----------
    wind_ds_or_path : str or xr.Dataset
        NetCDF 파일 경로(str) 또는 이미 열린 xr.Dataset 객체.
        Dataset 객체가 전달되면 내부에서 닫지 않습니다.
    target_time : str or datetime or pd.Timestamp
        추출할 시간 (예: '2024-01-15 12:00:00' 또는 datetime 객체)
    center_lon : float
        중심 경도
    center_lat : float
        중심 위도

    Returns
    -------
    list
        위도, 경도, 풍향, 풍속 정보가 포함된 리스트
    """
    range_km = WIND_JSON.RANGE_KM
    grid_spacing_km = WIND_JSON.GRID_SPACING_KM

    own_ds = isinstance(wind_ds_or_path, str)
    if own_ds:
        ds = xr.open_dataset(wind_ds_or_path)
    else:
        ds = wind_ds_or_path

    try:
        ds_time = ds.sel(time=target_time, method='nearest')

        lat_per_km = 1 / SIM.KM_PER_DEG_LAT
        lon_per_km = 1 / (SIM.KM_PER_DEG_LAT * np.cos(np.radians(center_lat)))

        lat_range = range_km * lat_per_km
        lon_range = range_km * lon_per_km

        lat_min = center_lat - lat_range
        lat_max = center_lat + lat_range
        lon_min = center_lon - lon_range
        lon_max = center_lon + lon_range

        ds_subset = ds_time.sel(
            lat=slice(lat_min, lat_max),
            lon=slice(lon_min, lon_max)
        )

        n_points = int(2 * range_km / grid_spacing_km) + 1
        new_lats = np.linspace(lat_min, lat_max, n_points)
        new_lons = np.linspace(lon_min, lon_max, n_points)

        ds_interp = ds_subset.interp(
            lat=new_lats,
            lon=new_lons,
            method='linear'
        )

        u = ds_interp['x_wind'].values
        v = ds_interp['y_wind'].values

        wind_speed = np.sqrt(u**2 + v**2)
        wind_direction = (270 - np.arctan2(v, u) * 180 / np.pi) % 360

        data = []
        for i, lat in enumerate(new_lats):
            for j, lon in enumerate(new_lons):
                ws = float(wind_speed[i, j]) if not np.isnan(wind_speed[i, j]) else None
                wd = float(wind_direction[i, j]) if not np.isnan(wind_direction[i, j]) else None
                data.append({
                    "lat": round(float(lat), 6),
                    "lon": round(float(lon), 6),
                    "wind_speed":      round(ws, 2) if ws is not None else None,
                    "wind_direction":  round(wd, 1) if wd is not None else None,
                })

        return data

    finally:
        if own_ds:
            ds.close()