"""CLI commands for deploying to AWS Bedrock AgentCore Runtime.

This command group provides integration with AWS Bedrock AgentCore,
allowing you to deploy LlamaIndex workflows as agents.

Uses direct boto3 API calls for full control and transparency.
Deploys via container build using AWS CodeBuild.
"""

from __future__ import annotations

import collections
import json
import random
import re
import string
import sys
import tempfile
import time
import uuid
import zipfile
from pathlib import Path
from typing import Any

import click
from llama_deploy.core.config import DEFAULT_DEPLOYMENT_FILE_PATH
from llama_deploy.core.deployment_config import (
    read_deployment_config_from_git_root_or_cwd,
)
from rich import print as rprint

if sys.version_info >= (3, 11):
    import tomllib
else:
    import tomli as tomllib  # type: ignore[import-not-found]

from ..app import app
from ..styles import WARNING

# Type for workflow discovery
WorkflowInfo = dict[str, Any]  # {name, module_path, attr, start_event_class}

# Type for log position tracking
Position = collections.namedtuple("Position", ["timestamp", "skip"])


def _discover_workflows_from_pyproject(project_dir: Path) -> dict[str, WorkflowInfo]:
    """Discover workflows from [tool.llamadeploy.workflows] in pyproject.toml.

    Returns dict like:
        {"process-file": {"module_path": "...", "attr": "workflow", "start_event_class": "FileEvent"}}
    """
    pyproject_path = project_dir / "pyproject.toml"
    if not pyproject_path.exists():
        return {}

    with open(pyproject_path, "rb") as f:
        pyproject = tomllib.load(f)

    workflows_config = (
        pyproject.get("tool", {}).get("llamadeploy", {}).get("workflows", {})
    )
    if not workflows_config:
        return {}

    workflows: dict[str, WorkflowInfo] = {}
    for name, dotted_path in workflows_config.items():
        # Parse "extraction_review.process_file:workflow" -> (module, attr)
        module_path, _, attr = dotted_path.partition(":")
        attr = attr or "workflow"

        # Infer start event class from workflow name
        # Convention: workflows with "file" or "process" use FileEvent
        if "file" in name.lower() or "process" in name.lower():
            start_event_class = "FileEvent"
        else:
            start_event_class = "StartEvent"

        workflows[name] = {
            "module_path": module_path,
            "attr": attr,
            "start_event_class": start_event_class,
        }

    return workflows


def _generate_entrypoint_from_workflows(
    workflows: dict[str, WorkflowInfo],
    fallback_app_import: str | None = None,
) -> str:
    """Generate AgentCore entrypoint code from discovered workflows.

    If no workflows found, falls back to dynamic loading via APP_IMPORT_PATH.
    """
    template_path = (
        Path(__file__).parent.parent / "templates" / "agentcore_entrypoint.py.template"
    )

    if not template_path.exists():
        raise FileNotFoundError(f"Template not found: {template_path}")

    with open(template_path) as f:
        template = f.read()

    if not workflows:
        # Fallback to dynamic loading (original behavior)
        return _generate_dynamic_entrypoint(fallback_app_import)

    # Generate static imports
    imports = []
    entries = []
    has_metadata = any("metadata" in name.lower() for name in workflows)
    metadata_workflow_var = None

    for name, info in workflows.items():
        module_path = info["module_path"]
        attr = info["attr"]
        start_event_class = info["start_event_class"]
        safe_name = name.replace("-", "_")

        # Import the workflow
        imports.append(f"from {module_path} import {attr} as {safe_name}_workflow")

        # Import FileEvent if needed
        if start_event_class == "FileEvent":
            imports.append(f"from {module_path} import FileEvent")

        # Track metadata workflow for caching
        if "metadata" in name.lower():
            metadata_workflow_var = f"{safe_name}_workflow"

        # Add workflow entry
        entries.append(
            f'    "{name}": {{\n'
            f'        "workflow": {safe_name}_workflow,\n'
            f'        "start_event_class": {start_event_class},\n'
            f"    }},"
        )

    # Add StartEvent import (always needed as fallback)
    imports.append("from workflows.events import StartEvent")

    # Generate payload routing logic
    routing_lines = []
    for name, info in workflows.items():
        if info["start_event_class"] == "FileEvent":
            routing_lines.append('        if "file_id" in payload:')
            routing_lines.append(f'            workflow_name = "{name}"')
            routing_lines.append(
                '            event_data = {"file_id": payload["file_id"]}'
            )

    # Default routing
    default_workflow = "metadata" if has_metadata else next(iter(workflows.keys()))
    if routing_lines:
        routing_lines.append("        else:")
        routing_lines.append(f'            workflow_name = "{default_workflow}"')
    else:
        routing_lines.append(f'        workflow_name = "{default_workflow}"')

    # Generate metadata workflow call
    if metadata_workflow_var:
        metadata_call = f"""    result = await {metadata_workflow_var}.run(start_event=StartEvent())
    if hasattr(result, "model_dump"):
        _metadata_cache = result.model_dump()
    else:
        _metadata_cache = result"""
    else:
        metadata_call = "    _metadata_cache = {}"

    # Fill in template
    code = template.replace("{{workflow_imports}}", "\n".join(imports))
    code = code.replace("{{workflow_entries}}", "\n".join(entries))
    code = code.replace("{{metadata_workflow_call}}", metadata_call)
    code = code.replace("{{payload_routing}}", "\n".join(routing_lines))

    return code


