"""Infrastructure provisioning — create the Bedrock RAG pipeline via boto3."""

from __future__ import annotations

import json
import logging
import time
import uuid
from pathlib import Path
from typing import Any, Callable, Optional

import boto3
import botocore.exceptions

from nimbus_core.bedrock import EMBEDDING_MODEL_ARN
from nimbus_core.models import InitResult

logger = logging.getLogger(__name__)

# Documents bundled with the package
_DOCS_DIR = Path(__file__).parent / "documents"
# Fallback: repo root when running from source
_REPO_DOCS_DIR = Path(__file__).parent.parent.parent / "document"


def _boto(service: str, region: str):
    return boto3.client(service, region_name=region)


def check_aws_credentials(region: Optional[str] = None) -> tuple[str, str]:
    """Validate AWS credentials and resolve region.

    Returns (account_id, region).
    Raises RuntimeError on failure.
    """
    session = boto3.session.Session()
    region = region or session.region_name or "us-east-1"

    try:
        sts = session.client("sts", region_name=region)
        identity = sts.get_caller_identity()
        account_id = identity["Account"]
        return account_id, region
    except botocore.exceptions.NoCredentialsError as e:
        raise RuntimeError("No AWS credentials found. Run `aws configure` or set environment variables.") from e
    except botocore.exceptions.ClientError as e:
        code = e.response["Error"]["Code"]
        if code in ("InvalidClientTokenId", "AuthFailure"):
            raise RuntimeError("AWS credentials are invalid or expired.") from e
        elif code == "ExpiredTokenException":
            raise RuntimeError("AWS session token has expired.") from e
        elif code == "AccessDenied":
            raise RuntimeError("Access denied — credentials lack sts:GetCallerIdentity permission.") from e
        raise RuntimeError(f"AWS error: {e.response['Error']['Message']}") from e


def create_docs_bucket(region: str, project: str = "nimbus") -> str:
    """Create an S3 bucket for KB documents. Returns bucket name."""
    s3 = _boto("s3", region)
    bucket_name = f"{project}-kb-docs-{uuid.uuid4().hex[:8]}"
    logger.info("Creating S3 docs bucket: %s", bucket_name)

    if region == "us-east-1":
        s3.create_bucket(Bucket=bucket_name)
    else:
        s3.create_bucket(
            Bucket=bucket_name,
            CreateBucketConfiguration={"LocationConstraint": region},
        )
    s3.put_public_access_block(
        Bucket=bucket_name,
        PublicAccessBlockConfiguration={
            "BlockPublicAcls": True, "IgnorePublicAcls": True,
            "BlockPublicPolicy": True, "RestrictPublicBuckets": True,
        },
    )
    s3.put_bucket_versioning(
        Bucket=bucket_name,
        VersioningConfiguration={"Status": "Enabled"},
    )
    return bucket_name


def upload_documents(bucket_name: str, region: str, docs_dir: Optional[Path] = None) -> int:
    """Upload security documents to the KB S3 bucket. Returns file count."""
    s3 = _boto("s3", region)

    if docs_dir is None:
        docs_dir = _DOCS_DIR if _DOCS_DIR.exists() else _REPO_DOCS_DIR
    if not docs_dir.exists():
        logger.warning("No documents directory found at %s", docs_dir)
        return 0

    files = [f for f in docs_dir.rglob("*") if f.is_file()]
    if not files:
        return 0

    for f in files:
        key = str(f.relative_to(docs_dir))
        s3.upload_file(str(f), bucket_name, key)

    logger.info("Uploaded %d documents to %s", len(files), bucket_name)
    return len(files)


def create_vector_store(region: str, project: str = "nimbus") -> tuple[str, str]:
    """Create S3 Vectors bucket + index. Returns (vector_bucket_name, index_name)."""
    s3v = _boto("s3vectors", region)
    vbucket = f"{project}-kb-vectors-{uuid.uuid4().hex[:8]}"
    index_name = f"{project}-kb-index"

    s3v.create_vector_bucket(vectorBucketName=vbucket)
    s3v.create_index(
        vectorBucketName=vbucket,
        indexName=index_name,
        dataType="float32",
        dimension=1024,
        distanceMetric="cosine",
    )
    return vbucket, index_name


