#!/usr/bin/env python3
"""
Doormile Full System Test Suite v2
Fixed: correct routes, 2-step login, pricing payload, health endpoint
"""
import asyncio
import aiohttp
import json
import time
import websockets
import nats
import redis.asyncio as aioredis
from datetime import datetime

# ── Config ────────────────────────────────────────────────────────────────────
API_BASE       = "https://api.doormile.com"
NATS_URL       = "nats://66.116.226.161:4223"
NATS_USER      = "doormile"
NATS_PASSWORD  = "Package@321#"
REDIS_HOST     = "66.116.226.255"
REDIS_PORT     = 6380
REDIS_PASSWORD = "Package@321#"
INTERNAL_KEY   = "doormile-internal-2024"

# Test customer — uses the one created by Window 1 test
TEST_PHONE     = "9900000001"
TEST_PIN       = "9999"

# Colors
GREEN  = "\033[92m"
RED    = "\033[91m"
YELLOW = "\033[93m"
BLUE   = "\033[94m"
CYAN   = "\033[96m"
BOLD   = "\033[1m"
RESET  = "\033[0m"

results = []

def log(msg, color=RESET):
    ts = datetime.now().strftime("%H:%M:%S")
    print(f"{color}[{ts}] {msg}{RESET}")

def pass_test(name, detail=""):
    results.append(("PASS", name))
    log(f"✅ PASS — {name} {detail}", GREEN)

def fail_test(name, detail=""):
    results.append(("FAIL", name))
    log(f"❌ FAIL — {name} {detail}", RED)

def info(msg):
    log(f"   ℹ  {msg}", CYAN)

def section(title):
    print(f"\n{BOLD}{YELLOW}{'='*60}{RESET}")
    print(f"{BOLD}{YELLOW}  {title}{RESET}")
    print(f"{BOLD}{YELLOW}{'='*60}{RESET}\n")

# ── Tests ─────────────────────────────────────────────────────────────────────

async def test_health(session):
    section("1. HEALTH CHECK")
    for path in ["/api/v1/health", "/api/v1/ready"]:
        try:
            async with session.get(f"{API_BASE}{path}",
                                   timeout=aiohttp.ClientTimeout(total=10)) as r:
                data = await r.json()
                if r.status == 200:
                    pass_test(f"GET {path}", str(data))
                else:
                    fail_test(f"GET {path}", f"status={r.status}")
        except Exception as e:
            fail_test(f"GET {path}", str(e))

async def login(session):
    section("2. CUSTOMER LOGIN (2-step)")
    token = None
    try:
        # Step 1
        async with session.post(f"{API_BASE}/api/v1/customer/login",
                                json={"phone": TEST_PHONE},
                                timeout=aiohttp.ClientTimeout(total=10)) as r:
            data = await r.json()
            if r.status == 200:
                pass_test("Login step 1 — phone", f"phone={TEST_PHONE}")
                info(str(data))
            else:
                fail_test("Login step 1 — phone", f"status={r.status} {data}")
                return None

        # Step 2
        async with session.post(f"{API_BASE}/api/v1/customer/verify-pin",
                                json={"phone": TEST_PHONE, "pin": TEST_PIN},
                                timeout=aiohttp.ClientTimeout(total=10)) as r:
            data = await r.json()
            if r.status == 200:
                token = data.get("token") or data.get("jwt") or data.get("accesstoken")
                if token:
                    pass_test("Login step 2 — PIN verify", "JWT token obtained")
                    info(f"Token: {token[:50]}...")
                else:
                    fail_test("Login step 2 — PIN verify", f"No token in resp: {data}")
            else:
                fail_test("Login step 2 — PIN verify", f"status={r.status} {data}")
    except Exception as e:
        fail_test("Login", str(e))
    return token

async def test_profile(session, token):
    section("3. CUSTOMER PROFILE")
    try:
        headers = {"Authorization": f"Bearer {token}"}
        async with session.get(f"{API_BASE}/api/v1/customer/profile",
                               headers=headers,
                               timeout=aiohttp.ClientTimeout(total=10)) as r:
            data = await r.json()
            if r.status == 200:
                pass_test("GET /customer/profile", f"name={data.get('firstname','')} {data.get('lastname','')}")
            else:
                fail_test("GET /customer/profile", f"status={r.status} {data}")
    except Exception as e:
        fail_test("Customer profile", str(e))