def _generate_dynamic_entrypoint(app_import: str | None) -> str:
    """Generate entrypoint with dynamic workflow loading (fallback mode)."""
    return f'''# SPDX-License-Identifier: MIT
"""
AgentCore entrypoint for LlamaIndex Workflows (dynamic loading mode).
Auto-generated by llamactl.
"""
from __future__ import annotations

import importlib
import logging
import os
import sys

# Add paths for imports (handles both deployed and local dev scenarios)
_current_dir = os.path.dirname(os.path.abspath(__file__))
_parent_dir = os.path.dirname(_current_dir)
sys.path.insert(0, _current_dir)
sys.path.insert(0, os.path.join(_current_dir, "src"))
# For local dev when running from .agentcore/ folder
sys.path.insert(0, _parent_dir)
sys.path.insert(0, os.path.join(_parent_dir, "src"))

# Try loading .env (optional - not needed in production where env vars are set via runtime config)
try:
    from dotenv import load_dotenv
    load_dotenv(os.path.join(_current_dir, ".env"), override=True)
    load_dotenv(os.path.join(_parent_dir, ".env"), override=True)
except ImportError:
    pass  # dotenv not available, env vars should be set via AgentCore runtime config

try:
    from bedrock_agentcore import BedrockAgentCoreApp
except ImportError:
    # Try alternate import paths
    try:
        from bedrock_agentcore.app import BedrockAgentCoreApp
    except ImportError:
        raise ImportError(
            "bedrock_agentcore not found. Ensure it is listed in requirements.txt"
        )
from workflows.events import StartEvent

logger = logging.getLogger(__name__)
app = BedrockAgentCoreApp()

# Load workflows dynamically from APP_IMPORT_PATH
APP_IMPORT_PATH = os.environ.get("APP_IMPORT_PATH", "{app_import or ""}")
WORKFLOWS = {{}}

def _load_workflows():
    if not APP_IMPORT_PATH:
        logger.warning("No APP_IMPORT_PATH set - no workflows available")
        return
    try:
        module_path, app_name = APP_IMPORT_PATH.split(":")
        module = importlib.import_module(module_path)
        app_obj = getattr(module, app_name)
        if hasattr(app_obj, "workflows"):
            for name, wf in app_obj.workflows.items():
                WORKFLOWS[name] = {{"workflow": wf, "start_event_class": StartEvent}}
            logger.info(f"Loaded workflows: {{list(WORKFLOWS.keys())}}")
    except Exception as e:
        logger.error(f"Failed to load workflows: {{e}}")

_load_workflows()


def _serialize(result):
    if hasattr(result, "model_dump"):
        return result.model_dump()
    if isinstance(result, (str, int, float, bool, dict, list, type(None))):
        return result
    return str(result)


@app.entrypoint
async def invoke(payload: dict, context):
    workflow_name = payload.get("workflow") or (next(iter(WORKFLOWS.keys())) if WORKFLOWS else None)
    if not workflow_name or workflow_name not in WORKFLOWS:
        return {{"error": f"Unknown workflow '{{workflow_name}}'", "available": list(WORKFLOWS.keys())}}

    event_data = payload.get("start_event", {{}})
    entry = WORKFLOWS[workflow_name]
    try:
        start_event = entry["start_event_class"](**event_data)
        result = await entry["workflow"].run(start_event=start_event)
        return {{
            "workflow": workflow_name,
            "status": "completed",
            "result": _serialize(result),
            "session_id": context.session_id,
        }}
    except Exception as e:
        logger.error(f"Workflow failed: {{e}}", exc_info=True)
        return {{"error": str(e)}}


if __name__ == "__main__":
    app.run()
'''


def _sanitize_agent_name(name: str) -> str:
    """Sanitize agent name - AgentCore requires underscores only, no hyphens."""
    return re.sub(r"[^a-zA-Z0-9]", "_", name)


