"""
Session management module for maintaining Schwab authenticated sessions.
This module provides functionality to refresh session state through browser navigation
without requiring 2FA approval for active sessions.
"""

import json
import logging
import time
from typing import List, Dict, Any, Optional
from datetime import datetime

from playwright.async_api import async_playwright
from ..core.config import load_config, get_playwright_url, get_cookies_path
from .client import new_context, new_page
from ..core import ErrorType, Envelope, fail, ok


async def refresh_session_state(cookies: Optional[List[Dict[str, Any]]] = None) -> bool:
    """
    Refresh session state through browser navigation.
    
    This function maintains active sessions by navigating to a Schwab page,
    which updates cookie expiration times and session state without requiring
    2FA approval for active sessions.
    
    Args:
        cookies: Optional list of cookies to use. If None, loads from cookies.json
        
    Returns:
        bool: True if session refresh was successful, False otherwise
    """
    logger = logging.getLogger(__name__)
    
    try:
        logger.info("Starting session refresh through navigation")
        
        # Load cookies if not provided
        if cookies is None:
            cookies_path = get_cookies_path()
            try:
                with open(cookies_path, 'r') as f:
                    cookies = json.load(f)
                logger.info(f"Loaded {len(cookies) if cookies else 0} cookies from {cookies_path}")
            except (FileNotFoundError, json.JSONDecodeError) as e:
                logger.error(f"Could not load cookies: {e}")
                return False
        
        if not cookies:
            logger.error("No cookies available for session refresh")
            return False
        
        config = load_config()
        playwright_url = get_playwright_url(config)
        
        async with async_playwright() as p:
            try:
                browser = await p.chromium.connect(playwright_url)
            except Exception as e:
                logger.error(f"Failed to connect to browser: {e}")
                return False
            
            try:
                # Create context with existing cookies
                context = await new_context(browser, cookies=cookies)
                page = await new_page(context)
                
                # Navigate to refresh session state
                logger.info("Navigating to Schwab research page to refresh session")
                await page.goto("https://client.schwab.com/app/research/#/stocks/AAPL", timeout=30000)
                await page.wait_for_timeout(2000)  # Let page settle and cookies update
                
                # Check if navigation was successful (no redirect to login)
                current_url = page.url
                is_redirected = any(pattern in current_url for pattern in [
                    '/login', '/signin', '/auth', '/Access/'
                ])
                
                if is_redirected:
                    logger.warning(f"Session refresh failed: redirected to login page")
                    logger.debug(f"Current URL: {current_url}")
                    await context.close()
                    await browser.close()
                    return False
                
                # Get updated cookies after navigation
                new_cookies = await context.cookies()
                logger.info(f"Retrieved {len(new_cookies)} cookies after navigation")
                
                # Check if we still have critical session cookies
                critical_session_cookies = ['LVAL', 'NS2', 'sstate']
                missing_critical_cookies = []
                
                for cookie_name in critical_session_cookies:
                    old_cookie = next((c for c in cookies if c['name'] == cookie_name), None)
                    new_cookie = next((c for c in new_cookies if c['name'] == cookie_name), None)
                    
                    if not new_cookie:
                        missing_critical_cookies.append(cookie_name)
                    elif old_cookie and new_cookie.get('expires') != -1:
                        # Session cookies should have expires = -1
                        missing_critical_cookies.append(f"{cookie_name} (invalid session cookie)")
                
                if missing_critical_cookies:
                    logger.warning(f"Session refresh failed: missing critical session cookies: {missing_critical_cookies}")
                    await context.close()
                    await browser.close()
                    return False
                
                # Compare cookie states to detect changes
                changes = []
                old_dict = {c['name']: c for c in cookies}
                new_dict = {c['name']: c for c in new_cookies}
                
                # Check for modified cookies (especially expiration changes)
                for name in old_dict:
                    if name in new_dict:
                        old_cookie = old_dict[name]
                        new_cookie = new_dict[name]
                        
                        # Check if expiration changed
                        old_expires = old_cookie.get('expires', -1)
                        new_expires = new_cookie.get('expires', -1)
                        if old_expires != new_expires:
                            changes.append({
                                'type': 'expiration_changed',
                                'name': name,
                                'old_expires': old_expires,
                                'new_expires': new_expires
                            })
                
                if changes:
                    logger.info(f"Detected {len(changes)} cookie changes (session refreshed)")
                    for change in changes[:3]:  # Show first 3
                        logger.debug(f"  {change['name']}: expiration updated")
                else:
                    logger.info("No cookie changes detected (session maintained)")
                
                # Save updated cookies
                cookies_path = get_cookies_path()
                with open(cookies_path, 'w') as f:
                    json.dump(new_cookies, f, indent=2)
                logger.info(f"Saved {len(new_cookies)} updated cookies")
                
                await context.close()
                await browser.close()
                
                return True
                
            except Exception as e:
                logger.error(f"Error during session refresh: {e}")
                try:
                    await context.close()
                except:
                    pass
                await browser.close()
                return False
                
    except Exception as e:
        logger.error(f"Session refresh failed: {e}")
        return False


