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

app = FastAPI()
logging.basicConfig(level=logging.INFO)

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

configure_cors(app)

# Precompute Base62 index mapping for O(1) lookup
BASE_62_MAP = {c: i for i, c in enumerate(
    "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
)}

async def get_client() -> httpx.AsyncClient:
    if not hasattr(app.state, "client"):
        app.state.client = httpx.AsyncClient(timeout=10.0)
    return app.state.client

# Convert base62 string to integer
def base62_to_int(token: str) -> int:
    result = 0
    for ch in token:
        result = result * 62 + BASE_62_MAP[ch]
    return result

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

# Common headers and payload for iCloud requests
ICLOUD_HEADERS = {"Origin": "https://www.icloud.com", "Content-Type": "text/plain"}
ICLOUD_PAYLOAD = '{"streamCtag":null}'

async def get_redirected_base_url(base_url: str, token: str) -> str:
    client = await get_client()
    resp = await client.post(
        f"{base_url}webstream", headers=ICLOUD_HEADERS, data=ICLOUD_PAYLOAD, follow_redirects=False
    )
    if resp.status_code == 330:
        body = resp.json()
        host = body.get("X-Apple-MMe-Host")
        return f"https://{host}/{token}/sharedstreams/"
    return base_url

async def post_json(path: str, base_url: str, payload: str) -> dict:
    client = await get_client()
    resp = await client.post(
        f"{base_url}{path}", headers=ICLOUD_HEADERS, data=payload
    )
    resp.raise_for_status()
    return resp.json()

async def get_metadata(base_url: str) -> list:
    data = await post_json("webstream", base_url, ICLOUD_PAYLOAD)
    return data.get("photos", [])

async def get_asset_urls(base_url: str, guids: list) -> dict:
    payload = json.dumps({"photoGuids": guids})
    data = await post_json("webasseturls", base_url, payload)
    return data.get("items", {})

@app.get("/album/{token}")
async def get_album(token: str):
    try:
        base_url = await get_base_url(token)
        base_url = await get_redirected_base_url(base_url, token)

        metadata = await get_metadata(base_url)
        guids = [photo["photoGuid"] for photo in metadata]
        asset_map = await get_asset_urls(base_url, guids)

        videos = []
        for photo in metadata:
            if photo.get("mediaAssetType", "").lower() != "video":
                logging.info(f"Photo {photo.get('photoGuid')} is not a video.")
                continue

            derivatives = photo.get("derivatives", {})
            best = max(
                (d for k, d in derivatives.items() if k.lower() != "posterframe"),
                key=lambda d: int(d.get("fileSize") or 0),
                default=None
            )
            if not best:
                logging.info(f"No video derivative for photo {photo.get('photoGuid')}")
                continue

            checksum = best.get("checksum")
            info = asset_map.get(checksum)
            if not info:
                logging.info(f"Missing asset for checksum {checksum}")
                continue
            video_url = f"https://{info['url_location']}{info['url_path']}"

            poster = None
            pf = derivatives.get("PosterFrame")
            if pf:
                pf_info = asset_map.get(pf.get("checksum"))
                if pf_info:
                    poster = f"https://{pf_info['url_location']}{pf_info['url_path']}"

            videos.append({
                "caption": photo.get("caption", ""),
                "url": video_url,
                "poster": poster
            })

        return {"videos": videos}
    except Exception as e:
        logging.error(f"Error in get_album: {e}")
        return {"error": str(e)}

@app.get("/album/{token}/raw")
async def get_album_raw(token: str):
    try:
        base_url = await get_base_url(token)
        base_url = await get_redirected_base_url(base_url, token)
        metadata = await get_metadata(base_url)
        guids = [photo["photoGuid"] for photo in metadata]
        asset_map = await get_asset_urls(base_url, guids)
        return {"metadata": metadata, "asset_urls": asset_map}
    except Exception as e:
        logging.error(f"Error in get_album_raw: {e}")
        return {"error": str(e)}