//! Stub extension server — handles LS connections in headless mode.

use crate::proto::wire::{encode_proto_string, encode_varint, extract_proto_string};
use std::io::{BufRead, BufReader, Read, Write};

/// Handle a single connection from the LS to the stub extension server.
///
/// The LS uses Connect RPC (HTTP/1.1, ServerStream) to call ExtensionServerService methods.
/// ALL methods are ServerStream — responses use Connect streaming envelope framing:
///   [0x00 | len(4) | protobuf_data]...  (0+ data messages)
///   [0x02 | len(4) | json_trailer]      (exactly 1 end-of-stream)
///
/// IMPORTANT: `SubscribeToUnifiedStateSyncTopic` is a long-lived stream.
/// If we immediately close it, the LS reconnects in a tight loop and never
/// proceeds to fetch OAuth tokens. We keep subscription connections OPEN.
pub fn stub_handle_connection(
    conn: std::net::TcpStream,
    oauth_token: &str,
    oauth_topic_bytes: &Option<Vec<u8>>,
) {
    let mut reader = BufReader::new(match conn.try_clone() {
        Ok(c) => c,
        Err(_) => return,
    });
    let mut writer = conn;

    // Read the HTTP request line
    let mut request_line = String::new();
    match reader.read_line(&mut request_line) {
        Ok(0) | Err(_) => return,
        _ => {}
    }

    // Extract method path for logging
    let path = request_line
        .split_whitespace()
        .nth(1)
        .unwrap_or("/unknown")
        .to_string();

    // Read headers
    let mut content_len: usize = 0;
    loop {
        let mut line = String::new();
        if reader.read_line(&mut line).unwrap_or(0) == 0 {
            return;
        }
        if line.trim().is_empty() {
            break;
        }
        if line.to_lowercase().starts_with("content-length:") {
            content_len = line
                .split(':')
                .nth(1)
                .and_then(|v| v.trim().parse().ok())
                .unwrap_or(0);
        }
    }

    // Read body
    let mut body = Vec::new();
    if content_len > 0 {
        body.resize(content_len, 0u8);
        if Read::read_exact(&mut reader, &mut body).is_err() {
            return;
        }
    }

    // ─── Long-lived streams ──────────────────────────────────────────────
    // SubscribeToUnifiedStateSyncTopic must stay open — the LS subscribes
    // once and expects updates (OAuth, settings) delivered over this stream.
    // If we close immediately, the LS reconnects in a tight loop (~30/sec).
    if path.contains("SubscribeToUnifiedStateSyncTopic") {
        handle_subscribe_stream(&mut writer, &body, &path, oauth_token, oauth_topic_bytes);
        return;
    }

    // ─── Short-lived methods (everything else) ───────────────────────────
    handle_short_lived(&mut writer, &body, &path, oauth_token);
}

