import os
import sys
import time
import json
import uuid
import base64
import threading
import concurrent.futures
import pymysql
import urllib.request
import urllib.error

BASE_URL = "http://127.0.0.1:8080"

def load_env_db_config():
    pw = os.environ.get("DB_PASSWORD", "")
    if not pw and os.path.exists(".env"):
        with open(".env") as f:
            for line in f:
                if line.startswith("DB_PASSWORD="):
                    pw = line.strip().split("=", 1)[1]
    return {
        "host": os.environ.get("DB_HOST", "94.136.190.60"),
        "user": os.environ.get("DB_USER", "chatpingdart_ravi"),
        "password": pw,
        "database": os.environ.get("DB_NAME", "chatpingdart_ucall"),
    }

DB_CONFIG = load_env_db_config()

def http_req(path, method="GET", payload=None, token=None):
    url = f"{BASE_URL}{path}"
    data = json.dumps(payload).encode() if payload is not None else None
    headers = {"Content-Type": "application/json"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=15) as resp:
            return resp.status, json.loads(resp.read().decode())
    except urllib.error.HTTPError as e:
        body = e.read().decode()
        try:
            return e.code, json.loads(body)
        except:
            return e.code, body

def register_user(label):
    phone = f"+1{int(time.time() * 1000 + hash(label)) % 10000000000:010d}"
    st, res = http_req("/api/v1/auth/register", "POST", {"phone": phone, "display_name": label})
    assert st == 200, f"Register failed: {res}"
    return {
        "label": label,
        "user_id": res["data"]["user_id"],
        "token": res["data"]["access_token"],
        "universal_id": res["data"]["universal_id"],
    }