async def test_pricing(session, token):
    section("4. PRICING CHECK (Redis-backed DoormilePricing)")
    try:
        # Correct payload — zone directly, not pincodes
        payload = {
            "zone": "OtherState",
            "service_type": "Normal",
            "weight": 1.0,
            "itemcategory": "Documents"
        }
        headers = {"Authorization": f"Bearer {token}"}
        async with session.post(f"{API_BASE}/api/v1/pricing/check",
                                json=payload,
                                headers=headers,
                                timeout=aiohttp.ClientTimeout(total=10)) as r:
            data = await r.json()
            if r.status == 200:
                pass_test("Pricing check — OtherState/Normal", f"resp={data}")
            else:
                fail_test("Pricing check", f"status={r.status} resp={data}")

        # Also test Local pricing
        payload2 = {"zone": "Local", "service_type": "Normal",
                    "weight": 0.5, "itemcategory": "Documents"}
        async with session.post(f"{API_BASE}/api/v1/pricing/check",
                                json=payload2,
                                headers=headers,
                                timeout=aiohttp.ClientTimeout(total=10)) as r:
            data = await r.json()
            if r.status == 200:
                pass_test("Pricing check — Local/Normal", f"resp={data}")
            else:
                fail_test("Pricing check — Local/Normal", f"status={r.status} {data}")
    except Exception as e:
        fail_test("Pricing check", str(e))

async def test_city_gate(session, token):
    section("5. CITY GATING — should REJECT Mumbai (400xxx)")
    try:
        payload = {
            "pickupaddress": "MG Road, Mumbai",
            "pickuppincode": "400001",
            "pickuplatitude": 18.9388,
            "pickuplongitude": 72.8354,
            "deliveryaddress": "Hitech City, Hyderabad",
            "deliverypincode": "500032",
            "serviceoption": "Normal",
            "parcels": [{"itemcategory": "Documents",
                         "itemdescription": "test", "weight": 0.5}]
        }
        headers = {"Authorization": f"Bearer {token}"}
        async with session.post(f"{API_BASE}/api/v1/customer/bookings",
                                json=payload, headers=headers,
                                timeout=aiohttp.ClientTimeout(total=10)) as r:
            data = await r.json()
            if r.status == 400:
                pass_test("City gate — Mumbai rejected ✓", str(data))
            else:
                fail_test("City gate — Mumbai should be rejected",
                          f"got status={r.status} — city gate may not be active")
    except Exception as e:
        fail_test("City gate test", str(e))

async def create_booking(session, token):
    section("6. CREATE BOOKING — Coimbatore → Hyderabad")
    booking_id = None
    try:
        payload = {
            "pickupaddress": "Gandhipuram, Coimbatore, Tamil Nadu",
            "pickuppincode": "641001",
            "pickuplatitude": 11.0168,
            "pickuplongitude": 76.9558,
            "deliveryaddress": "Hitech City, Hyderabad, Telangana",
            "deliverypincode": "500032",
            "serviceoption": "Normal",
            "parcels": [
                {
                    "itemcategory": "Documents",
                    "itemdescription": "Doormile E2E test parcel",
                    "weight": 0.5,
                    "length": 20,
                    "width": 15,
                    "height": 5
                }
            ]
        }
        headers = {"Authorization": f"Bearer {token}"}
        async with session.post(f"{API_BASE}/api/v1/customer/bookings",
                                json=payload, headers=headers,
                                timeout=aiohttp.ClientTimeout(total=15)) as r:
            data = await r.json()
            if r.status in [200, 201]:
                # Try common field names for booking id
                booking_id = (data.get("bookingid") or
                              data.get("booking_id") or
                              data.get("id") or
                              data.get("data", {}).get("bookingid"))
                pass_test("Booking created", f"id={booking_id}")
                info(f"Response: {json.dumps(data, indent=2)}")
            else:
                fail_test("Create booking", f"status={r.status} resp={data}")
    except Exception as e:
        fail_test("Create booking", str(e))
    return booking_id

async def test_assignment(session, token, booking_id):
    section("7. AUTO-ASSIGNMENT ENGINE (waiting 12s...)")
    if not booking_id:
        fail_test("Assignment check", "No booking_id from step 6")
        return

    info("Giving assignment engine 12 seconds to fire goroutine...")
    for i in range(12, 0, -3):
        info(f"  {i}s remaining...")
        await asyncio.sleep(3)

    try:
        headers = {"Authorization": f"Bearer {token}"}
        async with session.get(f"{API_BASE}/api/v1/customer/bookings/{booking_id}",
                               headers=headers,
                               timeout=aiohttp.ClientTimeout(total=10)) as r:
            data = await r.json()
            if r.status == 200:
                status = (data.get("status") or
                          data.get("booking", {}).get("status", "unknown"))
                miler  = (data.get("assignedmileruserid") or
                          data.get("booking", {}).get("assignedmileruserid", "none"))
                info(f"Booking status: {status}")
                info(f"Assigned miler: {miler}")
                if status in ["Miler_Assigned", "Pickup_Scheduled"]:
                    pass_test("Auto-assignment fired ✓", f"status={status} miler={miler}")
                elif status == "Created":
                    fail_test("Auto-assignment", "Still 'Created' — check kubectl logs for assignment engine errors")
                else:
                    pass_test("Booking updated", f"status={status}")
            else:
                fail_test("Check booking", f"status={r.status} {data}")
    except Exception as e:
        fail_test("Assignment check", str(e))

