import httpx
import psycopg2
import time

def run_tests():
    print("=" * 60)
    print("          TA-ATS SYSTEM END-TO-END VALIDATION")
    print("=" * 60)

    # 1. Check Gateway and Backend Health
    print("\n[1/4] Checking health endpoints...")
    try:
        r = httpx.get("http://localhost:8000/health", timeout=20.0)
        print(f"  Gateway Health: {r.status_code} - {r.json()}")
    except Exception as e:
        print(f"  Gateway Health Failed: {e}")

    try:
        r = httpx.get("http://localhost:8002/health/", timeout=20.0)
        print(f"  Backend Health: {r.status_code} - {r.json()}")
    except Exception as e:
        print(f"  Backend Health Failed: {e}")

    # 2. Trigger Login
    print("\n[2/4] Triggering login for superuser...")
    email = "kunal.verma@indovisionservices.in"
    password = "Admin@123"
    
    login_data = {
        "email": email,
        "password": password
    }
    
    try:
        r = httpx.post("http://localhost:8000/api/v1/auth/login/", json=login_data, timeout=20.0)
        res_json = r.json()
        print(f"  Status Code: {r.status_code}")
        print(f"  Response Message: {res_json.get('message')}")
        print(f"  Response Data: {res_json.get('data')}")
        mfa_required = res_json.get("data", {}).get("mfa_required", False)
        
        if not mfa_required:
            print("  FAIL: MFA was not required! Check database policy.")
            return
    except Exception as e:
        print(f"  Login Trigger Failed: {e}")
        return

    # 3. Retrieve OTP from PostgreSQL Database
    print("\n[3/4] Connecting to PostgreSQL to fetch generated OTP...")
    time.sleep(1) # Let database commit the record
    try:
        conn = psycopg2.connect(
            host="127.0.0.1",
            port=5432,
            dbname="ats_mcp_db",
            user="postgres",
            password="1234"
        )
        with conn.cursor() as cur:
            cur.execute("SELECT otp_code, otp_created FROM mfa_users WHERE user_key=%s ORDER BY updated_at DESC LIMIT 1;", (email,))
            row = cur.fetchone()
            if row:
                otp_code, created_at = row
                print(f"  Successfully retrieved OTP: {otp_code} (generated at {created_at})")
            else:
                print("  FAIL: No OTP record found in table 'mfa_users'.")
                return
        conn.close()
    except Exception as e:
        print(f"  Database connection failed: {e}")
        return

    # 4. Verify MFA
    print("\n[4/4] Submitting verification code via gateway...")
    verify_data = {
        "email": email,
        "mfa_code": otp_code
    }
    
    try:
        r = httpx.post("http://localhost:8000/api/v1/auth/verify-mfa/", json=verify_data, timeout=20.0)
        print(f"  Status Code: {r.status_code}")
        try:
            res_json = r.json()
            print(f"  Response Message: {res_json.get('message')}")
            tokens = res_json.get("data", {})
            if "access" in tokens and "refresh" in tokens:
                print("\n" + "=" * 60)
                print("  SUCCESS: Both ACCESS and REFRESH tokens received!")
                print(f"  Access Token: {tokens['access'][:50]}...[truncated]")
                print(f"  Refresh Token: {tokens['refresh'][:50]}...[truncated]")
                print("=" * 60)
            else:
                print(f"  FAIL: Token not received! Response: {res_json}")
        except Exception as json_err:
            print(f"  FAIL to parse JSON: {json_err}")
            print(f"  Response Text: {r.text}")
    except Exception as e:
        print(f"  Verification Call Failed: {e}")

if __name__ == "__main__":
    run_tests()