def main():
    print("=" * 75)
    print("UCALL TRANSACTIONAL ROLLBACK, CONCURRENCY & BOUNDARY VERIFICATION")
    print("=" * 75)

    # 1. Setup users
    print("\n[+] Setting up test users (Alice, Bob, Carol)...")
    alice = register_user(f"Alice_{uuid.uuid4().hex[:6]}")
    bob = register_user(f"Bob_{uuid.uuid4().hex[:6]}")
    carol = register_user(f"Carol_{uuid.uuid4().hex[:6]}")

    # Connect to MySQL directly for transaction rollback tests
    conn = pymysql.connect(
        host=DB_CONFIG["host"],
        user=DB_CONFIG["user"],
        password=DB_CONFIG["password"],
        database=DB_CONFIG["database"],
        autocommit=False
    )
    cursor = conn.cursor()

    # -------------------------------------------------------------
    # CASE 1: Transaction allocates sequence but fails / rolls back before insert
    # -------------------------------------------------------------
    print("\n--- CASE 1: Sequence Allocation Rollback Test ---")
    cursor.execute("SELECT last_seq FROM user_mailboxes WHERE user_id = %s", (bob["user_id"],))
    initial_mailbox = cursor.fetchone()
    initial_seq = initial_mailbox[0] if initial_mailbox else 0
    print(f"Initial Bob mailbox last_seq: {initial_seq}")

    # Transaction A: Allocates sequence with ON DUPLICATE KEY UPDATE / FOR UPDATE, then simulates failure & ROLLBACK
    try:
        cursor.execute("""
            INSERT INTO user_mailboxes (user_id, last_seq, updated_at)
            VALUES (%s, 1, NOW(6))
            ON DUPLICATE KEY UPDATE last_seq = last_seq + 1, updated_at = NOW(6)
        """, (bob["user_id"],))
        cursor.execute("SELECT last_seq FROM user_mailboxes WHERE user_id = %s FOR UPDATE", (bob["user_id"],))
        allocated_seq = cursor.fetchone()[0]
        print(f"Transaction A allocated sequence {allocated_seq}, now simulating catastrophic failure & ROLLBACK...")
        # Intentionally roll back without inserting message
        conn.rollback()
        print("Transaction A ROLLED BACK successfully.")
    except Exception as e:
        conn.rollback()
        raise e

    # Verify Bob's last_seq is completely restored to initial_seq
    cursor.execute("SELECT last_seq FROM user_mailboxes WHERE user_id = %s", (bob["user_id"],))
    row_after_rollback = cursor.fetchone()
    seq_after_rollback = row_after_rollback[0] if row_after_rollback else 0
    print(f"Bob mailbox last_seq after rollback: {seq_after_rollback}")
    assert seq_after_rollback == initial_seq, f"Expected {initial_seq} after rollback, got {seq_after_rollback}!"
    print("  ✓ SUCCESS: Sequence counter rolled back cleanly with ZERO ghost increments or gaps.")

    # Now send an actual message via API to confirm sequence begins exactly at initial_seq + 1
    st, res = http_req("/api/v1/messages", "POST", {
        "recipient_universal_id": bob["universal_id"],
        "message_type": "e2ee",
        "payload": json.dumps({"text": "Hello Bob after rollback", "nonce": str(uuid.uuid4())})
    }, token=alice["token"])
    assert st == 200, f"Send message failed: {res}"
    msg_seq = res["data"]["recipient_seq"]
    print(f"First committed message allocated seq: {msg_seq}")
    assert msg_seq == initial_seq + 1, f"Expected seq {initial_seq + 1}, got {msg_seq}"
    print("  ✓ SUCCESS: Subsequent committed message claimed the exact sequence without gap.")
    print(f"First committed message allocated seq: {msg_seq}")
    assert msg_seq == initial_seq + 1, f"Expected seq {initial_seq + 1}, got {msg_seq}"
    print("  ✓ SUCCESS: Subsequent committed message claimed the exact sequence without gap.")

    # -------------------------------------------------------------
    # CASE 2: Two concurrent transactions sending to the same recipient
    # -------------------------------------------------------------
    print("\n--- CASE 2: Concurrent Transactions to Same Recipient ---")
    def send_one(thread_id):
        return http_req("/api/v1/messages", "POST", {
            "recipient_universal_id": bob["universal_id"],
            "message_type": "e2ee",
            "payload": json.dumps({"text": f"Concurrent msg {thread_id}", "nonce": str(uuid.uuid4())})
        }, token=alice["token"])

    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as ex:
        futures = [ex.submit(send_one, i) for i in range(1, 6)]
        results = [f.result() for f in concurrent.futures.as_completed(futures)]

    seqs = sorted([r[1]["data"]["recipient_seq"] for r in results if r[0] == 200])
    print(f"Concurrently allocated sequences: {seqs}")
    # Verify strict contiguity: seqs should be [2, 3, 4, 5, 6] (since initial msg was 1)
    expected_seqs = list(range(2, 7))
    assert seqs == expected_seqs, f"Expected {expected_seqs}, got {seqs}"
    print("  ✓ SUCCESS: Concurrent transactions serialized correctly with strictly contiguous sequences.")

    # -------------------------------------------------------------
    # CASE 3: Client sync while transaction is committing (Isolation check)
    # -------------------------------------------------------------
    print("\n--- CASE 3: Isolation & Sync Boundary During In-Flight Transaction ---")
    # Open uncommitted transaction in raw MySQL
    conn_writer = pymysql.connect(
        host=DB_CONFIG["host"],
        user=DB_CONFIG["user"],
        password=DB_CONFIG["password"],
        database=DB_CONFIG["database"],
        autocommit=False
    )
    c_writer = conn_writer.cursor()
    c_writer.execute("""
        INSERT INTO user_mailboxes (user_id, last_seq, updated_at)
        VALUES (%s, 1, NOW(6))
        ON DUPLICATE KEY UPDATE last_seq = last_seq + 1, updated_at = NOW(6)
    """, (bob["user_id"],))
    c_writer.execute("SELECT last_seq FROM user_mailboxes WHERE user_id = %s FOR UPDATE", (bob["user_id"],))
    uncommitted_seq = c_writer.fetchone()[0]

    dummy_msg_id = str(uuid.uuid4())
    c_writer.execute("""
        INSERT INTO messages (id, sender_user_id, recipient_user_id, message_type, payload, status, created_at, recipient_seq)
        VALUES (%s, %s, %s, 'text', 'Uncommitted shadow', 'sent', NOW(6), %s)
    """, (dummy_msg_id, alice["user_id"], bob["user_id"], uncommitted_seq))

    # While writer has not committed, query /sync API as Bob
    st_sync, res_sync = http_req("/api/v1/messages/sync?since_seq=0&limit=50", "GET", token=bob["token"])
    assert st_sync == 200
    synced_ids = [m["id"] for m in res_sync["data"]["messages"]]
    print(f"Synced message IDs while transaction in-flight: {len(synced_ids)}")
    assert dummy_msg_id not in synced_ids, "Uncommitted message leaked through sync API!"
    print("  ✓ SUCCESS: Reader snapshot isolation preserved: uncommitted message is never returned to client.")

    # Now commit the transaction and sync again
    conn_writer.commit()
    conn_writer.close()

    st_sync2, res_sync2 = http_req(f"/api/v1/messages/sync?since_seq=6&limit=50", "GET", token=bob["token"])
    assert st_sync2 == 200
    synced_ids2 = [m["id"] for m in res_sync2["data"]["messages"]]
    assert dummy_msg_id in synced_ids2, "Committed message missing from post-commit sync!"
    print("  ✓ SUCCESS: Committed message immediately discoverable by client upon reconnect/sync.")

    # -------------------------------------------------------------
    # CASE 4: Page boundary falling between concurrent message submissions
    # -------------------------------------------------------------
    print("\n--- CASE 4: Pagination Across Sequence Boundaries ---")
    # Send 10 more messages
    for i in range(10):
        http_req("/api/v1/messages", "POST", {
            "recipient_universal_id": bob["universal_id"],
            "message_type": "e2ee",
            "payload": json.dumps({"text": f"Page boundary test {i}", "nonce": str(uuid.uuid4())})
        }, token=alice["token"])

    # Now fetch using page size of 4
    all_retrieved_seqs = []
    cursor_seq = 0
    pages = 0
    while True:
        pages += 1
        st_p, res_p = http_req(f"/api/v1/messages/sync?since_seq={cursor_seq}&limit=4", "GET", token=bob["token"])
        assert st_p == 200
        msgs = res_p["data"]["messages"]
        if not msgs:
            break
        for m in msgs:
            all_retrieved_seqs.append(m["recipient_seq"])
        cursor_seq = msgs[-1]["recipient_seq"]
        if len(msgs) < 4:
            break

    print(f"Retrieved {len(all_retrieved_seqs)} messages across {pages} pages with cursor progression: {all_retrieved_seqs}")
    # Verify no duplicates and strictly sorted
    assert len(all_retrieved_seqs) == len(set(all_retrieved_seqs)), "Duplicate sequences detected across page boundaries!"
    assert all_retrieved_seqs == sorted(all_retrieved_seqs), "Out-of-order sequence progression across page boundaries!"
    print("  ✓ SUCCESS: Multi-page cursor traversal guarantees zero skipped or duplicate messages across boundaries.")

    # -------------------------------------------------------------
    # CASE 5: Multiple recipients / devices fanout
    # -------------------------------------------------------------
    print("\n--- CASE 5: Multi-Recipient Mailbox Isolation ---")
    # Alice sends a message to Carol
    st_c, res_c = http_req("/api/v1/messages", "POST", {
        "recipient_universal_id": carol["universal_id"],
        "message_type": "e2ee",
        "payload": json.dumps({"text": "Hello Carol", "nonce": str(uuid.uuid4())})
    }, token=alice["token"])
    assert st_c == 200
    carol_seq = res_c["data"]["recipient_seq"]
    print(f"Carol message allocated seq: {carol_seq}")
    assert carol_seq == 1, f"Carol's independent mailbox should start at seq 1, got {carol_seq}!"

    # Bob's next message should continue Bob's independent counter (which was > 16)
    st_b, res_b = http_req("/api/v1/messages", "POST", {
        "recipient_universal_id": bob["universal_id"],
        "message_type": "e2ee",
        "payload": json.dumps({"text": "Next to Bob", "nonce": str(uuid.uuid4())})
    }, token=alice["token"])
    assert st_b == 200
    bob_seq = res_b["data"]["recipient_seq"]
    print(f"Bob message allocated seq: {bob_seq}")
    assert bob_seq > carol_seq, f"Bob's mailbox should have higher sequence ({bob_seq}) than Carol's ({carol_seq})"
    print("  ✓ SUCCESS: Mailbox sequences are strictly partitioned per recipient without cross-talk.")
    print(f"Bob message allocated seq: {bob_seq}")
    assert bob_seq > carol_seq, f"Bob's mailbox should have higher sequence ({bob_seq}) than Carol's ({carol_seq})"
    print("  ✓ SUCCESS: Mailbox sequences are strictly partitioned per recipient without cross-talk.")

    conn.close()
    print("\n" + "=" * 75)
    print("ALL 5 TRANSACTION & CURSOR BOUNDARY TESTS PASSED PERFECTLY!")
    print("=" * 75)

if __name__ == "__main__":
    main()