async def maintain_session_health() -> bool:
    """
    Check if the current session is healthy by attempting a simple navigation.
    
    Returns:
        bool: True if session is healthy, False if refresh is needed
    """
    logger = logging.getLogger(__name__)
    
    try:
        logger.info("Checking session health")
        
        # Load current cookies
        cookies_path = get_cookies_path()
        try:
            with open(cookies_path, 'r') as f:
                cookies = json.load(f)
        except (FileNotFoundError, json.JSONDecodeError):
            logger.error("No valid cookies found")
            return False
        
        if not cookies:
            logger.error("No cookies available")
            return False
        
        # First, check if we have valid session cookies (basic check)
        current_time = int(time.time())
        has_valid_session_cookies = False
        
        for cookie in cookies:
            name = cookie.get('name', '')
            expires = cookie.get('expires', -1)
            
            # Check for actual Schwab session cookies
            if name in ['auth', 'ASP.NET_SessionId', 'SessionInfo', '__RequestVerificationToken']:
                # Session cookies (expires=-1) are valid until browser closes
                # Other cookies must not be expired
                if expires == -1 or (expires and expires > current_time):
                    has_valid_session_cookies = True
                    break
        
        if not has_valid_session_cookies:
            logger.warning("Session health check: FAILED - no valid session cookies found")
            return False
        
        config = load_config()
        playwright_url = get_playwright_url(config)
        
        async with async_playwright() as p:
            browser = await p.chromium.connect(playwright_url)
            
            try:
                context = await new_context(browser, cookies=cookies)
                page = await new_page(context)
                
                # Navigate to a simple page to test session
                await page.goto("https://client.schwab.com/app/research/#/stocks/AAPL", timeout=30000)
                
                # Check if we're still authenticated by URL pattern
                current_url = page.url
                logger.debug(f"Current URL after navigation: {current_url}")
                
                is_authenticated_by_url = any(pattern in current_url for pattern in [
                    '/app/', '/Apps/', '/accounts/', '/Areas/Accounts', '/summary'
                ])
                
                # Check for login redirect patterns
                is_redirected = any(pattern in current_url for pattern in [
                    '/login', '/signin', '/auth', '/Access/'
                ])
                
                logger.debug(f"Authenticated by URL pattern: {is_authenticated_by_url}")
                logger.debug(f"Redirected to login: {is_redirected}")
                
                # Primary check: If we're not redirected and have a good URL pattern, we're authenticated
                if is_authenticated_by_url and not is_redirected:
                    logger.info("Session health check: PASSED - authenticated URL detected")
                    result = True
                elif is_redirected:
                    logger.warning("Session health check: FAILED - redirect to login detected")
                    result = False
                else:
                    # Secondary check: Look for any page content that indicates we're not on a login page
                    try:
                        # Check for login form elements
                        login_indicators = [
                            'input[type="password"]',
                            'input[name*="login"]',
                            'input[name*="user"]',
                            'input[id*="login"]',
                            'input[id*="user"]',
                            'button:has-text("Log In")',
                            'button:has-text("Sign In")'
                        ]
                        
                        login_found = False
                        for selector in login_indicators:
                            login_element = await page.query_selector(selector)
                            if login_element:
                                login_found = True
                                break
                        
                        if login_found:
                            logger.warning("Session health check: FAILED - login form detected")
                            result = False
                        else:
                            logger.info("Session health check: PASSED - no login form detected")
                            result = True
                            
                    except Exception as e:
                        logger.debug(f"Login form check error: {e}")
                        # If we can't check, assume healthy if we have valid cookies and no redirect
                        logger.info("Session health check: PASSED - based on cookies and URL")
                        result = True
                
                await context.close()
                await browser.close()
                
                return result
                
            except Exception as e:
                logger.error(f"Session health check error: {e}")
                try:
                    await context.close()
                except:
                    pass
                await browser.close()
                return False
                
    except Exception as e:
        logger.error(f"Session health check failed: {e}")
        return False


