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 to int using dict lookup
def base62_to_int(token: str) -> int:
    result = 0
    for ch in token:
        result = result * 62 + BASE_62_MAP[ch]
    return result

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

# Static request details
ICLOUD_HEADERS = {'Origin': 'https://www.icloud.com', 'Content-Type': 'text/plain'}
ICLOUD_PAYLOAD = '{"streamCtag":null}'

# Handle possible 330 redirect
async def get_redirected_base_url(base_url: str, token: str) -> str:
    client = await get_client()
    resp = await client.post(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

# Generic function to POST and parse JSON
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()

# Fetch metadata
async def get_metadata(base_url: str) -> list:
    return (await post_json('webstream', base_url, ICLOUD_PAYLOAD)).get('photos', [])

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

# Main endpoint\@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 = [p['photoGuid'] for p in metadata]
        asset_map = await get_asset_urls(base_url, guids)

        videos = []
        for p in metadata:
            if p.get('mediaAssetType', '').lower() != 'video':
                continue

            # select the largest derivative (exclude poster)
            derivatives = p.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')
            if checksum not in asset_map:
                logging.info(f"Missing asset for checksum {checksum}")
                continue

            info = asset_map[checksum]
            video_url = f"https://{info['url_location']}{info['url_path']}"

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

            videos.append({
                'caption': p.get('caption', ''),
                'url': video_url,
                'poster': poster
            })

        return {'videos': videos}

    except Exception as e:
        logging.error(f"Error fetching album {token}: {e}")
        return {'error': str(e)}

# Debug endpoint
@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 = [p['photoGuid'] for p in metadata]
        asset_map = await get_asset_urls(base_url, guids)
        return {'metadata': metadata, 'asset_urls': asset_map}
    except Exception as e:
        return {'error': str(e)}