#!/usr/bin/env python3
"""
Stage 2 Bootstrap - Requires dependencies
Run this after 'poetry install' and 'sam pipeline bootstrap'
"""

import json
import os
import subprocess
from pathlib import Path

import click
from dotenv import load_dotenv

# Get the project root directory (parent of scripts directory)
SCRIPT_DIR = Path(__file__).parent.resolve()
PROJECT_ROOT = SCRIPT_DIR.parent

# Try to import toml
try:
    import toml
except ImportError:
    try:
        import tomli as toml
    except ImportError:
        click.echo("❌ toml module not found. Please install with: pip install toml")
        raise ImportError("toml module required for parsing SAM config")


def run_command(cmd, check=True):
    """Run command and return output"""
    result = subprocess.run(cmd, check=False, shell=True, capture_output=True, text=True)
    if check and result.returncode != 0:
        click.echo(f"Error: {result.stderr}", err=True)
        raise click.Abort()
    return result.stdout.strip()


def extract_aws_config_from_sam():
    """Extract AWS configuration from SAM pipeline config files"""
    sam_dir = PROJECT_ROOT / ".aws-sam" / "pipeline"
    if not sam_dir.exists():
        click.echo("❌ No .aws-sam/pipeline directory found")
        click.echo("Please run 'sam pipeline bootstrap' first")
        return {}

    # Find pipelineconfig.toml files recursively
    config_files = list(sam_dir.rglob("pipelineconfig.toml"))
    click.echo(f"🔍 Looking for pipelineconfig.toml files in: {sam_dir}")
    click.echo(
        f"🔍 Found {len(config_files)} pipelineconfig.toml files: {[str(f) for f in config_files]}"
    )

    if not config_files:
        click.echo("❌ No pipelineconfig.toml files found")
        click.echo("💡 Checking for alternative config files...")

        # Check for .toml files in subdirectories
        toml_files = list(sam_dir.rglob("*.toml"))
        click.echo(f"🔍 All .toml files found: {[str(f) for f in toml_files]}")

        # Try to extract from AWS directly
        click.echo("📝 Trying to find AWS resources directly...")
        return extract_aws_config_from_roles()

    # If we found config files, process the first one
    config_file = config_files[0]
    click.echo(f"📄 Reading config from: {config_file}")

    # Check if file exists and is readable
    if not config_file.exists():
        click.echo(f"❌ Config file does not exist: {config_file}")
        return {}

    click.echo(f"📄 File exists, size: {config_file.stat().st_size} bytes")

    try:
        # Try to read the file content first
        with open(config_file) as f:
            content = f.read()
        click.echo(f"📄 Raw file content (first 500 chars):\n{content[:500]}")

        config = toml.load(config_file)
        click.echo(f"📊 Config structure: {list(config.keys())}")
        click.echo("🔍 Full config content:")
        for key, value in config.items():
            click.echo(
                f"  {key}: {type(value)} - {value if not isinstance(value, dict) else list(value.keys())}"
            )

        # Try different config structures
        updates = {}

        # Try the standard structure first
        stage_config = config.get("default", {}).get("pipeline_bootstrap", {}).get("parameters", {})
        click.echo(f"🔍 Default config parameters: {stage_config}")

        # Check if default config has the AWS resources we need
        if not any(
            param in stage_config
            for param in [
                "pipeline_execution_role",
                "cloudformation_execution_role",
                "artifacts_bucket",
            ]
        ):
            click.echo(
                "🔍 Default config doesn't have AWS resources, searching for stage-specific configs..."
            )
            # Try alternative structures - look for any stage that has pipeline_bootstrap
            for key in config:
                click.echo(f"🔍 Checking key: {key}")
                if isinstance(config[key], dict):
                    click.echo(f"  -> Contains: {list(config[key].keys())}")
                    if "pipeline_bootstrap" in config[key]:
                        stage_params = config[key]["pipeline_bootstrap"].get("parameters", {})
                        click.echo(
                            f"  -> Pipeline bootstrap parameters: {list(stage_params.keys())}"
                        )
                        if any(
                            param in stage_params
                            for param in [
                                "pipeline_execution_role",
                                "cloudformation_execution_role",
                                "artifacts_bucket",
                            ]
                        ):
                            stage_config = stage_params
                            click.echo(f"📋 Using config from stage: {key}")
                            break

        click.echo(f"🔍 Final stage_config: {stage_config}")

        if "pipeline_execution_role" in stage_config:
            updates["TESTING_PIPELINE_EXECUTION_ROLE"] = stage_config["pipeline_execution_role"]
        if "cloudformation_execution_role" in stage_config:
            updates["TESTING_CLOUDFORMATION_EXECUTION_ROLE"] = stage_config[
                "cloudformation_execution_role"
            ]
        if "artifacts_bucket" in stage_config:
            updates["TESTING_ARTIFACTS_BUCKET"] = stage_config["artifacts_bucket"]

        click.echo(f"📋 Extracted {len(updates)} values from config: {updates}")
        return updates
    except Exception as e:
        click.echo(f"❌ Error reading config: {e}")
        import traceback

        click.echo(f"🔍 Full traceback: {traceback.format_exc()}")
        return {}


