"""Browser-based login flow: custom login UI → Cognito USER_PASSWORD_AUTH → API key.

Endpoints:
  GET  /auth/login?session_id=<id>   — render custom login page
  POST /auth/authenticate            — authenticate via Cognito InitiateAuth
  GET  /auth/session/<session_id>    — agent polls for the provisioned API key
  GET  /setup                        — onboarding guide page
"""

from __future__ import annotations

import logging
import time
from typing import Any

import boto3
from jose import jwt
from starlette.requests import Request
from starlette.responses import HTMLResponse, JSONResponse

from cube_cloud.auth.api_keys import create_key
from cube_cloud.config import settings

logger = logging.getLogger(__name__)

# In-memory session store: session_id → {"api_key": str, "created_at": float}
_sessions: dict[str, dict[str, Any]] = {}

SESSION_TTL = 300  # 5 minutes


def _cleanup_expired() -> None:
    """Remove sessions older than SESSION_TTL."""
    now = time.time()
    expired = [sid for sid, data in _sessions.items() if now - data["created_at"] > SESSION_TTL]
    for sid in expired:
        del _sessions[sid]


# ---------------------------------------------------------------------------
# Cognito helpers (USER_PASSWORD_AUTH)
# ---------------------------------------------------------------------------


def _initiate_auth(email: str, password: str) -> dict:
    """Call Cognito InitiateAuth with USER_PASSWORD_AUTH flow."""
    client = boto3.client("cognito-idp", region_name=settings.aws_region)
    try:
        resp = client.initiate_auth(
            ClientId=settings.cognito_app_client_id,
            AuthFlow="USER_PASSWORD_AUTH",
            AuthParameters={"USERNAME": email, "PASSWORD": password},
        )
    except client.exceptions.NotAuthorizedException:
        return {"error": "Incorrect email or password."}
    except client.exceptions.UserNotFoundException:
        return {"error": "Incorrect email or password."}
    except client.exceptions.UserNotConfirmedException:
        return {"error": "Account not confirmed. Please contact your administrator."}
    except Exception as e:
        logger.error("InitiateAuth failed: %s", e)
        return {"error": "Authentication failed. Please try again."}

    if resp.get("ChallengeName") == "NEW_PASSWORD_REQUIRED":
        return {"challenge": "NEW_PASSWORD_REQUIRED", "session": resp["Session"]}

    auth_result = resp.get("AuthenticationResult", {})
    return {"id_token": auth_result.get("IdToken", "")}


def _respond_to_challenge(session: str, email: str, new_password: str) -> dict:
    """Respond to NEW_PASSWORD_REQUIRED challenge."""
    client = boto3.client("cognito-idp", region_name=settings.aws_region)
    try:
        resp = client.respond_to_auth_challenge(
            ClientId=settings.cognito_app_client_id,
            ChallengeName="NEW_PASSWORD_REQUIRED",
            Session=session,
            ChallengeResponses={
                "USERNAME": email,
                "NEW_PASSWORD": new_password,
            },
        )
    except client.exceptions.InvalidPasswordException as e:
        msg = str(e)
        # Extract the useful part of the error message
        if ":" in msg:
            msg = msg.split(":", 1)[1].strip()
        return {"error": msg}
    except Exception as e:
        logger.error("RespondToAuthChallenge failed: %s", e)
        return {"error": "Failed to set password. Please try again."}

    auth_result = resp.get("AuthenticationResult", {})
    return {"id_token": auth_result.get("IdToken", "")}


# ---------------------------------------------------------------------------
# Shared CSS design system
# ---------------------------------------------------------------------------

