import os
import rasterio
import requests
from shapely.geometry import box, mapping
from datetime import datetime, timezone
from pystac import Item, Asset, MediaType, Collection, Extent, SpatialExtent, TemporalExtent
from pyproj import Transformer

# === 配置部分 ===
STAC_API_URL = "http://localhost:8082"
IMG_DIR = "test_geosat1_imgs"
COLLECTION_ID = "geosat1"


def create_stac_item_from_tif(tif_path: str) -> Item:
    with rasterio.open(tif_path) as src:
        bounds = src.bounds
        src_crs = src.crs
        dst_crs = "EPSG:4326"  # WGS84

        # 坐标转换器
        transformer = Transformer.from_crs(src_crs, dst_crs, always_xy=True)
        minx, miny = transformer.transform(bounds.left, bounds.bottom)
        maxx, maxy = transformer.transform(bounds.right, bounds.top)

        bbox = [minx, miny, maxx, maxy]
        geom = mapping(box(minx, miny, maxx, maxy))
        dt = datetime.fromtimestamp(os.path.getmtime(tif_path), timezone.utc)

        item_id = os.path.splitext(os.path.basename(tif_path))[0]
        item = Item(
            id=item_id,
            geometry=geom,
            bbox=bbox,
            datetime=dt,
            properties={},
            stac_extensions=[]
        )

        item.add_asset(
            "image",
            Asset(
                href=os.path.abspath(tif_path),  # 注意路径服务器可访问性
                media_type=MediaType.COG,
                roles=["data"],
                title="GeoTIFF image"
            )
        )

        item.collection_id = COLLECTION_ID
        return item


def create_collection_if_needed():
    url = f"{STAC_API_URL}/collections/{COLLECTION_ID}"
    try:
        r = requests.get(url)
        if r.status_code == 404:
            print(f"Collection `{COLLECTION_ID}` 不存在，正在创建...")
            coll = Collection(
                id=COLLECTION_ID,
                description="Test collection for GeoSat-1 imagery",
                extent=Extent(
                    SpatialExtent([[-180, -90, 180, 90]]),
                    TemporalExtent([[None, None]])
                ),
                license="proprietary"
            )
            r = requests.post(
                f"{STAC_API_URL}/collections", json=coll.to_dict())
            r.raise_for_status()
            print(f"创建成功: {r.status_code}")
        else:
            print(f"Collection `{COLLECTION_ID}` 已存在。")
    except requests.RequestException as e:
        print(f"请求失败: {e}")


def register_item(item: Item):
    url = f"{STAC_API_URL}/collections/{COLLECTION_ID}/items"
    try:
        r = requests.post(url, json=item.to_dict())
        r.raise_for_status()
        print(f"✅ 已成功导入: {item.id}")
    except requests.RequestException as e:
        print(f"❌ 失败: {item.id} | 错误信息: {e}")


def main():
    create_collection_if_needed()

    for filename in os.listdir(IMG_DIR):
        if filename.lower().endswith(".tif"):
            tif_path = os.path.join(IMG_DIR, filename)
            item = create_stac_item_from_tif(tif_path)
            register_item(item)


if __name__ == "__main__":
    main()