/// Handle the long-lived SubscribeToUnifiedStateSyncTopic stream.
fn handle_subscribe_stream(
    writer: &mut std::net::TcpStream,
    body: &[u8],
    path: &str,
    oauth_token: &str,
    oauth_topic_bytes: &Option<Vec<u8>>,
) {
    // Parse the request body to extract the topic name.
    // Connect envelope: [flag(1)] [len(4)] [proto(N)]
    let proto_body = if body.len() > 5 { &body[5..] } else { body };

    // SubscribeToUnifiedStateSyncTopicRequest { string topic = 1; }
    let mut topic_name = String::new();
    let mut i = 0;
    while i < proto_body.len() {
        let tag_byte = proto_body[i];
        let field_num = tag_byte >> 3;
        let wire_type = tag_byte & 0x07;
        i += 1;
        if wire_type == 2 && i < proto_body.len() {
            let len = proto_body[i] as usize;
            i += 1;
            if i + len <= proto_body.len() {
                if field_num == 1 {
                    topic_name = String::from_utf8_lossy(&proto_body[i..i + len]).to_string();
                }
                i += len;
            } else {
                break;
            }
        } else {
            break;
        }
    }

    eprintln!("[stub-ext] STREAM → {path} topic={topic_name:?}");

    // Build initial_state bytes
    let initial_state_bytes = build_initial_state(&topic_name, oauth_token, oauth_topic_bytes);

    // Helper: wrap protobuf bytes in a Connect data envelope
    let make_envelope = |proto: &[u8]| -> Vec<u8> {
        let mut env = Vec::with_capacity(5 + proto.len());
        env.push(0x00u8); // data flag
        env.extend_from_slice(&(proto.len() as u32).to_be_bytes());
        env.extend_from_slice(proto);
        env
    };

    // Helper: write a chunk
    let send_chunk = |w: &mut std::net::TcpStream, data: &[u8]| -> bool {
        let hdr = format!("{:x}\r\n", data.len());
        w.write_all(hdr.as_bytes()).is_ok()
            && w.write_all(data).is_ok()
            && w.write_all(b"\r\n").is_ok()
            && w.flush().is_ok()
    };

    // Build UnifiedStateSyncUpdate { initial_state = initial_state_bytes }
    let mut initial_proto = Vec::new();
    initial_proto.push(0x0A); // field 1 (initial_state), LEN
    encode_varint(&mut initial_proto, initial_state_bytes.len() as u64);
    initial_proto.extend_from_slice(&initial_state_bytes);

    let initial_env = make_envelope(&initial_proto);

    let header = "HTTP/1.1 200 OK\r\n\
         Content-Type: application/connect+proto\r\n\
         Transfer-Encoding: chunked\r\n\
         \r\n"
        .to_string();
    if writer.write_all(header.as_bytes()).is_err() {
        return;
    }

    if !send_chunk(writer, &initial_env) {
        return;
    }
    eprintln!(
        "[stub-ext] STREAM → sent initial_state ({} bytes)",
        initial_state_bytes.len()
    );

    // Keep the stream alive with periodic valid messages.
    // The LS has a ~10s read timeout on streams. After the initial_state,
    // the LS only accepts AppliedUpdate (field 2 in the oneof).
    // We send an empty AppliedUpdate {} every 5s as keepalive.
    let keepalive_proto: &[u8] = &[0x12, 0x00]; // field 2 (applied_update), LEN=0
    let keepalive_env = make_envelope(keepalive_proto);
    loop {
        std::thread::sleep(std::time::Duration::from_secs(5));
        if !send_chunk(writer, &keepalive_env) {
            break;
        }
    }
}

/// Build the initial_state bytes for a USS topic subscription.
fn build_initial_state(
    topic_name: &str,
    oauth_token: &str,
    oauth_topic_bytes: &Option<Vec<u8>>,
) -> Vec<u8> {
    let mut initial_state_bytes = Vec::new();

    if topic_name == "uss-oauth" {
        if let Some(topic_bytes) = oauth_topic_bytes {
            // Use the exact Topic proto from Antigravity's state.vscdb.
            initial_state_bytes = topic_bytes.clone();
            eprintln!(
                "[stub-ext]   using state.vscdb topic ({} bytes)",
                topic_bytes.len()
            );
        } else if !oauth_token.is_empty() {
            // Manual token fallback — construct OAuthTokenInfo with far-future expiry
            let mut oauth_proto = Vec::new();
            // field 1 (access_token), LEN
            oauth_proto.push(0x0A);
            encode_varint(&mut oauth_proto, oauth_token.len() as u64);
            oauth_proto.extend_from_slice(oauth_token.as_bytes());
            // field 2 (token_type), LEN
            let token_type = b"Bearer";
            oauth_proto.push(0x12);
            encode_varint(&mut oauth_proto, token_type.len() as u64);
            oauth_proto.extend_from_slice(token_type);
            // field 4 (expiry) = Timestamp { seconds = 4_102_444_800 } (year 2099-12-31)
            let mut ts_proto = Vec::new();
            ts_proto.push(0x08); // field 1 (seconds), varint
            encode_varint(&mut ts_proto, 4_102_444_800u64);
            oauth_proto.push(0x22); // field 4 (expiry), LEN
            encode_varint(&mut oauth_proto, ts_proto.len() as u64);
            oauth_proto.extend_from_slice(&ts_proto);

            use base64::Engine;
            let b64_value = base64::engine::general_purpose::STANDARD.encode(&oauth_proto);

            // Build Row { value = b64_value, e_tag = 1 }
            let mut row = Vec::new();
            row.push(0x0A); // field 1 (value), LEN
            encode_varint(&mut row, b64_value.len() as u64);
            row.extend_from_slice(b64_value.as_bytes());
            row.push(0x10); // field 2 (e_tag), varint
            row.push(0x01);

            // Build map entry: { key = "oauthTokenInfoSentinelKey", value = row }
            let key_str = b"oauthTokenInfoSentinelKey";
            let mut map_entry = Vec::new();
            map_entry.push(0x0A); // field 1 (key), LEN
            encode_varint(&mut map_entry, key_str.len() as u64);
            map_entry.extend_from_slice(key_str);
            map_entry.push(0x12); // field 2 (value = Row), LEN
            encode_varint(&mut map_entry, row.len() as u64);
            map_entry.extend_from_slice(&row);

            // Build Topic { data = [map_entry] }
            initial_state_bytes.push(0x0A); // field 1 (data map), LEN
            encode_varint(&mut initial_state_bytes, map_entry.len() as u64);
            initial_state_bytes.extend_from_slice(&map_entry);
        }
    }

    initial_state_bytes
}

