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

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

# Enable CORS for all origins
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

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=15.0)
    return app.state.client

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/"

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:
        try:
            body = resp.json()
            host = body.get("X-Apple-MMe-Host")
            if not host:
                raise ValueError("Missing X-Apple-MMe-Host in 330 response")
            logging.info(f"Redirected to {host}")
            return f"https://{host}/{token}/sharedstreams/"
        except Exception as e:
            logging.error(f"Redirect parsing failed: {e}")
            raise
    elif resp.status_code == 200:
        return base_url
    else:
        resp.raise_for_status()

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":
                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:
                continue

            checksum = best.get("checksum")
            info = asset_map.get(checksum)
            if not info:
                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 or ""
            })

        return {"videos": videos}

    except Exception as e:
        logging.exception("Error in get_album")
        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.exception("Error in get_album_raw")
        return {"error": str(e)}