_GLASSMORPHISM_CSS = """\
* { margin: 0; padding: 0; box-sizing: border-box; }

body {
    font-family: -apple-system, BlinkMacSystemFont, 'Inter', 'Segoe UI', Roboto, sans-serif;
    min-height: 100vh;
    display: flex;
    align-items: center;
    justify-content: center;
    color: #fff;
    overflow: hidden;
}

.bg {
    position: fixed;
    inset: 0;
    background-size: cover;
    background-position: center;
    z-index: -2;
}

.bg::after {
    content: '';
    position: absolute;
    inset: 0;
    background: rgba(0, 0, 0, 0.55);
    z-index: -1;
}

.glass-card {
    background: rgba(255, 255, 255, 0.08);
    backdrop-filter: blur(20px);
    -webkit-backdrop-filter: blur(20px);
    border: 1px solid rgba(255, 255, 255, 0.12);
    border-radius: 20px;
    padding: 48px 40px;
    width: 100%;
    max-width: 420px;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3);
}

.logo {
    text-align: center;
    margin-bottom: 32px;
}

.logo h1 {
    font-size: 24px;
    font-weight: 600;
    letter-spacing: -0.5px;
}

.logo span {
    color: #6b8afd;
}

.logo p {
    font-size: 13px;
    color: rgba(255, 255, 255, 0.5);
    margin-top: 4px;
}

h2 {
    font-size: 20px;
    font-weight: 600;
    margin-bottom: 8px;
    letter-spacing: -0.3px;
}

.subtitle {
    font-size: 14px;
    color: rgba(255, 255, 255, 0.5);
    margin-bottom: 28px;
}

.form-group {
    margin-bottom: 18px;
}

.form-group label {
    display: block;
    font-size: 13px;
    font-weight: 500;
    color: rgba(255, 255, 255, 0.7);
    margin-bottom: 6px;
}

.form-group input {
    width: 100%;
    padding: 12px 14px;
    background: rgba(255, 255, 255, 0.06);
    border: 1px solid rgba(255, 255, 255, 0.12);
    border-radius: 10px;
    color: #fff;
    font-size: 15px;
    font-family: inherit;
    outline: none;
    transition: border-color 0.2s, background 0.2s;
}

.form-group input:focus {
    border-color: #6b8afd;
    background: rgba(255, 255, 255, 0.1);
}

.form-group input::placeholder {
    color: rgba(255, 255, 255, 0.3);
}

.btn {
    width: 100%;
    padding: 13px;
    background: #6b8afd;
    color: #fff;
    border: none;
    border-radius: 10px;
    font-size: 15px;
    font-weight: 600;
    font-family: inherit;
    cursor: pointer;
    transition: background 0.2s, box-shadow 0.2s;
    margin-top: 8px;
}

.btn:hover {
    background: #5a7bf0;
    box-shadow: 0 4px 20px rgba(107, 138, 253, 0.35);
}

.error-box {
    background: rgba(255, 59, 48, 0.15);
    border: 1px solid rgba(255, 59, 48, 0.3);
    border-radius: 10px;
    padding: 12px 14px;
    margin-bottom: 20px;
    font-size: 13px;
    color: #ff6b6b;
}

.hint {
    font-size: 12px;
    color: rgba(255, 255, 255, 0.4);
    margin-top: 6px;
}
"""


# ---------------------------------------------------------------------------
# HTML page builders
# ---------------------------------------------------------------------------


def _login_page(session_id: str, error: str = "") -> str:
    error_html = f'<div class="error-box">{error}</div>' if error else ""
    return f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>Sign In — EdgescaleAI Cube</title>
    <style>{_GLASSMORPHISM_CSS}</style>
</head>
<body>
    <div class="bg" style="background-image: url('/static/underground-mining-excavator-operations.png');"></div>
    <div class="glass-card">
        <div class="logo">
            <h1>Edgescale<span>AI</span></h1>
            <p>Cube Platform</p>
        </div>
        {error_html}
        <form method="POST" action="/auth/authenticate">
            <input type="hidden" name="session_id" value="{session_id}">
            <div class="form-group">
                <label for="email">Email</label>
                <input type="email" id="email" name="email" placeholder="you@company.com" required autofocus>
            </div>
            <div class="form-group">
                <label for="password">Password</label>
                <input type="password" id="password" name="password" placeholder="Enter your password" required>
            </div>
            <button type="submit" class="btn">Sign In</button>
        </form>
    </div>
</body>
</html>"""


def _new_password_page(session_id: str, email: str, challenge_session: str, error: str = "") -> str:
    error_html = f'<div class="error-box">{error}</div>' if error else ""
    return f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>Set Your Password — EdgescaleAI Cube</title>
    <style>{_GLASSMORPHISM_CSS}</style>
</head>
<body>
    <div class="bg" style="background-image: url('/static/underground-mining-excavator-operations.png');"></div>
    <div class="glass-card">
        <div class="logo">
            <h1>Edgescale<span>AI</span></h1>
            <p>Cube Platform</p>
        </div>
        <h2>Set Your Password</h2>
        <p class="subtitle">Choose a new password for your account.</p>
        {error_html}
        <form method="POST" action="/auth/authenticate">
            <input type="hidden" name="session_id" value="{session_id}">
            <input type="hidden" name="email" value="{email}">
            <input type="hidden" name="challenge_session" value="{challenge_session}">
            <div class="form-group">
                <label for="new_password">New Password</label>
                <input type="password" id="new_password" name="new_password" placeholder="Choose a strong password" required autofocus>
                <p class="hint">At least 12 characters with uppercase, lowercase, and numbers.</p>
            </div>
            <div class="form-group">
                <label for="confirm_password">Confirm Password</label>
                <input type="password" id="confirm_password" name="confirm_password" placeholder="Confirm your password" required>
            </div>
            <button type="submit" class="btn" onclick="return validatePasswords()">Continue</button>
        </form>
    </div>
    <script>
    function validatePasswords() {{
        var pw = document.getElementById('new_password').value;
        var cpw = document.getElementById('confirm_password').value;
        if (pw !== cpw) {{
            alert('Passwords do not match.');
            return false;
        }}
        return true;
    }}
    </script>
</body>
</html>"""