@app.group(
    help="Deploy workflows to AWS Bedrock AgentCore Runtime.",
    no_args_is_help=True,
)
def agentcore() -> None:
    """Manage AgentCore deployments."""
    pass


@agentcore.command("deploy")
@click.argument(
    "deployment_file",
    required=False,
    default=DEFAULT_DEPLOYMENT_FILE_PATH,
    type=click.Path(exists=True, dir_okay=True, resolve_path=True, path_type=Path),
)
@click.option(
    "--project-name",
    help="Name for the AgentCore project (defaults to deployment name)",
)
@click.option(
    "--execution-role",
    required=True,
    help="ARN of IAM role for AgentCore Runtime execution (required)",
)
@click.option(
    "--region",
    help="AWS region for deployment (defaults to AWS_DEFAULT_REGION)",
)
@click.option(
    "--identity-provider",
    help="AgentCore Identity credential provider name for API keys (e.g., 'partner-llamaindex-api-key')",
)
@click.option(
    "--deployment-role",
    required=True,
    help="ARN of IAM role for deployment operations (CodeBuild, ECR push)",
)
@click.option(
    "--oauth-discovery-url",
    help="OIDC discovery URL for inbound OAuth (e.g., https://cognito-idp.../.../.well-known/openid-configuration)",
)
@click.option(
    "--oauth-client-id",
    multiple=True,
    help="Allowed OAuth client ID (can specify multiple times)",
)
@click.option(
    "--s3-bucket",
    help="S3 bucket for build artifacts (auto-created if not provided)",
)
def deploy(
    deployment_file: Path,
    project_name: str | None,
    execution_role: str,
    region: str | None,
    identity_provider: str | None,
    deployment_role: str,
    oauth_discovery_url: str | None,
    oauth_client_id: tuple[str, ...],
    s3_bucket: str | None,
) -> None:
    """Deploy workflows to AWS Bedrock AgentCore Runtime using boto3.

    This command:
    1. Discovers workflows from pyproject.toml [tool.llamadeploy.workflows]
       (falls back to llama_deploy.yaml app if not found)
    2. Generates an AgentCore-compatible entrypoint using BedrockAgentCoreApp
    3. Builds a Docker image via CodeBuild and deploys to AgentCore

    Prerequisites:
    - AWS credentials configured (aws configure)
    - Deployment role with CodeBuild/ECR permissions
    - Execution role with AgentCore runtime permissions
    - For Identity integration: credential provider set up in AgentCore

    Examples:
        llamactl agentcore deploy \\
            --deployment-role arn:aws:iam::123456789012:role/DeployRole \\
            --execution-role arn:aws:iam::123456789012:role/ExecutionRole

        # With inbound OAuth authentication
        llamactl agentcore deploy \\
            --deployment-role arn:... --execution-role arn:... \\
            --oauth-discovery-url https://cognito-idp.../.../.well-known/openid-configuration \\
            --oauth-client-id my-client-id
    """
    try:
        import boto3

        # Initialize boto3 session
        session = boto3.Session(region_name=region)
        if not region:
            region = session.region_name or "us-east-1"

        # Get AWS account ID
        account_id = session.client("sts").get_caller_identity()["Account"]
        rprint(f"[cyan]Deploying to AWS account {account_id} in {region}[/cyan]")

        # Determine project directory
        config_dir = (
            deployment_file if deployment_file.is_dir() else deployment_file.parent
        )

        # Try to read deployment config (may not exist for pyproject-only projects)
        config = None
        try:
            config = read_deployment_config_from_git_root_or_cwd(
                Path.cwd(), deployment_file
            )
        except Exception:
            pass  # Config not found, will use pyproject.toml discovery

        # Discover workflows from pyproject.toml
        workflows = _discover_workflows_from_pyproject(config_dir)
        if workflows:
            rprint(
                f"[green]Discovered workflows from pyproject.toml: {list(workflows.keys())}[/green]"
            )
        elif config and config.app:
            rprint(
                f"[yellow]No tool.llamadeploy.workflows found, using app: {config.app}[/yellow]"
            )
        else:
            rprint(
                f"[{WARNING}]Error: No workflows found.[/]\n"
                "Either add tool.llamadeploy.workflows to pyproject.toml:\n"
                "  \[tool.llamadeploy.workflows\]\n"
                '  my-workflow = "my_package.workflow:workflow"\n'
                "\nOr add 'app' to llama_deploy.yaml:\n"
                "  app: path.to.module:app_variable"
            )
            raise click.Abort()

        # Determine project name
        raw_name = project_name or (config.name if config else None)
        if not raw_name:
            # Try to get from pyproject.toml
            pyproject_path = config_dir / "pyproject.toml"
            if pyproject_path.exists():
                with open(pyproject_path, "rb") as f:
                    pyproject = tomllib.load(f)
                raw_name = pyproject.get("project", {}).get(
                    "name", "llama-workflow-agent"
                )
            else:
                raw_name = "llama-workflow-agent"

        project_name = _sanitize_agent_name(raw_name)
        repository_name = f"bedrock-agentcore-{project_name}"
        runtime_name = f"{project_name}_runtime"

        # Create .agentcore directory
        agentcore_dir = config_dir / ".agentcore"
        agentcore_dir.mkdir(exist_ok=True)

        rprint(f"[green]Preparing deployment files in {agentcore_dir}[/green]")

        # Generate entrypoint from discovered workflows
        fallback_app = config.app if config else None
        entrypoint_code = _generate_entrypoint_from_workflows(workflows, fallback_app)
        entrypoint_path = agentcore_dir / "agentcore_entrypoint.py"
        with open(entrypoint_path, "w") as f:
            f.write(entrypoint_code)
        rprint("[green]✓ Generated agentcore_entrypoint.py[/green]")

        # Build environment variables
        env_vars: dict[str, str] = {"AWS_DEFAULT_REGION": region}
        if config and config.env:
            env_vars.update(config.env)
        if identity_provider:
            env_vars["AGENTCORE_CREDENTIAL_PROVIDER"] = identity_provider

        # Save deployment metadata
        _save_deployment_metadata(
            agentcore_dir,
            {
                "project_name": project_name,
                "runtime_name": runtime_name,
                "repository_name": repository_name,
                "region": region,
                "account_id": account_id,
                "oauth_enabled": bool(oauth_discovery_url and oauth_client_id),
            },
        )

        # Build container via CodeBuild
        rprint("[cyan]Building ARM64 container via CodeBuild...[/cyan]")

        # Prepare requirements and Dockerfile
        _prepare_requirements(config_dir, agentcore_dir)
        _generate_dockerfile(agentcore_dir)

        # Auto-create S3 bucket if not provided
        if not s3_bucket:
            s3_bucket = f"llamactl-agentcore-{account_id}-{region}"
            _ensure_s3_bucket(session, s3_bucket, region)

        ecr_uri = _build_and_push_image(
            session=session,
            source_dir=agentcore_dir,
            project_dir=config_dir,
            repository_name=repository_name,
            build_role=deployment_role,
            s3_bucket=s3_bucket,
        )

        rprint(f"[green]✓ Image built and pushed to {ecr_uri}[/green]")

        # Build OAuth authorizer config if provided
        authorizer_config: dict[str, Any] | None = None
        if oauth_discovery_url and oauth_client_id:
            authorizer_config = {
                "customJWTAuthorizer": {
                    "discoveryUrl": oauth_discovery_url,
                    "allowedClients": list(oauth_client_id),
                }
            }
            rprint(
                f"[green]OAuth enabled with discovery URL: {oauth_discovery_url}[/green]"
            )

        # Deploy to AgentCore Runtime
        rprint("[cyan]Creating/updating AgentCore Runtime...[/cyan]")
        agent_runtime_arn = _deploy_runtime(
            session=session,
            runtime_name=runtime_name,
            ecr_uri=ecr_uri,
            execution_role=execution_role,
            env_vars=env_vars,
            authorizer_config=authorizer_config,
        )

        rprint(
            f"\n[green]✓ Deployment successful![/green]\n"
            f"  Runtime: {runtime_name}\n"
            f"  ARN: {agent_runtime_arn}\n"
            f"  Region: {region}\n"
            f"\n[cyan]Test your workflow with:[/cyan]\n"
            f'  llamactl agentcore invoke "Hello, world!"'
        )

        if identity_provider:
            rprint(
                f"\n[yellow]Note: Ensure credential provider '{identity_provider}' "
                f"is set up in AgentCore Identity[/yellow]"
            )

    except click.Abort:
        raise
    except Exception as e:
        rprint(f"[red]Deployment failed: {e}[/red]")
        raise click.Abort()


