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

app = FastAPI()

# Set up logging for debugging
logging.basicConfig(level=logging.INFO)

# 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 from the token
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 from iCloud
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 (the "webstream" data)
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 mapping for a list of photoGuids
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 video URLs and include the poster (first frame)
@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 = []
        # Process each photo in the album
        for photo in metadata:
            # Check if the overall photo is a video
            if photo.get("mediaAssetType", "").lower() == "video":
                best_video = None
                best_size = 0
                # Look at each derivative except the poster frame
                derivatives = photo.get("derivatives", {})
                for key, derivative in derivatives.items():
                    if key.lower() == "posterframe":
                        continue
                    try:
                        file_size = int(derivative.get("fileSize", 0))
                    except (ValueError, TypeError):
                        file_size = 0
                    if file_size > best_size:
                        best_size = file_size
                        best_video = derivative

                # Also get the poster frame derivative if available
                poster_url = None
                poster_derivative = photo.get("derivatives", {}).get("PosterFrame")
                if poster_derivative:
                    poster_checksum = poster_derivative.get("checksum")
                    if poster_checksum in asset_urls:
                        poster_info = asset_urls[poster_checksum]
                        poster_url = f"https://{poster_info['url_location']}{poster_info['url_path']}"

                # Build the video URL if found
                if best_video:
                    checksum = best_video.get("checksum")
                    if checksum in asset_urls:
                        url_info = asset_urls[checksum]
                        video_url = f"https://{url_info['url_location']}{url_info['url_path']}"
                        video_list.append({
                            "caption": photo.get("caption", ""),
                            "url": video_url,
                            "poster": poster_url
                        })
                    else:
                        logging.info(f"Checksum {checksum} not found in asset_urls for photo {photo.get('photoGuid')}")
                else:
                    logging.info(f"No video derivative found for photo {photo.get('photoGuid')}")
            else:
                logging.info(f"Photo {photo.get('photoGuid')} is not a video. mediaAssetType: {photo.get('mediaAssetType')}")
        
        return {"videos": video_list}
    except Exception as e:
        return {"error": str(e)}

# Debug endpoint to return raw metadata for further inspection
@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)}