async def test_redis(booking_id):
    section("8. REDIS — GEO index, booking cache, pricing cache")
    try:
        r = aioredis.Redis(host=REDIS_HOST, port=REDIS_PORT,
                           password=REDIS_PASSWORD, decode_responses=True)
        await r.ping()
        pass_test("Redis connection", f"{REDIS_HOST}:{REDIS_PORT}")

        # Miler GEO index
        count = await r.zcard("milers:locations")
        if count > 0:
            pass_test("milers:locations GEO index", f"{count} milers indexed")
            # Find milers near Coimbatore
            nearby = await r.georadius(
                "milers:locations", 76.9558, 11.0168, 10, "km",
                withcoord=True, count=5, sort="ASC"
            )
            if nearby:
                pass_test("GEORADIUS near Coimbatore", f"found {len(nearby)} milers")
                for m in nearby:
                    info(f"  Miler: {m}")
            else:
                fail_test("GEORADIUS near Coimbatore", "No milers found in 10km radius")
        else:
            fail_test("milers:locations GEO index", "Empty — milers have no GPS in Redis")

        # Booking cache
        if booking_id:
            data = await r.hgetall(f"bookings:{booking_id}")
            if data:
                pass_test("Booking Redis cache", f"keys={list(data.keys())}")
            else:
                fail_test("Booking Redis cache",
                          f"bookings:{booking_id} not in Redis — NATS worker may not be running")

        # Pricing cache
        keys = await r.keys("doormile:pricing:*")
        if keys:
            pass_test("Pricing Redis cache", f"{len(keys)} slabs cached")
        else:
            fail_test("Pricing Redis cache", "No pricing keys in Redis")

        await r.aclose()
    except Exception as e:
        fail_test("Redis tests", str(e))

async def test_nats(booking_id):
    section("9. NATS JETSTREAM — streams + publish test")
    try:
        nc = await nats.connect(
            servers=[NATS_URL],
            user=NATS_USER,
            password=NATS_PASSWORD
        )
        js = nc.jetstream()
        pass_test("NATS connected", NATS_URL)

        streams = ["BOOKINGS", "TRACKING", "ASSIGNMENTS",
                   "NOTIFICATIONS", "CHAT", "STATUS"]
        for stream in streams:
            try:
                info_obj = await js.stream_info(stream)
                pass_test(f"Stream {stream}",
                          f"msgs={info_obj.state.messages}")
            except Exception as e:
                fail_test(f"Stream {stream}", str(e))

        # Publish test GPS event
        if booking_id:
            evt = json.dumps({
                "miler_id": "test_miler_001",
                "booking_id": booking_id,
                "lat": 11.0175,
                "lon": 76.9565,
                "timestamp": int(time.time())
            })
            ack = await js.publish("miler.location.updated", evt.encode())
            pass_test("Publish miler.location.updated", f"seq={ack.seq}")

        await nc.close()
    except Exception as e:
        fail_test("NATS tests", str(e))

async def test_internal_notify(session, booking_id):
    section("10. INTERNAL NOTIFY API (/api/v1/internal/notify)")
    if not booking_id:
        fail_test("Internal notify", "No booking_id")
        return
    try:
        payload = {
            "booking_id": booking_id,
            "target": "customer",
            "title": "Doormile Test ✅",
            "message": "E2E test notification — system is working!"
        }
        headers = {
            "X-Internal-Key": INTERNAL_KEY,
            "Content-Type": "application/json"
        }
        async with session.post(f"{API_BASE}/api/v1/internal/notify",
                                json=payload, headers=headers,
                                timeout=aiohttp.ClientTimeout(total=10)) as r:
            data = await r.json()
            if r.status == 200:
                pass_test("Internal notify", f"resp={data}")
            else:
                fail_test("Internal notify", f"status={r.status} resp={data}")
    except Exception as e:
        fail_test("Internal notify", str(e))