/// Handle short-lived extension server methods.
fn handle_short_lived(
    writer: &mut std::net::TcpStream,
    body: &[u8],
    path: &str,
    oauth_token: &str,
) {
    let is_noisy = path.contains("GetChromeDevtoolsMcpUrl")
        || path.contains("FetchMCPAuthToken")
        || path.contains("PushUnifiedStateSyncUpdate");
    if !is_noisy {
        eprintln!("[stub-ext] 200 OK → {path}");
    }

    // Build Connect streaming response body with proper envelope framing.
    let mut envelope = Vec::new();

    if path.contains("GetSecretValue") {
        // Parse request body to extract the key (protobuf: field 1 = key, string)
        let key = extract_proto_string(body, 1).unwrap_or_default();
        eprintln!("[stub-ext]   ← GetSecretValue key={key:?}");

        if !oauth_token.is_empty() {
            // Build protobuf: GetSecretValueResponse { string value = 1 }
            let proto = encode_proto_string(1, oauth_token.as_bytes());
            eprintln!(
                "[stub-ext]   → serving token ({} bytes) for key={key:?}",
                oauth_token.len()
            );

            // Data envelope: flag=0x00, length, data
            envelope.push(0x00u8);
            envelope.extend_from_slice(&(proto.len() as u32).to_be_bytes());
            envelope.extend_from_slice(&proto);
        } else {
            eprintln!("[stub-ext]   ⚠ no OAuth token available for key={key:?}");
        }
    } else if path.contains("StoreSecretValue") {
        // Parse and log what the LS is storing (for debugging)
        let key = extract_proto_string(body, 1).unwrap_or_default();
        let value = extract_proto_string(body, 2).unwrap_or_default();
        let val_preview = if value.len() > 32 {
            format!("{}...", &value[..32])
        } else {
            value
        };
        eprintln!("[stub-ext]   ← StoreSecretValue key={key:?} value={val_preview:?}");
    }

    if path.contains("PushUnifiedStateSyncUpdate") {
        // Unary proto — respond with empty PushUnifiedStateSyncUpdateResponse (0 bytes body)
        let header = "HTTP/1.1 200 OK\r\n\
                      Content-Type: application/proto\r\n\
                      Content-Length: 0\r\n\
                      Connection: close\r\n\
                      \r\n";
        let _ = writer.write_all(header.as_bytes());
        let _ = writer.flush();
        return;
    }

    // End-of-stream envelope: flag=0x02, length=2, data="{}"
    envelope.push(0x02u8);
    envelope.extend_from_slice(&2u32.to_be_bytes());
    envelope.extend_from_slice(b"{}");

    // Respond with 200 OK + Connection: close (one request per connection)
    let header = format!(
        "HTTP/1.1 200 OK\r\n\
         Content-Type: application/connect+proto\r\n\
         Content-Length: {}\r\n\
         Connection: close\r\n\
         \r\n",
        envelope.len()
    );
    let _ = writer.write_all(header.as_bytes());
    let _ = writer.write_all(&envelope);
    let _ = writer.flush();
}
