import boto3
import json
import os
import subprocess

# --- 配置参数 ---
APPLICATION_ID = None
REGION = None
USER_ID = None
USE_IAM_AUTH = True  # Set to True to use IAM credentials instead of SSO
AWS_ACCESS_KEY_ID = None  # Optional: set directly or use environment variables
AWS_SECRET_ACCESS_KEY = None  # Optional: set directly or use environment variables
# -----------------

qbusiness_client = None  # Will be initialized in main()


def auto_detect_parameters():
    """Auto-detect Q Business parameters from AWS CLI and local state files"""
    global APPLICATION_ID, REGION, USER_ID

    # Try to get Application ID from /tmp/qbusiness-app.json
    try:
        if os.path.exists('/tmp/qbusiness-app.json'):
            with open('/tmp/qbusiness-app.json', 'r') as f:
                app_data = json.load(f)
                APPLICATION_ID = app_data.get('applicationId')
                # Extract region from ARN (format: arn:aws:qbusiness:REGION:account:application/id)
                app_arn = app_data.get('applicationArn', '')
                if app_arn:
                    arn_parts = app_arn.split(':')
                    if len(arn_parts) >= 4:
                        REGION = arn_parts[3]
                print(f"✓ Found Application ID: {APPLICATION_ID}")
                if REGION:
                    print(f"✓ Detected region from application: {REGION}")
    except Exception as e:
        print(f"Warning: Could not read /tmp/qbusiness-app.json: {e}")

    # If not found, try AWS CLI
    if not APPLICATION_ID:
        try:
            result = subprocess.run(
                ['aws', 'qbusiness', 'list-applications', '--max-results', '1'],
                capture_output=True, text=True, timeout=10
            )
            if result.returncode == 0:
                apps = json.loads(result.stdout).get('applications', [])
                if apps:
                    APPLICATION_ID = apps[0].get('applicationId')
                    print(f"✓ Found Application ID from AWS CLI: {APPLICATION_ID}")
        except Exception as e:
            print(f"Warning: Could not list applications: {e}")

    # Get region from AWS CLI config if not already set
    if not REGION:
        try:
            result = subprocess.run(
                ['aws', 'configure', 'get', 'region'],
                capture_output=True, text=True, timeout=5
            )
            if result.returncode == 0 and result.stdout.strip():
                REGION = result.stdout.strip()
                print(f"✓ Using AWS region: {REGION}")
            else:
                REGION = "us-west-2"
                print(f"✓ Using default region: {REGION}")
        except Exception as e:
            REGION = "us-west-2"
            print(f"Warning: Could not get region, using default: {REGION}")

    # Try to get Identity Center user ID
    try:
        # First, get Identity Store ID from Identity Center
        result = subprocess.run(
            ['aws', 'sso-admin', 'list-instances', '--region', 'eu-north-1'],
            capture_output=True, text=True, timeout=10
        )
        if result.returncode == 0:
            instances = json.loads(result.stdout).get('Instances', [])
            if instances:
                identity_store_id = instances[0].get('IdentityStoreId')

                # Get first user from Identity Store
                result = subprocess.run(
                    ['aws', 'identitystore', 'list-users',
                     '--identity-store-id', identity_store_id,
                     '--region', 'eu-north-1', '--max-results', '1'],
                    capture_output=True, text=True, timeout=10
                )
                if result.returncode == 0:
                    users = json.loads(result.stdout).get('Users', [])
                    if users:
                        USER_ID = users[0].get('UserId')
                        username = users[0].get('UserName', 'unknown')
                        print(f"✓ Found Identity Center user: {username} ({USER_ID})")
    except Exception as e:
        print(f"Warning: Could not get Identity Center user: {e}")

    # Set default user ID if not found
    if not USER_ID:
        USER_ID = "cli-user-" + os.environ.get('USER', 'default')
        print(f"✓ Using generated user ID: {USER_ID}")

    # Validate required parameters
    if not APPLICATION_ID:
        print("\n⚠ ERROR: Could not find Q Business Application ID!")
        print("Please ensure:")
        print("  1. Q Business application is deployed")
        print("  2. /tmp/qbusiness-app.json exists, or")
        print("  3. You have AWS CLI configured with proper credentials")
        return False

    return True


def initialize_client(region=None, application_id=None, user_id=None, profile=None, use_iam=None):
    """Initialize the Q Business client with configuration"""
    global qbusiness_client, APPLICATION_ID, REGION, USER_ID

    if region:
        REGION = region
    if application_id:
        APPLICATION_ID = application_id
    if user_id:
        USER_ID = user_id

    # Determine authentication method
    auth_mode = use_iam if use_iam is not None else USE_IAM_AUTH

    if auth_mode:
        # Use IAM credentials (access keys)
        print("✓ Using IAM authentication (access keys)")

        # Create session with explicit credentials if provided
        session_kwargs = {'region_name': REGION}

        if AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY:
            session_kwargs['aws_access_key_id'] = AWS_ACCESS_KEY_ID
            session_kwargs['aws_secret_access_key'] = AWS_SECRET_ACCESS_KEY
            print("✓ Using explicitly configured access keys")
        else:
            # Use credentials from environment variables or ~/.aws/credentials
            print("✓ Using credentials from environment or ~/.aws/credentials")

        session = boto3.Session(**session_kwargs)
        qbusiness_client = session.client('qbusiness')

        # Verify credentials
        try:
            sts = session.client('sts')
            identity = sts.get_caller_identity()
            print(f"✓ Authenticated as: {identity['Arn']}")
        except Exception as e:
            print(f"⚠ Warning: Could not verify identity: {e}")

        return qbusiness_client
    else:
        # Use SSO/Profile authentication
        if profile or os.path.exists(os.path.expanduser('~/.aws/config')):
            try:
                session = boto3.Session(profile_name=profile or 'qbusiness-sso', region_name=REGION)
                qbusiness_client = session.client('qbusiness')
                print(f"✓ Using AWS profile: {profile or 'qbusiness-sso'}")
                return qbusiness_client
            except Exception as e:
                print(f"Note: Could not use SSO profile, falling back to default credentials")

        # Fall back to default credentials
        qbusiness_client = boto3.client('qbusiness', region_name=REGION)
        return qbusiness_client


