from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
import requests
import json

app = FastAPI()

# Allow all CORS so browser apps can call this
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

BASE_62_CHAR_SET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"

def base62_to_int(e):
    t = 0
    for n in range(len(e)):
        t = t * 62 + BASE_62_CHAR_SET.index(e[n])
    return t

def get_base_url(token):
    t = token[0]
    n = base62_to_int(token[1]) if t == "A" else base62_to_int(token[1:3])
    server_partition = n
    base_url = f"https://p{server_partition:02d}-sharedstreams.icloud.com/{token}/sharedstreams/"
    return base_url

def get_redirected_base_url(base_url, token):
    url = base_url + "webstream"
    headers = {
        'Origin': 'https://www.icloud.com',
        'Content-Type': 'text/plain',
    }
    data = '{"streamCtag":null}'
    response = requests.post(url, headers=headers, data=data, allow_redirects=False)

    if response.status_code == 330:
        body = response.json()
        new_base_url = f"https://{body['X-Apple-MMe-Host']}/{token}/sharedstreams/"
        return new_base_url
    return base_url

def get_metadata(base_url):
    url = base_url + "webstream"
    headers = {
        'Origin': 'https://www.icloud.com',
        'Content-Type': 'text/plain',
    }
    data = '{"streamCtag":null}'
    response = requests.post(url, headers=headers, data=data)
    return response.json().get("photos", [])

def get_asset_urls(base_url, guids):
    url = base_url + "webasseturls"
    headers = {
        'Origin': 'https://www.icloud.com',
        'Content-Type': 'text/plain',
    }
    data = json.dumps({"photoGuids": guids})
    response = requests.post(url, headers=headers, data=data)
    return response.json().get("items", {})

@app.get("/album/{token}")
def get_album(token: str):
    try:
        base_url = get_base_url(token)
        base_url = get_redirected_base_url(base_url, token)
        metadata = get_metadata(base_url)
        guids = [p["photoGuid"] for p in metadata]
        urls = get_asset_urls(base_url, guids)

        video_list = []
        for photo in metadata:
            for derivative in photo.get("derivatives", {}).values():
                if derivative.get("mediaAssetType") == "video":
                    checksum = derivative.get("checksum")
                    if checksum in urls:
                        video_url = f"https://{urls[checksum]['url_location']}{urls[checksum]['url_path']}"
                        video_list.append({
                            "caption": photo.get("caption", ""),
                            "url": video_url
                        })
        return {"videos": video_list}

    except Exception as e:
        return {"error": str(e)}