@agentcore.command("invoke")
@click.argument("prompt", required=True)
@click.option(
    "--workflow",
    help="Workflow name to invoke (uses first workflow if not specified)",
)
@click.option(
    "--project-dir",
    type=click.Path(exists=True, file_okay=False, path_type=Path),
    default=Path.cwd(),
    help="Project directory containing .agentcore/",
)
@click.option(
    "--region",
    help="AWS region (defaults to deployment region)",
)
def invoke(
    prompt: str, workflow: str | None, project_dir: Path, region: str | None
) -> None:
    """Invoke a deployed AgentCore workflow using boto3.

    Example:
        llamactl agentcore invoke "Hello, world!"
        llamactl agentcore invoke "Analyze this data" --workflow my-workflow
    """
    import boto3

    agentcore_dir = project_dir / ".agentcore"
    if not agentcore_dir.exists():
        rprint(
            f"[{WARNING}]Error: .agentcore directory not found. Run 'llamactl agentcore deploy' first.[/]"
        )
        raise click.Abort()

    try:
        # Load deployment metadata
        metadata = _load_deployment_metadata(agentcore_dir)
        runtime_name = metadata["runtime_name"]
        deployment_region = metadata["region"]

        # Use specified region or fall back to deployment region
        region = region or deployment_region

        # Initialize boto3 client
        session = boto3.Session(region_name=region)
        control_client = session.client("bedrock-agentcore-control")
        runtime_client = session.client("bedrock-agentcore")

        # Find the runtime ARN
        list_response = control_client.list_agent_runtimes()
        runtime_arn = None
        for runtime in list_response.get("agentRuntimes", []):
            if runtime["agentRuntimeName"] == runtime_name:
                runtime_arn = runtime["agentRuntimeArn"]
                break

        if not runtime_arn:
            rprint(
                f"[red]Error: Runtime '{runtime_name}' not found in region {region}[/red]"
            )
            raise click.Abort()

        # Build payload
        payload_data = {"prompt": prompt}
        if workflow:
            payload_data["workflow"] = workflow

        rprint("[cyan]Invoking workflow...[/cyan]")

        # Invoke the runtime
        response = runtime_client.invoke_agent_runtime(
            agentRuntimeArn=runtime_arn,
            qualifier="DEFAULT",
            runtimeSessionId=str(uuid.uuid4()),
            payload=json.dumps(payload_data).encode(),
        )

        # Process response
        content = []
        for chunk in response.get("response", []):
            content.append(chunk.decode("utf-8"))

        result = json.loads("".join(content))

        rprint("\n[green]Response:[/green]")
        rprint(json.dumps(result, indent=2))

    except Exception as e:
        rprint(f"[red]Invocation failed: {e}[/red]")
        raise click.Abort()