def send_message_to_q(message_text, current_session_id):
    """
    向 Amazon Q Business 发送消息并获取响应。
    """
    if qbusiness_client is None:
        initialize_client()

    request_params = {
        "applicationId": APPLICATION_ID,
        "userMessage": message_text
    }

    # For IAM authentication (non-Identity Center apps), userId must be passed
    # For Identity Center enabled apps with SSO, userId should not be passed
    if USE_IAM_AUTH and USER_ID:
        request_params["userId"] = USER_ID
        print(f"[DEBUG] Using userId: {USER_ID}")

    if current_session_id:
        request_params["conversationId"] = current_session_id

    try:
        response = qbusiness_client.chat_sync(**request_params)
        return response
    except Exception as e:
        print(f"Error communicating with Amazon Q Business: {e}")
        print(f"[DEBUG] Request params: {json.dumps({k: v for k, v in request_params.items() if k != 'userMessage'}, indent=2)}")
        return None


def print_response(response):
    """Print the response from Q Business"""
    if not response:
        print("Could not get a response from Amazon Q Business.")
        return None

    # Get the system message (Q Business response)
    system_message = response.get('systemMessage')

    if system_message:
        print(f"\nAmazon Q: {system_message}")
    else:
        print("\nAmazon Q: (No textual response received)")
        # 打印完整的响应以帮助调试
        print("Full Q Business response for debugging:")
        print(json.dumps(response, indent=2, ensure_ascii=False))

    # 可选：打印引用的文档源
    if 'sourceAttributions' in response and response['sourceAttributions']:
        print("\nSources:")
        for i, source in enumerate(response['sourceAttributions']):
            title = source.get('title', 'N/A')
            url = source.get('url', 'N/A')
            snippet = source.get('snippet', '')
            print(f"  {i+1}. {title}")
            if url != 'N/A':
                print(f"     URL: {url}")
            if snippet:
                print(f"     Snippet: {snippet[:100]}...")

    return response.get('conversationId')


def check_sso_login():
    """Check if SSO is logged in"""
    try:
        result = subprocess.run(
            ['aws', 'sts', 'get-caller-identity', '--profile', 'qbusiness-sso'],
            capture_output=True, text=True, timeout=5
        )
        return result.returncode == 0
    except:
        return False


def main():
    """Main interactive chat loop"""
    print("=" * 60)
    print("Amazon Q Business CLI Chat")
    print("=" * 60)
    print("\nInitializing and detecting configuration...\n")

    # Auto-detect parameters
    if not auto_detect_parameters():
        print("\nFailed to initialize. Exiting...")
        return

    print(f"\nConfiguration Summary:")
    print(f"  Application ID: {APPLICATION_ID}")
    print(f"  Region: {REGION}")
    print(f"  User ID: {USER_ID}")
    print(f"  Auth Mode: {'IAM (Access Keys)' if USE_IAM_AUTH else 'SSO (Identity Center)'}")
    print("=" * 60)

    # Check SSO login only if not using IAM auth
    if not USE_IAM_AUTH:
        if not check_sso_login():
            print("\n⚠ WARNING: SSO is not logged in!")
            print("For Identity Center enabled Q Business apps, you need to login with SSO.")
            print("\nTo login, run:")
            print("  aws sso login --profile qbusiness-sso")
            print("\nOr visit: https://d-c367615483.awsapps.com/start")
            print("\n" + "=" * 60)
            response = input("\nDo you want to continue anyway? (y/N): ")
            if response.lower() != 'y':
                print("Exiting... Please login and try again.")
                return
            print("\nContinuing with default credentials (may fail)...\n")
    else:
        print("\n✓ Using IAM authentication - SSO login not required")
        print("  Ensure you have:")
        print("  - AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY in environment, or")
        print("  - Valid credentials in ~/.aws/credentials, or")
        print("  - Access keys configured in the script")
        print("=" * 60)

    initialize_client()

    session_id = None  # 用于存储会话 ID，实现多轮对话

    print("\n✓ Ready! Type your questions below.")
    print("Type 'exit' or 'quit' to end the conversation.\n")

    while True:
        try:
            user_input = input("\nYou: ")
            if user_input.lower() in ['exit', 'quit']:
                print("\nEnding conversation. Goodbye!")
                break

            if not user_input.strip():
                continue

            response = send_message_to_q(user_input, session_id)

            if response:
                # 从响应中提取并更新 session_id
                if 'conversationId' in response:
                    session_id = response['conversationId']

                print_response(response)
            else:
                print("Could not get a response from Amazon Q Business.")

        except KeyboardInterrupt:
            print("\n\nInterrupted by user. Goodbye!")
            break
        except EOFError:
            print("\nEnd of input. Goodbye!")
            break
        except Exception as e:
            print(f"\nUnexpected error: {e}")


if __name__ == "__main__":
    main()