def extract_aws_config_from_roles():
    """Try to extract AWS resources by looking for SAM-created roles"""
    load_dotenv(PROJECT_ROOT / ".env")
    gh_repo = os.getenv("GH_REPO", "")

    if not gh_repo:
        click.echo("❌ GH_REPO not found in .env")
        return {}

    # Look for SAM-created roles
    project_prefix = gh_repo.replace("-", "")[:10]  # SAM truncates names

    try:
        # Find pipeline execution role
        role_query = f"Roles[?contains(RoleName, '{project_prefix}') && contains(RoleName, 'PipelineExecutionRole')].RoleName"
        pipeline_role = run_command(
            f'aws iam list-roles --query "{role_query}" --output text', check=False
        )

        if pipeline_role and pipeline_role != "None":
            click.echo(f"📋 Found pipeline role: {pipeline_role}")

            # Get the role ARN
            role_arn = run_command(
                f'aws iam get-role --role-name {pipeline_role} --query "Role.Arn" --output text',
                check=False,
            )

            updates = {}
            if role_arn:
                updates["TESTING_PIPELINE_EXECUTION_ROLE"] = role_arn

                # Try to find related resources
                cf_role_query = f"Roles[?contains(RoleName, '{project_prefix}') && contains(RoleName, 'CloudFormationExecutionRole')].RoleName"
                cf_role = run_command(
                    f'aws iam list-roles --query "{cf_role_query}" --output text',
                    check=False,
                )

                if cf_role and cf_role != "None":
                    cf_role_arn = run_command(
                        f'aws iam get-role --role-name {cf_role} --query "Role.Arn" --output text',
                        check=False,
                    )
                    if cf_role_arn:
                        updates["TESTING_CLOUDFORMATION_EXECUTION_ROLE"] = cf_role_arn

                # Try to find S3 bucket
                bucket_query = f"Buckets[?contains(Name, '{project_prefix}')].Name"
                bucket = run_command(
                    f'aws s3api list-buckets --query "{bucket_query}" --output text',
                    check=False,
                )

                if bucket and bucket != "None":
                    updates["TESTING_ARTIFACTS_BUCKET"] = bucket

                click.echo(f"📋 Found {len(updates)} AWS resources")
                return updates

        click.echo("❌ Could not find SAM-created resources")
        return {}

    except Exception as e:
        click.echo(f"❌ Error searching for AWS resources: {e}")
        return {}


def update_env_file(updates):
    """Update .env file with new values"""
    env_file = PROJECT_ROOT / ".env"
    if not env_file.exists():
        click.echo("❌ .env file not found. Run bootstrap-stage1.py first")
        raise click.Abort()

    # Read existing content
    with open(env_file, encoding="utf-8") as f:
        lines = f.readlines()

    # Update values
    for key, value in updates.items():
        found = False
        for i, line in enumerate(lines):
            if line.startswith(f"{key}="):
                lines[i] = f"{key}={value}\n"
                found = True
                break

        if not found:
            # Add new key if not found
            lines.append(f"{key}={value}\n")

        click.echo(f"✅ Updated {key}")

    # Write back
    with open(env_file, "w", encoding="utf-8") as f:
        f.writelines(lines)