@agentcore.command("destroy")
@click.option(
    "--project-dir",
    type=click.Path(exists=True, file_okay=False, path_type=Path),
    default=Path.cwd(),
    help="Project directory containing .agentcore/",
)
@click.option(
    "--region",
    help="AWS region (defaults to deployment region)",
)
@click.option(
    "--keep-ecr",
    is_flag=True,
    help="Keep ECR repository (only delete runtime)",
)
@click.confirmation_option(prompt="This will delete AgentCore resources. Continue?")
def destroy(project_dir: Path, region: str | None, keep_ecr: bool) -> None:
    """Destroy AgentCore deployment and clean up AWS resources using boto3.

    This removes:
    - AgentCore Runtime
    - ECR repository and images (unless --keep-ecr is specified)
    """
    import boto3

    agentcore_dir = project_dir / ".agentcore"
    if not agentcore_dir.exists():
        rprint(f"[{WARNING}]No .agentcore directory found - nothing to destroy.[/]")
        return

    try:
        # Load deployment metadata
        metadata = _load_deployment_metadata(agentcore_dir)
        runtime_name = metadata["runtime_name"]
        repository_name = metadata["repository_name"]
        deployment_region = metadata["region"]

        # Use specified region or fall back to deployment region
        region = region or deployment_region

        # Initialize boto3 client
        session = boto3.Session(region_name=region)
        control_client = session.client("bedrock-agentcore-control")
        ecr_client = session.client("ecr")

        # Find and delete the runtime
        rprint(f"[cyan]Deleting AgentCore Runtime: {runtime_name}...[/cyan]")
        list_response = control_client.list_agent_runtimes()
        runtime_id = None
        for runtime in list_response.get("agentRuntimes", []):
            if runtime["agentRuntimeName"] == runtime_name:
                runtime_id = runtime["agentRuntimeId"]
                break

        if runtime_id:
            control_client.delete_agent_runtime(agentRuntimeId=runtime_id)
            rprint(f"[green]✓ Runtime deleted: {runtime_name}[/green]")

            # Wait for deletion to complete
            rprint("[cyan]Waiting for deletion to complete...[/cyan]")
            while True:
                try:
                    status_response = control_client.get_agent_runtime(
                        agentRuntimeId=runtime_id
                    )
                    status = status_response["status"]
                    if status in ["DELETE_FAILED"]:
                        rprint(f"[red]Deletion failed with status: {status}[/red]")
                        break
                    rprint(f"  Status: {status}")
                    time.sleep(5)
                except control_client.exceptions.ResourceNotFoundException:
                    rprint("[green]✓ Runtime fully deleted[/green]")
                    break
        else:
            rprint(
                f"[yellow]Runtime '{runtime_name}' not found (may already be deleted)[/yellow]"
            )

        # Delete ECR repository unless --keep-ecr is specified
        if not keep_ecr:
            rprint(f"[cyan]Deleting ECR repository: {repository_name}...[/cyan]")
            try:
                ecr_client.delete_repository(repositoryName=repository_name, force=True)
                rprint(f"[green]✓ ECR repository deleted: {repository_name}[/green]")
            except ecr_client.exceptions.RepositoryNotFoundException:
                rprint(
                    f"[yellow]ECR repository '{repository_name}' not found (may already be deleted)[/yellow]"
                )
        else:
            rprint(f"[yellow]Keeping ECR repository: {repository_name}[/yellow]")

        rprint("\n[green]✓ Cleanup complete[/green]")

    except FileNotFoundError as e:
        rprint(f"[red]{e}[/red]")
        raise click.Abort()
    except Exception as e:
        rprint(f"[red]Destroy failed: {e}[/red]")
        raise click.Abort()


