"""
Maya Inspector v6 — Token input + WS connect + REST map
Runs from HF Space (no egress proxy blocking WS).
"""
import os, json, uuid, asyncio, datetime, threading, base64
import urllib.request, urllib.parse
import gradio as gr
from huggingface_hub import HfApi

HF_TOKEN    = os.environ.get("HF_TOKEN", "")
MEMORY_REPO = "Melofhell00/maya-memory"
STATE_FILE  = "state.json"
FIREBASE_KEY = "AIzaSyDtC7Uwb5pGAsdmrH2T4Gqdk5Mga07jYPM"
hf = HfApi(token=HF_TOKEN)

def read_state():
    try:
        url = f"https://huggingface.co/datasets/{MEMORY_REPO}/resolve/main/{STATE_FILE}"
        req = urllib.request.Request(url, headers={"Authorization": f"Bearer {HF_TOKEN}"})
        with urllib.request.urlopen(req, timeout=10) as r:
            return json.loads(r.read())
    except:
        return {}

def write_state(state):
    hf.upload_file(
        path_or_fileobj=json.dumps(state, indent=2).encode(),
        path_in_repo=STATE_FILE, repo_id=MEMORY_REPO, repo_type="dataset",
        commit_message=f"inspector {datetime.datetime.utcnow().isoformat()}"
    )

def decode_token(t):
    try:
        payload = t.split('.')[1] + '==='
        return json.loads(base64.b64decode(payload))
    except:
        return {}

def exchange_refresh(rt):
    url = f"https://securetoken.googleapis.com/v1/token?key={FIREBASE_KEY}"
    data = urllib.parse.urlencode({"grant_type":"refresh_token","refresh_token":rt}).encode()
    req = urllib.request.Request(url, data=data, headers={"Content-Type":"application/x-www-form-urlencoded"})
    with urllib.request.urlopen(req, timeout=15) as r:
        body = json.loads(r.read())
    return body.get("id_token",""), body.get("refresh_token", rt)