def _success_page() -> str:
    return f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>Login Successful — EdgescaleAI Cube</title>
    <style>
    {_GLASSMORPHISM_CSS}
    .success-icon {{
        width: 72px;
        height: 72px;
        border-radius: 50%;
        background: rgba(52, 199, 89, 0.15);
        border: 2px solid rgba(52, 199, 89, 0.4);
        display: flex;
        align-items: center;
        justify-content: center;
        margin: 0 auto 24px;
        animation: scaleIn 0.4s cubic-bezier(0.175, 0.885, 0.32, 1.275);
    }}
    .checkmark {{
        width: 32px;
        height: 32px;
        stroke: #34c759;
        stroke-width: 3;
        fill: none;
        stroke-linecap: round;
        stroke-linejoin: round;
    }}
    .checkmark path {{
        stroke-dasharray: 48;
        stroke-dashoffset: 48;
        animation: draw 0.5s 0.3s ease forwards;
    }}
    @keyframes scaleIn {{
        from {{ transform: scale(0); opacity: 0; }}
        to {{ transform: scale(1); opacity: 1; }}
    }}
    @keyframes draw {{
        to {{ stroke-dashoffset: 0; }}
    }}
    .success-text {{
        text-align: center;
    }}
    .success-text h2 {{
        margin-bottom: 12px;
    }}
    .success-text p {{
        color: rgba(255, 255, 255, 0.5);
        font-size: 14px;
        line-height: 1.5;
    }}
    </style>
</head>
<body>
    <div class="bg" style="background-image: url('/static/underground-mining-excavator-operations.png');"></div>
    <div class="glass-card">
        <div class="logo">
            <h1>Edgescale<span>AI</span></h1>
            <p>Cube Platform</p>
        </div>
        <div class="success-icon">
            <svg class="checkmark" viewBox="0 0 32 32">
                <path d="M6 17l7 7 13-13"/>
            </svg>
        </div>
        <div class="success-text">
            <h2>Login Successful!</h2>
            <p>You can close this tab and return to your terminal.</p>
        </div>
    </div>
</body>
</html>"""


def _setup_page_html() -> str:
    return f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>Get Started — EdgescaleAI Cube</title>
    <style>
    {_GLASSMORPHISM_CSS}
    body {{
        align-items: flex-start;
        justify-content: center;
        padding: 60px 20px;
        overflow: auto;
    }}
    .container {{
        width: 100%;
        max-width: 640px;
    }}
    .hero {{
        text-align: center;
        margin-bottom: 48px;
    }}
    .hero h1 {{
        font-size: 32px;
        font-weight: 700;
        letter-spacing: -0.5px;
        margin-bottom: 4px;
    }}
    .hero h1 span {{
        color: #6b8afd;
    }}
    .hero .tagline {{
        font-size: 15px;
        color: rgba(255, 255, 255, 0.5);
        margin-top: 8px;
    }}
    .step-card {{
        background: rgba(255, 255, 255, 0.08);
        backdrop-filter: blur(20px);
        -webkit-backdrop-filter: blur(20px);
        border: 1px solid rgba(255, 255, 255, 0.12);
        border-radius: 16px;
        padding: 28px 28px;
        margin-bottom: 20px;
        box-shadow: 0 4px 24px rgba(0, 0, 0, 0.2);
    }}
    .step-header {{
        display: flex;
        align-items: center;
        gap: 14px;
        margin-bottom: 14px;
    }}
    .step-num {{
        width: 32px;
        height: 32px;
        border-radius: 50%;
        background: rgba(107, 138, 253, 0.2);
        border: 1px solid rgba(107, 138, 253, 0.4);
        display: flex;
        align-items: center;
        justify-content: center;
        font-size: 14px;
        font-weight: 700;
        color: #6b8afd;
        flex-shrink: 0;
    }}
    .step-header h3 {{
        font-size: 16px;
        font-weight: 600;
    }}
    .step-card p {{
        font-size: 14px;
        color: rgba(255, 255, 255, 0.65);
        line-height: 1.6;
        margin-bottom: 8px;
    }}
    .step-card code {{
        background: rgba(107, 138, 253, 0.15);
        padding: 2px 7px;
        border-radius: 5px;
        font-size: 13px;
        color: #93aafe;
    }}
    pre {{
        background: rgba(0, 0, 0, 0.4);
        border: 1px solid rgba(255, 255, 255, 0.08);
        border-radius: 10px;
        padding: 14px 16px;
        overflow-x: auto;
        font-size: 13px;
        color: #c8d2e0;
        margin: 10px 0 4px;
        line-height: 1.5;
    }}
    </style>
</head>
<body>
    <div class="bg" style="background-image: url('/static/automated-guided-vehicles-warehouse-floor.jpg');"></div>
    <div class="container">
        <div class="hero">
            <h1>Edgescale<span>AI</span> Cube</h1>
            <p class="tagline">Setup Guide</p>
        </div>

        <div class="step-card">
            <div class="step-header">
                <div class="step-num">1</div>
                <h3>Install cube-mcp</h3>
            </div>
            <pre>pip install cube-mcp</pre>
            <p>Or with uv:</p>
            <pre>uv pip install cube-mcp</pre>
        </div>

        <div class="step-card">
            <div class="step-header">
                <div class="step-num">2</div>
                <h3>Login</h3>
            </div>
            <p>In Claude Code or your MCP client, run the <code>agent_login_browser</code> tool. This will open your browser for sign-in and automatically save your credentials.</p>
        </div>

        <div class="step-card">
            <div class="step-header">
                <div class="step-num">3</div>
                <h3>Connect to a Cluster</h3>
            </div>
            <p>Run the <code>cube_cluster_login</code> tool with your cluster name:</p>
            <pre>cube_cluster_login("staging-int")</pre>
            <p>This merges credentials into <code>~/.kube/config</code> automatically.</p>
        </div>

        <div class="step-card">
            <div class="step-header">
                <div class="step-num">4</div>
                <h3>Verify</h3>
            </div>
            <pre>kubectl get namespaces</pre>
            <p>You should see the namespaces on your cluster.</p>
        </div>
    </div>
</body>
</html>"""