def _ensure_s3_bucket(session: Any, bucket_name: str, region: str) -> None:
    """Ensure S3 bucket exists for CodeBuild artifacts."""
    s3 = session.client("s3")
    try:
        s3.head_bucket(Bucket=bucket_name)
        rprint(f"[green]Using existing S3 bucket: {bucket_name}[/green]")
    except Exception:
        try:
            if region == "us-east-1":
                s3.create_bucket(Bucket=bucket_name)
            else:
                s3.create_bucket(
                    Bucket=bucket_name,
                    CreateBucketConfiguration={"LocationConstraint": region},
                )
            rprint(f"[green]Created S3 bucket: {bucket_name}[/green]")
        except Exception as e:
            rprint(f"[red]Error creating S3 bucket: {e}[/red]")
            raise


def _generate_dockerfile(agentcore_dir: Path) -> None:
    """Generate Dockerfile for AgentCore container deployment."""
    dockerfile_content = """FROM public.ecr.aws/docker/library/python:3.12-slim-bookworm

WORKDIR /app

# Copy and install requirements
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy all application code
COPY . .

# Set Python path to include src
ENV PYTHONPATH=/app/src:/app

# Run BedrockAgentCoreApp
CMD ["python", "agentcore_entrypoint.py"]
"""

    dockerfile_path = agentcore_dir / "Dockerfile"
    with open(dockerfile_path, "w") as f:
        f.write(dockerfile_content)


def _save_deployment_metadata(agentcore_dir: Path, metadata: dict[str, Any]) -> None:
    """Save deployment metadata for invoke/destroy commands."""
    metadata_path = agentcore_dir / "deployment.json"
    with open(metadata_path, "w") as f:
        json.dump(metadata, f, indent=2)


def _load_deployment_metadata(agentcore_dir: Path) -> dict[str, Any]:
    """Load deployment metadata."""
    metadata_path = agentcore_dir / "deployment.json"
    if not metadata_path.exists():
        raise FileNotFoundError(
            "No deployment metadata found. Run 'llamactl agentcore deploy' first."
        )
    with open(metadata_path) as f:
        return json.load(f)


def _log_stream(
    client: Any, log_group: str, stream_name: str, position: Position
) -> Any:
    """Stream logs from CloudWatch."""
    start_time, skip = position
    next_token = None
    event_count = 1
    while event_count > 0:
        token_arg = {"nextToken": next_token} if next_token else {}
        response = client.get_log_events(
            logGroupName=log_group,
            logStreamName=stream_name,
            startTime=start_time,
            startFromHead=True,
            **token_arg,
        )
        next_token = response["nextForwardToken"]
        events = response["events"]
        event_count = len(events)
        if event_count > skip:
            events = events[skip:]
            skip = 0
        else:
            skip = skip - event_count
            events = []
        for ev in events:
            ts, count = position
            if ev["timestamp"] == ts:
                position = Position(timestamp=ts, skip=count + 1)
            else:
                position = Position(timestamp=ev["timestamp"], skip=1)
            yield ev, position


def _logs_for_build(build_id: str, session: Any, poll: int = 10) -> str:
    """Stream logs for CodeBuild and return final status."""
    import botocore.config

    codebuild = session.client("codebuild")
    description = codebuild.batch_get_builds(ids=[build_id])["builds"][0]
    status = description["buildStatus"]
    log_group = description["logs"].get("groupName")
    stream_name = description["logs"].get("streamName")
    position = Position(timestamp=0, skip=0)

    config = botocore.config.Config(retries={"max_attempts": 15})
    logs_client = session.client("logs", config=config)

    # Wait for log group/stream to be available
    while log_group is None and status == "IN_PROGRESS":
        time.sleep(poll)
        description = codebuild.batch_get_builds(ids=[build_id])["builds"][0]
        log_group = description["logs"].get("groupName")
        stream_name = description["logs"].get("streamName")
        status = description["buildStatus"]

    last_describe_call = time.time()

    # Stream logs
    while True:
        for event, position in _log_stream(
            logs_client, log_group, stream_name, position
        ):
            rprint(event["message"].rstrip())

        if status != "IN_PROGRESS":
            break

        time.sleep(poll)

        if time.time() - last_describe_call >= 30:
            description = codebuild.batch_get_builds(ids=[build_id])["builds"][0]
            status = description["buildStatus"]
            last_describe_call = time.time()
            if status != "IN_PROGRESS":
                break

    return status


