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 API
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Base62 conversion helper
BASE_62_CHAR_SET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
def base62_to_int(e):
    t = 0
    for char in e:
        t = t * 62 + BASE_62_CHAR_SET.index(char)
    return t

# Construct the base URL for the shared album
def get_base_url(token: str) -> str:
    first_char = token[0]
    n = base62_to_int(token[1]) if first_char == "A" else base62_to_int(token[1:3])
    base_url = f"https://p{n:02d}-sharedstreams.icloud.com/{token}/sharedstreams/"
    return base_url

# Handle 330 redirect by calling the webstream endpoint with allow_redirects=False
def get_redirected_base_url(base_url: str, token: str) -> str:
    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

# Fetch metadata from the webstream endpoint
def get_metadata(base_url: str):
    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", [])

# Fetch asset URLs for a list of photoGuids from the webasseturls endpoint
def get_asset_urls(base_url: str, guids: list):
    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", {})

# Main endpoint to extract videos with enhanced checks
@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 = [photo["photoGuid"] for photo in metadata]
        asset_urls = get_asset_urls(base_url, guids)

        video_list = []
        # Loop over each photo in the album
        for photo in metadata:
            best_video = None
            best_size = 0
            # First, check the derivatives field
            for derivative in photo.get("derivatives", {}).values():
                if (
                    derivative.get("mediaAssetType") == "video" or
                    derivative.get("type", "").startswith("video") or
                    derivative.get("filename", "").endswith(".mp4")
                ):
                    try:
                        file_size = int(derivative.get("fileSize", 0))
                    except ValueError:
                        file_size = 0
                    if file_size > best_size:
                        best_size = file_size
                        best_video = derivative
            # If no derivative qualifies, check a potential 'movies' field
            if not best_video and "movies" in photo:
                for movie in photo["movies"]:
                    if movie.get("filename", "").endswith(".mp4"):
                        try:
                            file_size = int(movie.get("fileSize", 0))
                        except ValueError:
                            file_size = 0
                        if file_size > best_size:
                            best_size = file_size
                            best_video = movie
            # Build the video URL if found and if its checksum is in the asset_urls
            if best_video and best_video.get("checksum") in asset_urls:
                url_info = asset_urls[best_video["checksum"]]
                video_url = f"https://{url_info['url_location']}{url_info['url_path']}"
                video_list.append({
                    "caption": photo.get("caption", ""),
                    "url": video_url
                })

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

# Debug endpoint to return raw metadata from iCloud
@app.get("/album/{token}/raw")
def get_album_raw(token: str):
    try:
        base_url = get_base_url(token)
        base_url = get_redirected_base_url(base_url, token)
        metadata = get_metadata(base_url)
        guids = [photo["photoGuid"] for photo in metadata]
        asset_urls = get_asset_urls(base_url, guids)
        return {"metadata": metadata, "asset_urls": asset_urls}
    except Exception as e:
        return {"error": str(e)}