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

app = FastAPI()

# Set up basic logging
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
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 extensive fallback 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

            # Check derivatives first
            derivatives = photo.get("derivatives", {})
            for key, derivative in derivatives.items():
                # Log the derivative keys for debugging
                logging.info(f"Photo {photo.get('photoGuid')} derivative key: {key} | content: {list(derivative.keys())}")

                # Standard checks
                if (derivative.get("mediaAssetType") == "video" or
                    derivative.get("type", "").startswith("video") or
                    derivative.get("filename", "").lower().endswith(".mp4")):
                    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
                else:
                    # Fallback: Check if any value in the derivative JSON contains "mp4"
                    derivative_str = json.dumps(derivative).lower()
                    if "mp4" in derivative_str:
                        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

            # Fallback to 'movies' field if present
            if not best_video and "movies" in photo:
                for movie in photo["movies"]:
                    logging.info(f"Photo {photo.get('photoGuid')} movie entry keys: {list(movie.keys())}")
                    if movie.get("filename", "").lower().endswith(".mp4"):
                        try:
                            file_size = int(movie.get("fileSize", 0))
                        except (ValueError, TypeError):
                            file_size = 0
                        if file_size > best_size:
                            best_size = file_size
                            best_video = movie

            # 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
                    })
                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')}")
        
        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)}