def _build_and_push_image(
    session: Any,
    source_dir: Path,
    project_dir: Path,
    repository_name: str,
    build_role: str,
    s3_bucket: str,
) -> str:
    """Build ARM64 Docker image via CodeBuild and push to ECR."""
    account_id = session.client("sts").get_caller_identity()["Account"]
    region = session.region_name

    # Create ECR repository
    ecr = session.client("ecr")
    try:
        ecr.create_repository(repositoryName=repository_name)
        rprint(f"[green]Created ECR repository: {repository_name}[/green]")
    except ecr.exceptions.RepositoryAlreadyExistsException:
        rprint(f"[green]Using existing ECR repository: {repository_name}[/green]")

    # Upload source code to S3
    random_suffix = "".join(random.choices(string.ascii_letters, k=16))
    key = f"codebuild-{random_suffix}.zip"

    with tempfile.TemporaryFile() as tmp:
        with zipfile.ZipFile(tmp, "w") as zip_file:
            # Add files from .agentcore/ (entrypoint, Dockerfile, requirements.txt)
            for file_path in source_dir.iterdir():
                if file_path.is_file():
                    zip_file.write(file_path, file_path.name)

            # Add root-level Python files from project
            for file_path in project_dir.glob("*.py"):
                if file_path.is_file():
                    zip_file.write(file_path, file_path.name)

            # Add src/ directory if it exists
            src_dir = project_dir / "src"
            if src_dir.exists():
                for file_path in src_dir.rglob("*"):
                    if file_path.is_file() and "__pycache__" not in str(file_path):
                        arcname = str(file_path.relative_to(project_dir))
                        zip_file.write(file_path, arcname)

            # Add configs/ directory if it exists
            configs_dir = project_dir / "configs"
            if configs_dir.exists():
                for file_path in configs_dir.rglob("*"):
                    if file_path.is_file():
                        arcname = str(file_path.relative_to(project_dir))
                        zip_file.write(file_path, arcname)

            # Add buildspec for ARM64 build
            buildspec = f"""version: 0.2
phases:
  pre_build:
    commands:
      - echo Logging in to Amazon ECR...
      - aws ecr get-login-password --region {region} | docker login --username AWS --password-stdin {account_id}.dkr.ecr.{region}.amazonaws.com
  build:
    commands:
      - echo Build started on `date`
      - echo Building the Docker image...
      - docker build -t {repository_name}:latest .
      - docker tag {repository_name}:latest {account_id}.dkr.ecr.{region}.amazonaws.com/{repository_name}:latest
  post_build:
    commands:
      - echo Build completed on `date`
      - echo Pushing the Docker image...
      - docker push {account_id}.dkr.ecr.{region}.amazonaws.com/{repository_name}:latest
"""
            zip_file.writestr("buildspec.yml", buildspec)

        tmp.seek(0)
        session.client("s3").upload_fileobj(tmp, s3_bucket, key)

    rprint(f"[green]Uploaded source to s3://{s3_bucket}/{key}[/green]")

    # Create and run CodeBuild project
    project_name = f"llama-build-{random_suffix}"
    codebuild = session.client("codebuild")

    codebuild.create_project(
        name=project_name,
        source={"type": "S3", "location": f"{s3_bucket}/{key}"},
        artifacts={"type": "NO_ARTIFACTS"},
        environment={
            "type": "ARM_CONTAINER",
            "image": "aws/codebuild/amazonlinux2-aarch64-standard:3.0",
            "computeType": "BUILD_GENERAL1_SMALL",
            "privilegedMode": True,
        },
        serviceRole=build_role,
    )

    rprint(f"[cyan]Starting CodeBuild project: {project_name}[/cyan]")

    build_id = codebuild.start_build(projectName=project_name)["build"]["id"]

    # Stream logs and wait for completion
    status = _logs_for_build(build_id, session)

    # Cleanup
    codebuild.delete_project(name=project_name)
    session.client("s3").delete_object(Bucket=s3_bucket, Key=key)

    if status != "SUCCEEDED":
        raise RuntimeError(f"CodeBuild failed with status: {status}")

    return f"{account_id}.dkr.ecr.{region}.amazonaws.com/{repository_name}:latest"