def fix_trust_policy():
    """Fix the SAM-generated trust policy for GitHub Actions"""
    load_dotenv(PROJECT_ROOT / ".env")

    # Get account ID
    account_id = run_command("aws sts get-caller-identity --query Account --output text")
    if not account_id or account_id == "Account":
        click.echo("❌ Could not get AWS account ID")
        click.echo("💡 Make sure AWS CLI is configured: aws configure")
        return

    # Update .env with account ID
    update_env_file({"AWS_ACCOUNT_ID": account_id})

    gh_account = os.getenv("GH_ACCOUNT")
    gh_repo = os.getenv("GH_REPO")

    if not gh_account or not gh_repo:
        click.echo("❌ Missing GH_ACCOUNT or GH_REPO in .env")
        return

    # Find the pipeline execution role
    project_prefix = gh_repo[:9]
    role_query = f"Roles[?starts_with(RoleName, 'aws-sam-cli-managed-{project_prefix}') && contains(RoleName, 'PipelineExecutionRole')].RoleName"
    role_name = run_command(f'aws iam list-roles --query "{role_query}" --output text')

    if not role_name:
        click.echo("❌ Could not find PipelineExecutionRole")
        click.echo("Make sure SAM pipeline bootstrap completed successfully")
        return

    click.echo(f"🔧 Updating trust policy for role: {role_name}")

    # Check if OIDC provider exists
    oidc_arn = f"arn:aws:iam::{account_id}:oidc-provider/token.actions.githubusercontent.com"
    try:
        run_command(
            f"aws iam get-open-id-connect-provider --open-id-connect-provider-arn {oidc_arn}",
            check=True,
        )
        click.echo("✅ GitHub OIDC provider exists")
    except Exception:
        click.echo("❌ GitHub OIDC provider not found")
        if click.confirm("Would you like to create the GitHub OIDC provider?"):
            try:
                run_command(
                    "aws iam create-open-id-connect-provider "
                    "--url https://token.actions.githubusercontent.com "
                    "--thumbprint-list 6938fd4d98bab03faadb97b34396831e3780aea1 "
                    "--client-id-list sts.amazonaws.com"
                )
                click.echo("✅ Created GitHub OIDC provider")
            except Exception:
                click.echo("❌ Failed to create OIDC provider - you may need admin permissions")
                click.echo("💡 Ask an admin to run:")
                click.echo("   aws iam create-open-id-connect-provider \\")
                click.echo("     --url https://token.actions.githubusercontent.com \\")
                click.echo("     --thumbprint-list 6938fd4d98bab03faadb97b34396831e3780aea1 \\")
                click.echo("     --client-id-list sts.amazonaws.com")
        else:
            click.echo("⚠️  Skipping OIDC provider creation - your pipeline may fail without it")

    # Ask before updating trust policy
    if not click.confirm(
        f"Update trust policy for role '{role_name}' to allow GitHub Actions OIDC?"
    ):
        click.echo("⚠️  Skipping trust policy update")
        return

    # Create trust policy
    trust_policy = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Principal": {
                    "Federated": f"arn:aws:iam::{account_id}:oidc-provider/token.actions.githubusercontent.com"
                },
                "Action": "sts:AssumeRoleWithWebIdentity",
                "Condition": {
                    "StringEquals": {
                        "token.actions.githubusercontent.com:aud": "sts.amazonaws.com"
                    },
                    "StringLike": {
                        "token.actions.githubusercontent.com:sub": [
                            f"repo:{gh_account}/{gh_repo}:ref:refs/heads/fix/*",
                            f"repo:{gh_account}/{gh_repo}:ref:refs/heads/feat/*",
                            f"repo:{gh_account}/{gh_repo}:ref:refs/heads/perf/*",
                            f"repo:{gh_account}/{gh_repo}:ref:refs/heads/docs/*",
                            f"repo:{gh_account}/{gh_repo}:ref:refs/heads/refactor/*",
                            f"repo:{gh_account}/{gh_repo}:ref:refs/heads/dev",
                            f"repo:{gh_account}/{gh_repo}:ref:refs/heads/main",
                        ]
                    },
                },
            }
        ],
    }

    # Write and apply trust policy
    trust_policy_path = PROJECT_ROOT / "trust-policy.json"
    with open(trust_policy_path, "w") as f:
        json.dump(trust_policy, f)

    try:
        run_command(
            f"aws iam update-assume-role-policy --role-name {role_name} "
            f"--policy-document file://{trust_policy_path}"
        )
        click.echo("✅ Updated trust policy for GitHub Actions OIDC")
    except Exception as e:
        click.echo(f"❌ Failed to update trust policy: {e}")
    finally:
        # Clean up the temp file
        if trust_policy_path.exists():
            trust_policy_path.unlink()