# ---------------------------------------------------------------------------
# Route handlers
# ---------------------------------------------------------------------------


async def auth_login(request: Request) -> HTMLResponse:
    """Render the custom login page."""
    session_id = request.query_params.get("session_id", "")
    if not session_id:
        return JSONResponse(status_code=400, content={"error": "session_id is required"})

    if not settings.cognito_app_client_id:
        return JSONResponse(
            status_code=503,
            content={"error": "Browser login not configured. Use agent_login with an API key."},
        )

    return HTMLResponse(_login_page(session_id))


async def auth_authenticate(request: Request) -> HTMLResponse:
    """Handle login form POST — authenticate via Cognito USER_PASSWORD_AUTH."""
    form = await request.form()
    email = form.get("email", "")
    password = form.get("password", "")
    session_id = form.get("session_id", "")
    challenge_session = form.get("challenge_session", "")
    new_password = form.get("new_password", "")

    if not session_id:
        return HTMLResponse(_login_page("", error="Invalid session."), status_code=400)

    # Handle NEW_PASSWORD_REQUIRED challenge response
    if challenge_session and new_password:
        result = _respond_to_challenge(challenge_session, email, new_password)
        if result.get("error"):
            return HTMLResponse(
                _new_password_page(session_id, email, challenge_session, error=result["error"])
            )
    else:
        # Initial authentication
        if not email or not password:
            return HTMLResponse(_login_page(session_id, error="Email and password are required."))

        result = _initiate_auth(email, password)
        if result.get("challenge") == "NEW_PASSWORD_REQUIRED":
            return HTMLResponse(_new_password_page(session_id, email, result["session"]))
        if result.get("error"):
            return HTMLResponse(_login_page(session_id, error=result["error"]))

    # Success — extract claims from ID token
    id_token = result.get("id_token", "")
    if not id_token:
        return HTMLResponse(_login_page(session_id, error="Authentication failed. Please try again."))

    try:
        claims = jwt.get_unverified_claims(id_token)
    except Exception as e:
        logger.error("Failed to decode ID token: %s", e)
        return HTMLResponse(_login_page(session_id, error="Invalid token. Please try again."))

    user_email = claims.get("email", "unknown")
    groups = claims.get("cognito:groups", [])
    profile = groups[0] if groups else "admin"

    # Create API key and store for agent polling
    api_key, _ = create_key(profile, description=f"Browser login: {user_email}")

    _cleanup_expired()
    _sessions[session_id] = {"api_key": api_key, "created_at": time.time()}

    return HTMLResponse(_success_page())


async def auth_session(request: Request) -> JSONResponse:
    """Agent polls this to retrieve the provisioned API key."""
    session_id = request.path_params.get("session_id", "")
    if not session_id:
        return JSONResponse(status_code=400, content={"error": "session_id is required"})

    _cleanup_expired()

    data = _sessions.get(session_id)
    if not data:
        return JSONResponse(content={"status": "pending"})

    # One-time use: delete after retrieval
    api_key = data["api_key"]
    del _sessions[session_id]
    return JSONResponse(content={"status": "ready", "api_key": api_key})


async def setup_page(request: Request) -> HTMLResponse:
    """Onboarding guide page with glassmorphism design."""
    return HTMLResponse(_setup_page_html())
