"""
Core Authenticator for Upstox Automation.
Handles headless login and token retrieval.
"""

import asyncio
import datetime
import logging
from typing import Dict, Optional
from urllib.parse import parse_qs, urlencode, urlparse

import pyotp
import pytz
import requests
from playwright.async_api import async_playwright

from .exceptions import InvalidCredentialsError, LoginTimeoutError, UpstoxAuthError
from .stealth import get_stealth_config
from .storage import BaseStorage, MemoryStorage

logger = logging.getLogger(__name__)


class UpstoxAuthenticator:
    """
    Automates the Upstox login process to retrieve an access token.
    """

    def __init__(
        self,
        api_key: str,
        api_secret: str,
        redirect_uri: str,
        mobile_no: str,
        pin: str,
        totp_key: str,
        storage: Optional[BaseStorage] = None,
        headless: bool = True,
        retries: int = 2,
    ):
        """
        Initialize the Upstox Authenticator.

        Args:
            api_key (str): Your Upstox API Key (from developer.upstox.com).
            api_secret (str): Your Upstox API Secret (from developer.upstox.com).
            redirect_uri (str): Your app's Redirect URI.
                Must match what's configured in the Upstox developer portal.
            mobile_no (str): Your registered Upstox mobile number.
            pin (str): Your 6-digit login PIN.
            totp_key (str): Your 2FA TOTP secret key from account settings.
            storage (BaseStorage, optional): Storage provider for tokens.
                Defaults to MemoryStorage.
            headless (bool, optional): Whether to run the browser in headless mode.
                Defaults to True.
            retries (int, optional): Number of times to retry automation on failure.
                Defaults to 2.
        """
        self.api_key = api_key
        self.api_secret = api_secret
        self.redirect_uri = redirect_uri
        self.mobile_no = mobile_no
        self.pin = pin
        self.totp_key = totp_key
        self.storage = storage or MemoryStorage()
        self.headless = headless
        self.retries = retries

    def generate_auth_url(self) -> str:
        """
        Generates the official Upstox authorization URL.

        This is the URL that the browser must visit to start the OAuth2 flow.
        """
        query_params = {
            "response_type": "code",
            "client_id": self.api_key,
            "redirect_uri": self.redirect_uri,
        }
        return (
            f"https://api.upstox.com/v2/login/authorization/dialog"
            f"?{urlencode(query_params)}"
        )

    async def get_auth_code(self, timeout_seconds: int = 60) -> str:
        """
        Runs Playwright automation to login and capture the auth code.

        This method automates the entire browser interaction: entering mobile number,
        requesting OTP, generating and entering TOTP, and entering the PIN.

        Returns:
            str: The authorization code captured from the redirect URL.

        Raises:
            LoginTimeoutError: If the automation or redirect takes too long.
            InvalidCredentialsError: If login fails due to wrong credentials.
            UpstoxAuthError: If any other automation error occurs.
        """
        auth_url = self.generate_auth_url()
        stealth = get_stealth_config()

        async with async_playwright() as p:
            browser = await p.chromium.launch(
                headless=self.headless, args=stealth["args"]
            )

            context = await browser.new_context(
                user_agent=stealth["user_agent"],
                viewport={"width": 1280, "height": 720},
            )

            await context.add_init_script(stealth["init_script"])
            page = await context.new_page()

            try:
                # 1. Navigate to Login
                await page.goto(auth_url, wait_until="domcontentloaded")

                # 2. Enter Mobile Number
                try:
                    await page.wait_for_selector("#mobileNum", timeout=15000)
                except Exception:
                    # Capture screenshot on failure for debugging
                    import os
                    os.makedirs(".tmp", exist_ok=True)
                    await page.screenshot(path=".tmp/error_mobile_selector.png")
                    content = await page.content()
                    title = await page.title()
                    logger.error(f"Failed to find #mobileNum. Page title: {title}")
                    if "zscaler" in content.lower() or "blocked" in content.lower():
                        raise UpstoxAuthError(
                            "Access blocked by network filter (e.g., Zscaler/Firewall)."
                        )
                    raise

                await page.fill("#mobileNum", self.mobile_no)
                await page.click("button:has-text('Get OTP')")

                # Detect immediate mobile number errors
                error_msg = await page.get_by_text("invalid", exact=False).is_visible()
                if error_msg:
                    raise InvalidCredentialsError("Invalid mobile number provided.")

                # 3. Enter TOTP
                await page.wait_for_selector("#otpNum", timeout=15000)
                otp = pyotp.TOTP(self.totp_key).now()
                await page.fill("#otpNum", otp)
                await page.click("button:has-text('Continue')")

                # 4. Enter PIN
                await page.wait_for_selector("input[type='password']", timeout=15000)
                # Upstox sometimes has 6 individual digit inputs or one single input
                pin_inputs = await page.locator("input[type='password']").all()
                if len(pin_inputs) >= len(self.pin):
                    for i, digit in enumerate(self.pin):
                        await pin_inputs[i].fill(digit)
                else:
                    await pin_inputs[0].fill(self.pin)
                await page.click("button:has-text('Continue')")

                # 5. Capture Redirect Code using wait_for_url
                try:
                    await page.wait_for_url(
                        lambda url: "code=" in url, timeout=timeout_seconds * 1000
                    )
                    parsed = urlparse(page.url)
                    code = parse_qs(parsed.query).get("code")
                    if code:
                        return code[0]
                except Exception:
                    # Check for generic error messages on the page before timing out
                    content = await page.content()
                    if any(x in content.lower() for x in ["incorrect", "invalid"]):
                        raise InvalidCredentialsError(
                            "Authentication failed. Check your PIN or TOTP key."
                        )
                    raise LoginTimeoutError(
                        "Timed out waiting for authorization redirect code."
                    )

                raise UpstoxAuthError("Failed to extract authorization code from URL.")

            except (LoginTimeoutError, InvalidCredentialsError):
                raise
            except Exception as e:
                if "timeout" in str(e).lower():
                    raise LoginTimeoutError(f"Automation step timed out: {e}")
                raise UpstoxAuthError(f"Automation failed: {e}")
            finally:
                await browser.close()

    def exchange_code_for_token(self, code: str) -> Dict:
        """
        Exchanges the captured auth code for a full access token.

        This calls the official Upstox API `/token` endpoint.

        Args:
            code (str): The authorization code from get_auth_code().

        Returns:
            Dict: The full JSON response from Upstox (contains access_token).
        """
        url = "https://api.upstox.com/v2/login/authorization/token"
        headers = {
            "accept": "application/json",
            "Content-Type": "application/x-www-form-urlencoded",
        }
        data = {
            "code": code,
            "client_id": self.api_key,
            "client_secret": self.api_secret,
            "redirect_uri": self.redirect_uri,
            "grant_type": "authorization_code",
        }

        try:
            response = requests.post(url, headers=headers, data=data, timeout=30)
        except requests.RequestException as exc:
            raise UpstoxAuthError(f"Token exchange request failed: {exc}") from exc
        if response.status_code != 200:
            raise UpstoxAuthError(f"Token exchange failed: {response.text}")

        token_data = response.json()
        self.storage.save_token(token_data)
        return token_data

    async def get_access_token(self, force_refresh: bool = False) -> str:
        """
        Returns a valid access token, automating login only if necessary.

        Workflow:
        1. Checks storage for an existing token.
        2. Validates if the token is still valid (not expired).
        3. If missing, expired, or `force_refresh` is True, runs Playwright
           automation with retries.
        4. Exchanges code for a new token and saves it to storage.

        Args:
            force_refresh (bool, optional): If True, bypasses storage and runs
                automation.

        Returns:
            str: The valid Upstox Access Token.
        """
        if not force_refresh:
            existing = self.storage.get_token()
            if existing and "access_token" in existing:
                # Upstox tokens expire daily at 3:30 AM IST.
                # We check if the token was fetched today (IST).
                ist = pytz.timezone("Asia/Kolkata")
                now_ist = datetime.datetime.now(ist)

                # Simple but effective check:
                # If a 'fetched_at' is not in storage, we assume it might be old.
                fetched_at = existing.get("fetched_at")
                if fetched_at:
                    fetched_dt = datetime.datetime.fromisoformat(fetched_at)
                    if fetched_dt.date() == now_ist.date():
                        # If fetched today, it's a safe approximation for daily tokens.
                        return existing["access_token"]

        # Run automation with retries
        last_err = None
        for attempt in range(self.retries + 1):
            try:
                if attempt > 0:
                    logger.info(
                        f"Retrying automation (attempt {attempt}/{self.retries})..."
                    )
                code = await self.get_auth_code()
                token_data = self.exchange_code_for_token(code)
                return token_data["access_token"]
            except (InvalidCredentialsError, UpstoxAuthError) as e:
                last_err = e
                if isinstance(e, InvalidCredentialsError):
                    # Don't retry on wrong password/PIN
                    raise
                await asyncio.sleep(2**attempt)  # Exponential backoff

        msg = "Failed to retrieve access token after retries."
        raise last_err or UpstoxAuthError(msg)