def _deploy_runtime(
    session: Any,
    runtime_name: str,
    ecr_uri: str,
    execution_role: str,
    env_vars: dict[str, str],
    authorizer_config: dict[str, Any] | None = None,
) -> str:
    """Create or update AgentCore Runtime."""
    client = session.client("bedrock-agentcore-control")

    # Check if runtime already exists
    existing_runtime = _find_existing_runtime(client, runtime_name)

    runtime_artifact = {"containerConfiguration": {"containerUri": ecr_uri}}

    # Build common kwargs
    common_kwargs: dict[str, Any] = {
        "roleArn": execution_role,
        "agentRuntimeArtifact": runtime_artifact,
        "networkConfiguration": {"networkMode": "PUBLIC"},
        "environmentVariables": env_vars,
        "lifecycleConfiguration": {"maxLifetime": 3600},
    }
    if authorizer_config:
        common_kwargs["authorizerConfiguration"] = authorizer_config

    if existing_runtime:
        rprint(f"[cyan]Updating existing runtime: {runtime_name}[/cyan]")
        client.update_agent_runtime(
            agentRuntimeId=existing_runtime["agentRuntimeId"],
            **common_kwargs,
        )
        agent_runtime_arn = existing_runtime["agentRuntimeArn"]
        agent_runtime_id = existing_runtime["agentRuntimeId"]
    else:
        rprint(f"[cyan]Creating new runtime: {runtime_name}[/cyan]")
        response = client.create_agent_runtime(
            agentRuntimeName=runtime_name,
            **common_kwargs,
        )
        agent_runtime_arn = response["agentRuntimeArn"]
        agent_runtime_id = response["agentRuntimeId"]

    # Wait for runtime to be ready
    _wait_for_runtime_ready(client, agent_runtime_id)

    return agent_runtime_arn


def _find_existing_runtime(client: Any, runtime_name: str) -> dict | None:
    """Find existing runtime by name."""
    try:
        list_response = client.list_agent_runtimes()
        for runtime in list_response.get("agentRuntimes", []):
            if runtime["agentRuntimeName"] == runtime_name:
                return runtime
    except Exception as e:
        rprint(f"[yellow]Could not list existing runtimes: {e}[/yellow]")
    return None


def _wait_for_runtime_ready(client: Any, runtime_id: str) -> None:
    """Wait for runtime to reach READY status."""
    rprint("[cyan]Waiting for runtime to be ready...[/cyan]")
    while True:
        status_response = client.get_agent_runtime(agentRuntimeId=runtime_id)
        status = status_response["status"]

        if status == "READY":
            rprint("[green]✓ Runtime is ready[/green]")
            break
        elif status in ["CREATE_FAILED", "UPDATE_FAILED", "DELETE_FAILED"]:
            raise RuntimeError(f"Runtime deployment failed with status: {status}")

        rprint(f"  Status: {status}")
        time.sleep(10)


def _prepare_requirements(config_dir: Path, agentcore_dir: Path) -> None:
    """Prepare requirements.txt for AgentCore container deployment.

    Priority order:
    1. pyproject.toml dependencies (if exists)
    2. requirements.txt (if exists)
    3. Minimal agentcore dependencies
    """
    pyproject_path = config_dir / "pyproject.toml"
    requirements_path = config_dir / "requirements.txt"
    target_path = agentcore_dir / "requirements.txt"

    # Required AgentCore dependencies
    agentcore_deps = [
        "bedrock-agentcore>=1.3.0",
        "python-dotenv>=1.0.0",
    ]

    deps: list[str] = []

    if pyproject_path.exists():
        # Extract dependencies from pyproject.toml
        with open(pyproject_path, "rb") as f:
            pyproject = tomllib.load(f)

        project_deps = pyproject.get("project", {}).get("dependencies", [])
        if project_deps:
            deps.extend(project_deps)
            rprint(
                f"[green]Extracted {len(project_deps)} dependencies from pyproject.toml[/green]"
            )

    if not deps and requirements_path.exists():
        # Fall back to requirements.txt
        with open(requirements_path) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#"):
                    deps.append(line)
        rprint(f"[green]Using {len(deps)} dependencies from requirements.txt[/green]")

    # Ensure AgentCore deps are included
    existing_names = {dep.split(">=")[0].split("==")[0].split("[")[0] for dep in deps}
    for dep in agentcore_deps:
        dep_name = dep.split(">=")[0]
        if dep_name not in existing_names:
            deps.append(dep)

    # Write requirements.txt
    with open(target_path, "w") as f:
        f.write("\n".join(deps) + "\n")