def get_session_info() -> Dict[str, Any]:
    """
    Get information about the current session state.
    
    Returns:
        Dict containing session information
    """
    cookies_path = get_cookies_path()
    try:
        with open(cookies_path, 'r') as f:
            cookies = json.load(f)
        
        session_cookies = []
        expiring_cookies = []
        current_time = datetime.now().timestamp()
        
        for cookie in cookies:
            name = cookie.get('name', '')
            expires = cookie.get('expires', -1)
            
            # Check if this is a session-related cookie
            if any(keyword in name.lower() for keyword in ['session', 'auth', 'token']):
                session_cookies.append({
                    'name': name,
                    'domain': cookie.get('domain', ''),
                    'expires': expires,
                    'is_session_cookie': expires == -1
                })
                
                if expires != -1 and expires > 0:
                    days_until_expire = (expires - current_time) / (24 * 3600)
                    if days_until_expire < 7:  # Expiring within a week
                        expiring_cookies.append({
                            'name': name,
                            'days_until_expire': days_until_expire
                        })
        
        return {
            'total_cookies': len(cookies),
            'session_cookies': len(session_cookies),
            'expiring_cookies': len(expiring_cookies),
            'expiring_soon': expiring_cookies,
            'session_status': 'active' if session_cookies else 'no_session_cookies'
        }
        
    except (FileNotFoundError, json.JSONDecodeError):
        return {
            'error': 'No valid cookies found',
            'total_cookies': 0,
            'session_cookies': 0,
            'expiring_cookies': 0,
            'expiring_soon': [],
            'session_status': 'missing_cookies'
        }


async def ensure_valid_session() -> bool:
    """
    Ensure we have a valid session, attempting refresh if needed.
    
    Returns:
        bool: True if a valid session exists or was successfully refreshed
    """
    logger = logging.getLogger(__name__)
    
    # First check if we have any cookies
    cookies_path = get_cookies_path()
    try:
        with open(cookies_path, 'r') as f:
            cookies = json.load(f)
        
        if not cookies:
            logger.error("No cookies available")
            return False
            
    except (FileNotFoundError, json.JSONDecodeError):
        logger.error("No valid cookies found")
        return False
    
    # Check session health
    if await maintain_session_health():
        logger.info("Session is healthy")
        return True
    
    # Session needs refresh
    logger.info("Session needs refresh, attempting navigation refresh")
    return await refresh_session_state(cookies)


async def get_session_status(debug: bool = False) -> Envelope[dict]:
    logger = logging.getLogger(__name__)

    try:
        # First get basic cookie information
        info = get_session_info()
        
        # If we have session cookies, validate they actually work with Schwab
        if info.get('session_status') == 'active':
            logger.debug("Session cookies found, validating with Schwab...")
            
            # Use maintain_session_health to actually test the session
            is_healthy = await maintain_session_health()
            
            if not is_healthy:
                # Update status to reflect that cookies exist but are invalid
                info['session_status'] = 'invalid'
                info['validation_error'] = 'Session cookies exist but Schwab authentication failed'
                logger.warning("Session validation failed: cookies present but not accepted by Schwab")
            else:
                logger.debug("Session validation succeeded")
        
        logger.debug("Session status info: %s", info)
        return ok(info)
    except Exception as exc:
        logger.exception("Failed to gather session status")
        return fail(str(exc), ErrorType.UNKNOWN, retryable=True)


async def refresh_session(debug: bool = False) -> Envelope[None]:
    logger = logging.getLogger(__name__)

    try:
        refreshed = await refresh_session_state()
        if refreshed:
            logger.info("Session refresh succeeded")
            return ok(None)
        logger.warning("Session refresh failed")
        return fail("Session refresh failed", ErrorType.AUTHENTICATION, retryable=True)
    except Exception as exc:
        logger.exception("Exception during session refresh")
        return fail(str(exc), ErrorType.UNKNOWN, retryable=True)


async def set_cookies_from_file(path: str, debug: bool = False) -> Envelope[None]:
    logger = logging.getLogger(__name__)

    try:
        with open(path, "r") as fh:
            cookies = json.load(fh)

        cookies_path = get_cookies_path()
        with open(cookies_path, "w") as fh:
            json.dump(cookies, fh, indent=2)

        logger.info("Imported %s cookies from %s", len(cookies), path)
        return ok(None)
    except (FileNotFoundError, json.JSONDecodeError) as exc:
        logger.error("Failed to load cookies from %s: %s", path, exc)
        return fail(str(exc), ErrorType.VALIDATION, retryable=False)
    except Exception as exc:
        logger.exception("Unexpected error importing cookies from %s", path)
        return fail(str(exc), ErrorType.UNKNOWN, retryable=True)


async def export_cookies(path: str, debug: bool = False) -> Envelope[None]:
    logger = logging.getLogger(__name__)

    cookies_path = get_cookies_path()
    try:
        with open(cookies_path, "r") as fh:
            cookies = json.load(fh)

        with open(path, "w") as fh:
            json.dump(cookies, fh, indent=2)

        logger.info("Exported %s cookies to %s", len(cookies), path)
        return ok(None)
    except (FileNotFoundError, json.JSONDecodeError) as exc:
        logger.error("Failed to read cookies for export: %s", exc)
        return fail(str(exc), ErrorType.AUTHENTICATION, retryable=False)
    except Exception as exc:
        logger.exception("Unexpected error exporting cookies to %s", path)
        return fail(str(exc), ErrorType.UNKNOWN, retryable=True)