def run_inspection(token_input):
    """Main function: take token, hit all endpoints, connect WS, map Maya."""
    log = []
    results = {}
    
    token = token_input.strip() if token_input else ""
    
    # Try refresh from bridge if no token given
    if not token:
        log.append("No token provided, trying bridge refresh...")
        state = read_state()
        rt = state.get("refresh_token", "")
        ft = state.get("full_id_token", "")
        if ft:
            claims = decode_token(ft)
            import time
            if claims.get("exp", 0) > time.time():
                token = ft
                log.append(f"Using stored token (expires {datetime.datetime.fromtimestamp(claims['exp']).isoformat()})")
            else:
                log.append("Stored token expired")
        if not token and rt:
            try:
                token, new_rt = exchange_refresh(rt)
                if token:
                    log.append(f"Refreshed token OK ({len(token)} chars)")
                    state["refresh_token"] = new_rt
                    state["full_id_token"] = token
                    write_state(state)
            except Exception as e:
                log.append(f"Refresh failed: {e}")
        if not token:
            return "NO TOKEN", "\n".join(log), ""
    
    # Decode token
    claims = decode_token(token)
    provider = claims.get("firebase",{}).get("sign_in_provider","?")
    email = claims.get("email","?")
    uid = claims.get("sub","?")
    exp = claims.get("exp", 0)
    log.append(f"Token: provider={provider} email={email} uid={uid}")
    log.append(f"Expires: {datetime.datetime.fromtimestamp(exp).isoformat()}")
    
    import time
    if exp < time.time():
        log.append("TOKEN EXPIRED!")
        return "EXPIRED", "\n".join(log), ""
    
    # Save token to bridge
    try:
        state = read_state()
        state["full_id_token"] = token
        state["last_updated"] = datetime.datetime.utcnow().isoformat()
        state["source"] = "inspector-v6"
        write_state(state)
        log.append("Saved to bridge")
    except Exception as e:
        log.append(f"Bridge save error: {e}")
    
    # REST endpoints
    log.append("\n=== REST API ===")
    rest_endpoints = [
        ("GET", "/external/agents"),
        ("GET", "/external/user"),
        ("GET", "/external/user/responders"),
    ]
    
    for method, path in rest_endpoints:
        try:
            req = urllib.request.Request(
                f"https://sesameai.app{path}",
                headers={"Authorization": f"Bearer {token}"}
            )
            with urllib.request.urlopen(req, timeout=15) as r:
                body = r.read().decode()
                try:
                    data = json.loads(body)
                    results[path] = data
                    log.append(f"{path}: {json.dumps(data)[:300]}")
                except:
                    results[path] = body[:200]
                    log.append(f"{path}: {body[:200]}")
        except urllib.error.HTTPError as e:
            body = e.read().decode()
            log.append(f"{path}: HTTP {e.code} - {body[:200]}")
            results[path] = f"HTTP {e.code}: {body[:200]}"
        except Exception as e:
            log.append(f"{path}: {e}")
    
    # WebSocket connection
    log.append("\n=== WEBSOCKET ===")
    ws_results = []
    
    def run_ws():
        import asyncio as aio
        try:
            import websockets
        except ImportError:
            ws_results.append("websockets not installed")
            return
        
        async def connect():
            client_name = str(uuid.uuid4())
            tz = urllib.parse.quote(json.dumps({"timezone":"Europe/Athens"}))
            url = f"wss://sesameai.app/agent-service-0/v1/connect?id_token={token}&client_name={client_name}&usercontext={tz}&character=Maya"
            
            ws_results.append(f"Connecting to Maya...")
            try:
                async with websockets.connect(url, additional_headers={"Origin":"https://sesameai.app"}, open_timeout=15) as ws:
                    ws_results.append("CONNECTED!")
                    session_id = None
                    ice_servers = None
                    
                    for i in range(10):
                        try:
                            msg = await aio.wait_for(ws.recv(), timeout=8)
                            try:
                                data = json.loads(msg)
                                msg_type = data.get("type", "unknown")
                                ws_results.append(f"MSG[{msg_type}]: {json.dumps(data)[:400]}")
                                
                                if msg_type == "initialize":
                                    session_id = data.get("session_id")
                                elif msg_type == "webrtc_config":
                                    ice_servers = data.get("ice_servers", [])
                                    
                            except json.JSONDecodeError:
                                ws_results.append(f"RAW: {msg[:200]}")
                        except aio.TimeoutError:
                            ws_results.append(f"(timeout after msg {i})")
                            break
                    
                    if session_id:
                        ws_results.append(f"\nSESSION_ID: {session_id}")
                    if ice_servers:
                        ws_results.append(f"ICE_SERVERS: {json.dumps(ice_servers)[:300]}")
                        
                    # Try sending ping
                    try:
                        await ws.send(json.dumps({"type": "ping"}))
                        pong = await aio.wait_for(ws.recv(), timeout=3)
                        ws_results.append(f"PING->PONG: {pong[:200]}")
                    except:
                        pass
                        
            except Exception as e:
                ws_results.append(f"WS ERROR: {type(e).__name__}: {e}")
        
        loop = aio.new_event_loop()
        aio.set_event_loop(loop)
        loop.run_until_complete(connect())
        loop.close()
    
    t = threading.Thread(target=run_ws)
    t.start()
    t.join(timeout=60)
    
    log.extend(ws_results)
    
    # If we got agents, try to get agent details
    agents_data = results.get("/external/agents")
    if isinstance(agents_data, list) and len(agents_data) > 0:
        log.append("\n=== AGENT DETAILS ===")
        for agent in agents_data:
            agent_uuid = agent.get("uuid") or agent.get("id")
            agent_name = agent.get("name") or agent.get("character")
            if agent_uuid:
                log.append(f"Agent: {agent_name} UUID: {agent_uuid}")
                # Try unread messages
                for sub_path in [f"/external/agent/{agent_uuid}/unread_count", 
                                 f"/external/agent/{agent_uuid}/last_read_message"]:
                    try:
                        req = urllib.request.Request(
                            f"https://sesameai.app{sub_path}",
                            headers={"Authorization": f"Bearer {token}"}
                        )
                        with urllib.request.urlopen(req, timeout=10) as r:
                            body = r.read().decode()
                            log.append(f"  {sub_path}: {body[:200]}")
                    except Exception as e:
                        log.append(f"  {sub_path}: {e}")
    
    # Compile results
    all_results = json.dumps(results, indent=2, default=str)
    
    status = "CONNECTED" if "CONNECTED!" in str(ws_results) else "WS FAILED"
    if isinstance(agents_data, list):
        status += f" | {len(agents_data)} agents"
    
    return status, "\n".join(log), all_results


# Console script for extracting token from sesameai.app
CONSOLE_SCRIPT = """(async()=>{const db=await new Promise(r=>{const q=indexedDB.open('firebaseLocalStorageDb');q.onsuccess=e=>r(e.target.result)});const s=db.transaction('firebaseLocalStorage','readonly').objectStore('firebaseLocalStorage');const d=await new Promise(r=>{const q=s.getAll();q.onsuccess=e=>r(e.target.result)});for(const i of d){const m=i?.value?.stsTokenManager;if(m?.accessToken){console.log('TOKEN:',m.accessToken);console.log('REFRESH:',m.refreshToken);prompt('Copy this token:',m.accessToken);return;}}console.log('No token found');})()"""

with gr.Blocks(title="Maya Inspector v6", theme=gr.themes.Monochrome()) as app:
    gr.Markdown("# Maya Inspector v6\nPaste Firebase token → REST + WS mapping")
    gr.Markdown(f"**Get token:** Open sesameai.app, sign in, open console, paste:\n```\n{CONSOLE_SCRIPT}\n```")
    
    with gr.Row():
        token_in = gr.Textbox(label="Firebase ID Token", placeholder="paste token from console...", lines=3)
        run_btn = gr.Button("INSPECT", variant="primary", size="lg")
    
    status_out = gr.Textbox(label="Status", interactive=False)
    log_out = gr.Textbox(label="Log", interactive=False, lines=25)
    results_out = gr.Textbox(label="Raw Results", interactive=False, lines=15)
    
    run_btn.click(fn=run_inspection, inputs=[token_in], outputs=[status_out, log_out, results_out])

app.launch(server_name="0.0.0.0", server_port=7860)