import os
import json
import numpy as np
import requests
from urllib.parse import urlparse, urljoin
from pathlib import Path
from pyproj import Transformer
from py3dtiles.tileset.content import B3dm
from py3dtiles.tilers.b3dm.wkb_utils import TriangleSoup

CACHE_DIR = "./_cache_tiles"
os.makedirs(CACHE_DIR, exist_ok=True)

def download_file(url, cache_dir=CACHE_DIR):
    os.makedirs(cache_dir, exist_ok=True)
    filename = os.path.basename(urlparse(url).path)
    local_path = os.path.join(cache_dir, filename)
    if not os.path.exists(local_path):
        print(f"📥 Downloading {url}")
        r = requests.get(url)
        r.raise_for_status()
        with open(local_path, "wb") as f:
            f.write(r.content)
    return local_path

def apply_transform(points, transform):
    n = points.shape[0]
    homo = np.hstack([points, np.ones((n,1))])
    transformed = (transform @ homo.T).T[:, :3]
    return transformed

def extract_positions_from_b3dm(b3dm_bytes, transform=None):
    soup = TriangleSoup.from_glb_bytes(b3dm_bytes)
    positions = soup.get_positions()  # Nx3 numpy
    if transform is not None:
        positions = apply_transform(positions, transform)
    return positions

def load_tileset_json(path_or_url):
    if path_or_url.startswith("http"):
        local_path = download_file(path_or_url)
    else:
        local_path = path_or_url
    with open(local_path, "r", encoding="utf-8") as f:
        js = json.load(f)
    if "geometricError" not in js:
        print("⚠️ tileset.json缺少geometricError，自动补 1000.0")
        js["geometricError"] = 1000.0
    return js, os.path.dirname(local_path), path_or_url if path_or_url.startswith("http") else None

def parse_transform(tile_json):
    # 返回 4x4 np.array
    if "transform" in tile_json:
        return np.array(tile_json["transform"]).reshape(4,4)
    else:
        return np.eye(4)

def recursive_load(tile_json, base_dir, base_url, bbox, transformer, transform_accum=None):
    if transform_accum is None:
        transform_accum = np.eye(4)

    points_list = []
    tile_transform = parse_transform(tile_json)
    transform_total = transform_accum @ tile_transform

    print(f"\n处理 Tile，当前累计 transform 矩阵:\n{transform_total}")
    print(f"Tile 的 bbox (输入): {bbox}")

    content = tile_json.get("content")
    if content:
        uri = content.get("uri")
        if uri and uri.endswith(".b3dm"):
            if base_url is not None:
                b3dm_url = urljoin(base_url + "/", uri)
                b3dm_path = download_file(b3dm_url)
            else:
                b3dm_path = os.path.join(base_dir, uri)
            if os.path.exists(b3dm_path):
                with open(b3dm_path, "rb") as f:
                    b3dm = B3dm.from_array(f.read())
                    positions = extract_positions_from_b3dm(b3dm.glb_bytes, transform_total)
                    print(f"加载 {b3dm_path}，原始顶点数量: {len(positions)}")

                    lon, lat, elev = transformer.transform(
                        positions[:,0], positions[:,1], positions[:,2]
                    )
                    print(f"转换后经度范围: {lon.min():.6f} ~ {lon.max():.6f}")
                    print(f"转换后纬度范围: {lat.min():.6f} ~ {lat.max():.6f}")

                    geo_pts = np.vstack([lon, lat, elev]).T
                    mask = (
                        (geo_pts[:,0] >= bbox[0]) & (geo_pts[:,0] <= bbox[2]) &
                        (geo_pts[:,1] >= bbox[1]) & (geo_pts[:,1] <= bbox[3])
                    )
                    in_bbox_count = np.sum(mask)
                    print(f"bbox内点数量: {in_bbox_count}")

                    points_list.append(geo_pts[mask])

        elif uri and uri.endswith(".json"):
            if base_url is not None:
                child_url = urljoin(base_url + "/", uri)
                child_json, child_base_dir, child_base_url = load_tileset_json(child_url)
            else:
                child_path = os.path.join(base_dir, uri)
                child_json, child_base_dir, child_base_url = load_tileset_json(child_path)
            pts = recursive_load(child_json, child_base_dir, child_base_url, bbox, transformer, transform_total)
            if pts.size > 0:
                points_list.append(pts)

    for child in tile_json.get("children", []):
        pts = recursive_load(child, base_dir, base_url, bbox, transformer, transform_total)
        if pts.size > 0:
            points_list.append(pts)

    if points_list:
        return np.vstack(points_list)
    else:
        return np.empty((0,3))

def load_tileset_all_points(path_or_url, bbox):
    transformer = Transformer.from_crs("EPSG:4979", "EPSG:4326", always_xy=True)
    tileset_json, base_dir, base_url = load_tileset_json(path_or_url)
    points = recursive_load(tileset_json, base_dir, base_url, bbox, transformer)
    print(f"\n总计提取 {len(points)} 个点在区域内")
    return points

if __name__ == "__main__":
    bbox = [116.391, 39.905, 116.405, 39.915]
    model1_url = "http://8.137.54.85:9000/300bdf2b-a150-406e-be63-d28bd29b409f/dszh/1748398014403562192_OUT/B3DM/tileset.json"
    model2_url = "http://8.137.54.85:9000/300bdf2b-a150-406e-be63-d28bd29b409f/dszh/1748325943733189898_OUT/B3DM/tileset.json"

    pts1 = load_tileset_all_points(model1_url, bbox)
    pts2 = load_tileset_all_points(model2_url, bbox)