#!/usr/bin/env python3
"""
Duo Manager for Media Downloader
Handles Duo Security 2FA Operations

Based on backup-central's implementation
Uses Duo Universal Prompt
"""

import sys
import sqlite3
import secrets
from datetime import datetime, timedelta
from typing import Optional, Dict
from pathlib import Path
import os

# Add parent path to allow imports from modules
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from modules.universal_logger import get_logger

logger = get_logger('DuoManager')


class DuoManager:
    def __init__(self, db_path: str = None):
        if db_path is None:
            db_path = str(Path(__file__).parent.parent.parent / 'database' / 'auth.db')

        self.db_path = db_path

        # Duo Configuration (from environment variables)
        self.client_id = os.getenv('DUO_CLIENT_ID')
        self.client_secret = os.getenv('DUO_CLIENT_SECRET')
        self.api_host = os.getenv('DUO_API_HOSTNAME')
        self.redirect_url = os.getenv('DUO_REDIRECT_URL', 'https://md.lic.ad/api/auth/2fa/duo/callback')

        # Check if Duo is configured
        self.is_configured = bool(self.client_id and self.client_secret and self.api_host)

        # State store (for OAuth flow)
        self.state_store = {}  # username → state mapping

        # Initialize database
        self._init_database()

        # Initialize Duo client if configured
        self.duo_client = None
        if self.is_configured:
            self._init_duo_client()

    def _init_database(self):
        """Initialize Duo-related database tables"""
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()

            # Ensure Duo columns exist in users table
            try:
                cursor.execute("ALTER TABLE users ADD COLUMN duo_enabled INTEGER NOT NULL DEFAULT 0")
            except sqlite3.OperationalError:
                pass  # Column already exists

            try:
                cursor.execute("ALTER TABLE users ADD COLUMN duo_username TEXT")
            except sqlite3.OperationalError:
                pass

            try:
                cursor.execute("ALTER TABLE users ADD COLUMN duo_enrolled_at TEXT")
            except sqlite3.OperationalError:
                pass

            # Duo audit log
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS duo_audit_log (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    username TEXT NOT NULL,
                    action TEXT NOT NULL,
                    success INTEGER NOT NULL,
                    ip_address TEXT,
                    user_agent TEXT,
                    details TEXT,
                    timestamp TEXT NOT NULL
                )
            """)

            conn.commit()

    def _init_duo_client(self):
        """Initialize Duo Universal Prompt client"""
        try:
            from duo_universal import Client
            self.duo_client = Client(
                client_id=self.client_id,
                client_secret=self.client_secret,
                host=self.api_host,
                redirect_uri=self.redirect_url
            )
            logger.info("Duo Universal Prompt client initialized successfully", module="Duo")
        except Exception as e:
            logger.error(f"Error initializing Duo client: {e}", module="Duo")
            self.is_configured = False

    def is_duo_configured(self) -> bool:
        """Check if Duo is properly configured"""
        return self.is_configured

    def get_configuration_status(self) -> Dict:
        """Get Duo configuration status"""
        return {
            'configured': self.is_configured,
            'hasClientId': bool(self.client_id),
            'hasClientSecret': bool(self.client_secret),
            'hasApiHostname': bool(self.api_host),
            'hasRedirectUrl': bool(self.redirect_url),
            'apiHostname': self.api_host if self.api_host else None
        }

    def generate_state(self, username: str, remember_me: bool = False) -> str:
        """
        Generate a state parameter for Duo OAuth flow

        Args:
            username: Username to associate with this state
            remember_me: Whether to remember the user for 30 days

        Returns:
            Random state string
        """
        state = secrets.token_urlsafe(32)
        self.state_store[state] = {
            'username': username,
            'remember_me': remember_me,
            'created_at': datetime.now(),
            'expires_at': datetime.now() + timedelta(minutes=10)
        }
        return state

    def verify_state(self, state: str) -> Optional[tuple]:
        """
        Verify state parameter and return associated username and remember_me

        Args:
            state: State parameter from Duo callback

        Returns:
            Tuple of (username, remember_me) if valid, None otherwise
        """
        if state not in self.state_store:
            return None

        state_data = self.state_store[state]

        # Check if expired
        if datetime.now() > state_data['expires_at']:
            del self.state_store[state]
            return None

        username = state_data['username']
        remember_me = state_data.get('remember_me', False)
        del self.state_store[state]  # One-time use

        return (username, remember_me)

    def create_auth_url(self, username: str, remember_me: bool = False) -> Dict:
        """
        Create Duo authentication URL for user

        Args:
            username: Username to authenticate
            remember_me: Whether to remember the user for 30 days

        Returns:
            Dict with authUrl and state
        """
        if not self.is_configured or not self.duo_client:
            raise Exception("Duo is not configured")

        try:
            # Generate state
            state = self.generate_state(username, remember_me)

            # Get Duo username for this user
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("SELECT duo_username FROM users WHERE username = ?", (username,))
                result = cursor.fetchone()

            duo_username = result[0] if result and result[0] else username

            # Create Duo Universal Prompt auth URL
            auth_url = self.duo_client.create_auth_url(duo_username, state)

            self._log_audit(username, 'duo_auth_url_created', True, None, None,
                          'Duo authentication URL created')

            return {
                'authUrl': auth_url,
                'state': state
            }
        except Exception as e:
            self._log_audit(username, 'duo_auth_url_failed', False, None, None,
                          f'Error: {str(e)}')
            raise

    def verify_duo_response(self, duo_code: str, username: str) -> bool:
        """
        Verify Duo authentication response (Universal Prompt)

        Args:
            duo_code: Authorization code from Duo
            username: Expected username

        Returns:
            True if authentication successful
        """
        if not self.is_configured or not self.duo_client:
            return False

        try:
            # Exchange authorization code for 2FA result
            decoded_token = self.duo_client.exchange_authorization_code_for_2fa_result(
                duo_code,
                username
            )

            # Check authentication result
            if decoded_token and decoded_token.get('auth_result', {}).get('result') == 'allow':
                authenticated_username = decoded_token.get('preferred_username')

                # Get Duo username for this user
                with sqlite3.connect(self.db_path) as conn:
                    cursor = conn.cursor()
                    cursor.execute("SELECT duo_username FROM users WHERE username = ?", (username,))
                    result = cursor.fetchone()

                duo_username = result[0] if result and result[0] else username

                # Check if username matches
                if authenticated_username == duo_username:
                    self._log_audit(username, 'duo_verify_success', True, None, None,
                                  'Duo authentication successful')
                    return True
                else:
                    self._log_audit(username, 'duo_verify_failed', False, None, None,
                                  f'Username mismatch: expected {duo_username}, got {authenticated_username}')
                    return False
            else:
                self._log_audit(username, 'duo_verify_failed', False, None, None,
                              'Duo authentication denied')
                return False

        except Exception as e:
            self._log_audit(username, 'duo_verify_failed', False, None, None,
                          f'Error: {str(e)}')
            logger.error(f"Duo verification error: {e}", module="Duo")
            return False

    def _get_application_key(self) -> str:
        """Get or generate application key for Duo"""
        key_file = Path(__file__).parent.parent.parent / '.duo_application_key'

        if key_file.exists():
            with open(key_file, 'r') as f:
                return f.read().strip()

        # Generate new key (at least 40 characters)
        key = secrets.token_urlsafe(40)
        try:
            with open(key_file, 'w') as f:
                f.write(key)
            os.chmod(key_file, 0o600)
        except Exception as e:
            logger.warning(f"Could not save Duo application key: {e}", module="Duo")

        return key

    def enroll_user(self, username: str, duo_username: Optional[str] = None) -> bool:
        """
        Enroll a user in Duo 2FA

        Args:
            username: Username
            duo_username: Optional Duo username (defaults to username)

        Returns:
            True if successful
        """
        if duo_username is None:
            duo_username = username

        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()

                cursor.execute("""
                    UPDATE users
                    SET duo_enabled = 1, duo_username = ?, duo_enrolled_at = ?
                    WHERE username = ?
                """, (duo_username, datetime.now().isoformat(), username))

                conn.commit()

            self._log_audit(username, 'duo_enrolled', True, None, None,
                          f'Duo enrolled with username: {duo_username}')

            return True
        except Exception as e:
            logger.error(f"Error enrolling user in Duo: {e}", module="Duo")
            return False

    def unenroll_user(self, username: str) -> bool:
        """
        Unenroll a user from Duo 2FA

        Args:
            username: Username

        Returns:
            True if successful
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()

                cursor.execute("""
                    UPDATE users
                    SET duo_enabled = 0, duo_username = NULL, duo_enrolled_at = NULL
                    WHERE username = ?
                """, (username,))

                conn.commit()

            self._log_audit(username, 'duo_unenrolled', True, None, None,
                          'Duo unenrolled')

            return True
        except Exception as e:
            logger.error(f"Error unenrolling user from Duo: {e}", module="Duo")
            return False

    def get_duo_status(self, username: str) -> Dict:
        """
        Get Duo status for a user

        Args:
            username: Username

        Returns:
            Dict with enabled, duoUsername, enrolledAt
        """
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()

            cursor.execute("""
                SELECT duo_enabled, duo_username, duo_enrolled_at
                FROM users
                WHERE username = ?
            """, (username,))

            row = cursor.fetchone()

        if not row:
            return {'enabled': False, 'duoUsername': None, 'enrolledAt': None}

        return {
            'enabled': bool(row[0]),
            'duoUsername': row[1],
            'enrolledAt': row[2]
        }

    def preauth_user(self, username: str) -> Dict:
        """
        Perform Duo preauth to check enrollment status

        Args:
            username: Username

        Returns:
            Dict with result and status_msg
        """
        if not self.is_configured or not self.duo_client:
            return {'result': 'error', 'status_msg': 'Duo not configured'}

        try:
            # Get Duo username
            user_info = self.get_duo_status(username)
            duo_username = user_info.get('duoUsername', username)

            # Perform preauth
            preauth_result = self.duo_client.preauth(username=duo_username)

            self._log_audit(username, 'duo_preauth', True, None, None,
                          f'Preauth result: {preauth_result.get("result")}')

            return preauth_result

        except Exception as e:
            self._log_audit(username, 'duo_preauth_failed', False, None, None,
                          f'Error: {str(e)}')
            return {'result': 'error', 'status_msg': 'Failed to check Duo enrollment status'}

    def _log_audit(self, username: str, action: str, success: bool,
                   ip_address: Optional[str], user_agent: Optional[str],
                   details: Optional[str] = None):
        """Log Duo audit event"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()

                cursor.execute("""
                    INSERT INTO duo_audit_log
                    (username, action, success, ip_address, user_agent, details, timestamp)
                    VALUES (?, ?, ?, ?, ?, ?, ?)
                """, (username, action, int(success), ip_address, user_agent, details,
                      datetime.now().isoformat()))

                conn.commit()
        except Exception as e:
            logger.error(f"Error logging Duo audit: {e}", module="Duo")

    def health_check(self) -> Dict:
        """
        Check Duo service health

        Returns:
            Dict with healthy status and details
        """
        if not self.is_configured:
            return {'healthy': False, 'error': 'Duo not configured'}

        try:
            # Simple ping to Duo API
            response = self.duo_client.ping()
            return {'healthy': True, 'response': response}
        except Exception as e:
            logger.error(f"Duo health check failed: {e}", module="Duo")
            return {'healthy': False, 'error': 'Duo service unavailable'}