def ensure_kb_role(region: str, project: str, docs_bucket: str) -> str:
    """Create or reuse the IAM role for Bedrock KB. Returns role ARN."""
    iam = _boto("iam", region)
    sts = _boto("sts", region)
    role_name = f"{project}-bedrock-kb-role"
    account_id = sts.get_caller_identity()["Account"]

    trust = {
        "Version": "2012-10-17",
        "Statement": [{
            "Effect": "Allow",
            "Principal": {"Service": "bedrock.amazonaws.com"},
            "Action": "sts:AssumeRole",
            "Condition": {"StringEquals": {"aws:SourceAccount": account_id}},
        }],
    }

    try:
        resp = iam.create_role(
            RoleName=role_name,
            AssumeRolePolicyDocument=json.dumps(trust),
            Description="Nimbus Bedrock Knowledge Base role",
        )
        role_arn = resp["Role"]["Arn"]
    except iam.exceptions.EntityAlreadyExistsException:
        role_arn = iam.get_role(RoleName=role_name)["Role"]["Arn"]

    policy = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Action": ["s3:GetObject", "s3:ListBucket"],
                "Resource": [f"arn:aws:s3:::{docs_bucket}", f"arn:aws:s3:::{docs_bucket}/*"],
            },
            {
                "Effect": "Allow",
                "Action": ["bedrock:InvokeModel", "bedrock:InvokeModelWithResponseStream"],
                "Resource": "*",
            },
            {"Effect": "Allow", "Action": ["s3vectors:*"], "Resource": "*"},
        ],
    }
    iam.put_role_policy(
        RoleName=role_name,
        PolicyName=f"{project}-kb-policy",
        PolicyDocument=json.dumps(policy),
    )
    time.sleep(10)  # wait for IAM propagation
    return role_arn


def create_knowledge_base(
    region: str,
    project: str,
    role_arn: str,
    vbucket: str,
    index_name: str,
    docs_bucket: str,
) -> tuple[str, str]:
    """Create Bedrock KB + data source. Returns (kb_id, data_source_id)."""
    bedrock_agent = _boto("bedrock-agent", region)
    s3v = _boto("s3vectors", region)

    idx = s3v.get_index(vectorBucketName=vbucket, indexName=index_name)
    index_arn = idx["index"]["indexArn"]

    kb_name = f"{project}-security-kb"
    resp = bedrock_agent.create_knowledge_base(
        name=kb_name,
        description="Nimbus AWS security documentation knowledge base",
        roleArn=role_arn,
        knowledgeBaseConfiguration={
            "type": "VECTOR",
            "vectorKnowledgeBaseConfiguration": {
                "embeddingModelArn": EMBEDDING_MODEL_ARN.format(region=region),
                "embeddingModelConfiguration": {
                    "bedrockEmbeddingModelConfiguration": {
                        "dimensions": 1024,
                        "embeddingDataType": "FLOAT32",
                    }
                },
            },
        },
        storageConfiguration={
            "type": "S3_VECTORS",
            "s3VectorsConfiguration": {"indexArn": index_arn},
        },
    )
    kb_id = resp["knowledgeBase"]["knowledgeBaseId"]

    # Wait for ACTIVE
    for _ in range(30):
        status = bedrock_agent.get_knowledge_base(knowledgeBaseId=kb_id)["knowledgeBase"]["status"]
        if status == "ACTIVE":
            break
        time.sleep(5)

    ds_resp = bedrock_agent.create_data_source(
        knowledgeBaseId=kb_id,
        name=f"{project}-security-docs",
        dataSourceConfiguration={
            "type": "S3",
            "s3Configuration": {"bucketArn": f"arn:aws:s3:::{docs_bucket}"},
        },
    )
    ds_id = ds_resp["dataSource"]["dataSourceId"]
    return kb_id, ds_id


def start_ingestion(kb_id: str, ds_id: str, region: str) -> None:
    """Start a KB ingestion job."""
    bedrock_agent = _boto("bedrock-agent", region)
    bedrock_agent.start_ingestion_job(knowledgeBaseId=kb_id, dataSourceId=ds_id)


def provision_pipeline(
    region: Optional[str] = None,
    project: str = "nimbus",
    docs_dir: Optional[Path] = None,
    on_step: Optional[Callable[[int, str], None]] = None,
) -> InitResult:
    """Provision the full RAG pipeline end-to-end.

    Args:
        region: AWS region (auto-detected if None).
        project: Project name prefix for resources.
        docs_dir: Path to documents to upload (auto-detected if None).
        on_step: Optional callback(step_number, description) for progress reporting.

    Returns:
        InitResult with all provisioned resource identifiers.
    """
    def _step(n: int, msg: str) -> None:
        if on_step:
            on_step(n, msg)

    account_id, region = check_aws_credentials(region)

    _step(1, "Creating S3 documents bucket")
    docs_bucket = create_docs_bucket(region, project)

    _step(2, "Uploading security documents")
    upload_documents(docs_bucket, region, docs_dir)

    _step(3, "Creating S3 Vectors store")
    vbucket, index_name = create_vector_store(region, project)

    _step(4, "Creating IAM role for Bedrock")
    role_arn = ensure_kb_role(region, project, docs_bucket)

    _step(5, "Creating Bedrock Knowledge Base")
    kb_id, ds_id = create_knowledge_base(
        region, project, role_arn, vbucket, index_name, docs_bucket
    )
    start_ingestion(kb_id, ds_id, region)

    return InitResult(
        account_id=account_id,
        region=region,
        kb_id=kb_id,
        docs_bucket=docs_bucket,
        vector_bucket=vbucket,
        ds_id=ds_id,
    )