async def test_reassign(session, booking_id):
    section("11. INTERNAL REASSIGN API")
    if not booking_id:
        fail_test("Internal reassign", "No booking_id")
        return
    try:
        headers = {
            "X-Internal-Key": INTERNAL_KEY,
            "Content-Type": "application/json"
        }
        async with session.post(
                f"{API_BASE}/api/v1/internal/bookings/{booking_id}/reassign",
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=10)) as r:
            data = await r.json()
            if r.status == 200:
                pass_test("Internal reassign", f"resp={data}")
            elif r.status == 400:
                info(f"Reassign returned 400 — booking may not be in assignable state: {data}")
                pass_test("Internal reassign endpoint reachable", f"status=400 (expected if not assigned yet)")
            else:
                fail_test("Internal reassign", f"status={r.status} resp={data}")
    except Exception as e:
        fail_test("Internal reassign", str(e))

async def test_websocket(booking_id):
    section("12. WEBSOCKET LIVE TRACKING")
    if not booking_id:
        fail_test("WebSocket", "No booking_id")
        return
    try:
        ws_url = f"wss://api.doormile.com/ws/bookings/{booking_id}/track"
        info(f"Connecting to {ws_url}")
        async with websockets.connect(ws_url, open_timeout=10) as ws:
            pass_test("WebSocket connected", ws_url)
            try:
                msg = await asyncio.wait_for(ws.recv(), timeout=5)
                data = json.loads(msg)
                pass_test("WebSocket frame received", str(data))
            except asyncio.TimeoutError:
                info("No frame in 5s — miler GPS not streaming yet (normal for test)")
                pass_test("WebSocket connection works", "no GPS frame yet — expected")
    except Exception as e:
        fail_test("WebSocket", str(e))

async def test_booking_cache_api(session, token, booking_id):
    section("13. BOOKING CACHE API (Redis GET endpoint)")
    if not booking_id:
        fail_test("Booking cache API", "No booking_id")
        return
    try:
        headers = {"Authorization": f"Bearer {token}"}
        async with session.get(
                f"{API_BASE}/api/v1/bookings/cache/{booking_id}",
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=10)) as r:
            data = await r.json()
            if r.status == 200:
                pass_test("GET /bookings/cache/:id", f"fields={list(data.keys()) if isinstance(data,dict) else data}")
            else:
                fail_test("GET /bookings/cache/:id", f"status={r.status} {data}")
    except Exception as e:
        fail_test("Booking cache API", str(e))

def print_summary():
    section("FINAL TEST SUMMARY")
    passed = [r for r in results if r[0] == "PASS"]
    failed = [r for r in results if r[0] == "FAIL"]

    for status, name in results:
        icon = "✅" if status == "PASS" else "❌"
        color = GREEN if status == "PASS" else RED
        print(f"  {color}{icon} {name}{RESET}")

    pct = int(len(passed) / len(results) * 100) if results else 0
    print(f"\n{BOLD}{'='*60}{RESET}")
    print(f"{BOLD}  {len(results)} tests  |  "
          f"{GREEN}{len(passed)} passed{RESET}{BOLD}  |  "
          f"{RED}{len(failed)} failed{RESET}{BOLD}  |  "
          f"{YELLOW}{pct}% success{RESET}")
    print(f"{BOLD}{'='*60}{RESET}\n")

    if failed:
        print(f"{RED}Failed tests to investigate:{RESET}")
        for _, name in failed:
            print(f"  • {name}")
        print()

async def main():
    print(f"\n{BOLD}{BLUE}")
    print("╔══════════════════════════════════════════════════════════╗")
    print("║        DOORMILE FULL SYSTEM TEST SUITE v2               ║")
    print("║  Backend · NATS · Redis · Assignment · WebSocket · FCM  ║")
    print("╚══════════════════════════════════════════════════════════╝")
    print(f"{RESET}\n")

    async with aiohttp.ClientSession() as session:
        await test_health(session)
        token = await login(session)
        if not token:
            fail_test("STOPPING — no auth token")
            print_summary()
            return
        await test_profile(session, token)
        await test_pricing(session, token)
        await test_city_gate(session, token)
        booking_id = await create_booking(session, token)
        await test_assignment(session, token, booking_id)
        await test_redis(booking_id)
        await test_nats(booking_id)
        await test_internal_notify(session, booking_id)
        await test_reassign(session, booking_id)
        await test_websocket(booking_id)
        await test_booking_cache_api(session, token, booking_id)

    print_summary()

if __name__ == "__main__":
    asyncio.run(main())