@click.command()
def main():
    """Stage 2 Bootstrap - Complete AWS setup after SAM pipeline bootstrap"""
    click.echo("🚀 Augint Library Bootstrap - Stage 2")
    click.echo("=====================================\n")

    # Change to project root directory for consistent execution
    original_cwd = os.getcwd()
    os.chdir(PROJECT_ROOT)

    try:
        # Step 1: Extract AWS config from SAM
        click.echo("📦 Extracting AWS configuration from SAM...")
        updates = extract_aws_config_from_sam()

        if updates:
            # Step 2: Update .env file
            click.echo("\n📝 Updating .env file...")
            update_env_file(updates)
            click.echo("✅ Updated .env with AWS resources:")
            for key, value in updates.items():
                click.echo(f"  {key}: {value}")
        else:
            click.echo("⚠️  Could not extract AWS resources from SAM config")
            click.echo("💡 You may need to manually add these to .env:")

        # Step 3: Fix trust policy
        click.echo("\n🔐 Configuring IAM trust policy for GitHub Actions...")
        fix_trust_policy()

        # Done!
        click.echo("\n" + "=" * 50)
        click.echo("✅ Stage 2 Bootstrap Complete!")
        click.echo("=" * 50)

        click.echo("\n📋 Next Steps:\n")
        click.echo("1. Review your .env file to ensure all values are correct")

        click.echo("\n2. Secure your .env file with chezmoi:")
        click.echo("   chezmoi add .env")
        click.echo("   chezmoi git add .")
        click.echo(
            f'   chezmoi git commit -- -am "{os.getenv("GH_REPO", "project")} .env updated with AWS resources"'
        )
        click.echo("   chezmoi git push")

        click.echo("\n3. Set up development environment:")
        click.echo("   poetry run pre-commit install")
        click.echo("   pre-commit run --all-files  # Initial run to fix any issues")

        click.echo("\n4. Run initial tests to verify setup:")
        click.echo("   poetry run pytest")
        click.echo("   poetry run ruff check")

        click.echo("\n5. Configure GitHub repository:")
        click.echo(
            "   - Configure branch protection (see .github/BRANCH_PROTECTION.md for detailed instructions)"
        )
        click.echo("   - Set up PyPI trusted publishing at:")
        click.echo("     https://pypi.org/manage/account/publishing/")

        click.echo("\n6. Create initial commit and push:")
        click.echo("   git add -A")
        click.echo("   git commit -m 'feat: initialize project from augint-library template'")
        click.echo("   git push")

        click.echo("\n7. Monitor your first pipeline run:")
        click.echo("   - Check GitHub Actions for the pipeline status")
        click.echo("   - First run may take longer due to cache warming")

        click.echo("\n8. Optional: Configure MCP servers for enhanced Claude Code experience:")
        click.echo("   Add these to your Claude Desktop config for this project:")
        click.echo("   - File system access: npx @modelcontextprotocol/server-filesystem")
        click.echo("   - Python development: npx @modelcontextprotocol/server-python")
        click.echo("   - Database access: npx @modelcontextprotocol/server-sqlite")
        click.echo("   - AWS integration: npx @modelcontextprotocol/server-aws")
        click.echo("   - GitHub integration: npx @modelcontextprotocol/server-github")
        click.echo("   - Web search: npx @modelcontextprotocol/server-brave-search")

        click.echo("\n📋 All commands for remaining setup:")
        click.echo("─" * 45)
        click.echo("chezmoi add .env")
        click.echo("chezmoi git add .")
        click.echo(f'chezmoi git commit -- -am "{os.getenv("GH_REPO", "project")} .env updated"')
        click.echo("chezmoi git push")

        click.echo("poetry run pre-commit install")
        click.echo("pre-commit run --all-files")
        click.echo("poetry run pytest")
        click.echo("poetry run ruff check")
        click.echo("ai-gh-push")

        click.echo("git add -A")
        click.echo("git commit -m 'feat: initialize project from augint-library template'")
        click.echo("git push")

        click.echo("npx @modelcontextprotocol/server-filesystem")
        click.echo("npx @modelcontextprotocol/server-python")
        click.echo("npx @modelcontextprotocol/server-aws")
        click.echo("npx @modelcontextprotocol/server-github")
        click.echo("npx @modelcontextprotocol/server-brave-search")
        click.echo("─" * 45)

        click.echo("\n🎉 Your project is ready for development!")
    finally:
        # Restore original working directory
        os.chdir(original_cwd)


if __name__ == "__main__":
    main()
