"""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.config
import botocore.exceptions

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

logger = logging.getLogger(__name__)

# Public S3 bucket containing curated security documents
SOURCE_DOCS_BUCKET = "nimbus-security-docs"
SOURCE_DOCS_REGION = "us-east-1"


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:
    """Sync security documents from the public source bucket to the user's KB bucket.

    Downloads from the public nimbus-security-docs bucket and uploads
    to the user's private KB docs bucket. Falls back to a local directory
    if provided via docs_dir.

    Returns file count.
    """
    dest_s3 = _boto("s3", region)

    # If a local docs directory is provided, upload from there instead
    if docs_dir and docs_dir.exists():
        files = [f for f in docs_dir.rglob("*") if f.is_file()]
        for f in files:
            key = str(f.relative_to(docs_dir))
            dest_s3.upload_file(str(f), bucket_name, key)
        logger.info("Uploaded %d local documents to %s", len(files), bucket_name)
        return len(files)

    # Sync from public source bucket (no credentials needed for public reads)
    source_s3 = boto3.client("s3", region_name=SOURCE_DOCS_REGION,
                             config=botocore.config.Config(signature_version=botocore.UNSIGNED))

    try:
        paginator = source_s3.get_paginator("list_objects_v2")
        count = 0
        for page in paginator.paginate(Bucket=SOURCE_DOCS_BUCKET):
            for obj in page.get("Contents", []):
                key = obj["Key"]
                # Download from public bucket, upload to user's private bucket
                response = source_s3.get_object(Bucket=SOURCE_DOCS_BUCKET, Key=key)
                body = response["Body"].read()
                dest_s3.put_object(Bucket=bucket_name, Key=key, Body=body)
                count += 1

        logger.info("Synced %d documents from %s to %s", count, SOURCE_DOCS_BUCKET, bucket_name)
        return count
    except Exception as e:
        logger.warning("Could not sync documents from source bucket: %s", e)
        return 0


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 _load_existing_config() -> dict:
    """Load existing config from ~/.nimbus-cli.json if it exists."""
    cfg_path = Path.home() / ".nimbus-cli.json"
    if cfg_path.exists():
        return json.loads(cfg_path.read_text())
    return {}


def _save_config(cfg: dict) -> None:
    """Save config to ~/.nimbus-cli.json."""
    cfg_path = Path.home() / ".nimbus-cli.json"
    cfg_path.write_text(json.dumps(cfg, indent=2) + "\n")


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.

    Idempotent — skips resources that already exist in the saved config.
    Safe to run multiple times.

    Args:
        region: AWS region. Defaults to us-east-1 for best Bedrock support.
        project: Project name prefix for resources.
        docs_dir: Path to documents to upload (syncs from public bucket 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)

    cfg = _load_existing_config()

    # If already fully initialized, tell the user
    if cfg.get("kb_id") and cfg.get("docs_bucket") and cfg.get("vector_bucket"):
        logger.info("Pipeline already provisioned (kb_id=%s). Skipping.", cfg["kb_id"])
        return InitResult(
            account_id=cfg.get("account_id", ""),
            region=cfg.get("region", region or "us-east-1"),
            kb_id=cfg["kb_id"],
            docs_bucket=cfg.get("docs_bucket", ""),
            vector_bucket=cfg.get("vector_bucket", ""),
            ds_id=cfg.get("ds_id", ""),
        )

    # Default to us-east-1 for Bedrock compatibility
    region = region or "us-east-1"
    account_id, region = check_aws_credentials(region)

    # Step 1 — S3 Documents Bucket
    docs_bucket = cfg.get("docs_bucket")
    if docs_bucket:
        _step(1, f"S3 docs bucket already exists ({docs_bucket})")
    else:
        _step(1, "Creating S3 documents bucket")
        docs_bucket = create_docs_bucket(region, project)
        cfg["docs_bucket"] = docs_bucket
        cfg["region"] = region
        cfg["project"] = project
        cfg["account_id"] = account_id
        _save_config(cfg)

    # Step 2 — Upload/sync documents
    _step(2, "Syncing security documents")
    upload_documents(docs_bucket, region, docs_dir)

    # Step 3 — S3 Vectors
    vbucket = cfg.get("vector_bucket")
    index_name = cfg.get("vector_index", f"{project}-kb-index")
    if vbucket:
        _step(3, f"S3 Vectors already exists ({vbucket})")
    else:
        _step(3, "Creating S3 Vectors store")
        vbucket, index_name = create_vector_store(region, project)
        cfg["vector_bucket"] = vbucket
        cfg["vector_index"] = index_name
        _save_config(cfg)

    # Step 4 — IAM Role
    _step(4, "Ensuring IAM role for Bedrock")
    role_arn = ensure_kb_role(region, project, docs_bucket)
    cfg["kb_role_arn"] = role_arn
    _save_config(cfg)

    # Step 5 — Bedrock Knowledge Base
    kb_id = cfg.get("kb_id")
    ds_id = cfg.get("ds_id")
    if kb_id:
        _step(5, f"Knowledge Base already exists ({kb_id})")
    else:
        _step(5, "Creating Bedrock Knowledge Base")
        kb_id, ds_id = create_knowledge_base(
            region, project, role_arn, vbucket, index_name, docs_bucket
        )
        cfg["kb_id"] = kb_id
        cfg["ds_id"] = ds_id
        _save_config(cfg)
        start_ingestion(kb_id, ds_id, region)

    # Save final config
    cfg.update({
        "region": region,
        "bedrock_region": region,
        "account_id": account_id,
        "project": project,
    })
    _save_config(cfg)

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


def destroy_pipeline(
    config: Optional[dict] = None,
    on_step: Optional[Callable[[int, str], None]] = None,
) -> None:
    """Tear down all AWS resources created by provision_pipeline.

    Reads resource identifiers from the provided config dict
    (or from ~/.nimbus-cli.json if None).

    Args:
        config: Dict with keys like kb_id, ds_id, docs_bucket, vector_bucket, region.
        on_step: Optional callback(step_number, description) for progress reporting.
    """
    if config is None:
        cfg_path = Path.home() / ".nimbus-cli.json"
        if not cfg_path.exists():
            raise RuntimeError("No config found at ~/.nimbus-cli.json. Nothing to destroy.")
        config = json.loads(cfg_path.read_text())

    region = config.get("bedrock_region") or config.get("region", "us-east-1")
    project = config.get("project", "nimbus")

    def _step(n: int, msg: str) -> None:
        if on_step:
            on_step(n, msg)

    # Step 1 — Delete Bedrock KB (this also deletes data sources)
    kb_id = config.get("kb_id")
    if kb_id:
        _step(1, f"Deleting Bedrock Knowledge Base ({kb_id})")
        try:
            bedrock_agent = _boto("bedrock-agent", region)
            # Delete data source first if we have the ID
            ds_id = config.get("ds_id")
            if ds_id:
                try:
                    bedrock_agent.delete_data_source(knowledgeBaseId=kb_id, dataSourceId=ds_id)
                except Exception:
                    pass  # may already be deleted
            bedrock_agent.delete_knowledge_base(knowledgeBaseId=kb_id)
            logger.info("Deleted Knowledge Base: %s", kb_id)
        except Exception as e:
            logger.warning("Could not delete KB %s: %s", kb_id, e)
    else:
        _step(1, "No Knowledge Base to delete")

    # Step 2 — Delete S3 Vectors bucket
    vector_bucket = config.get("vector_bucket")
    if vector_bucket:
        _step(2, f"Deleting S3 Vectors bucket ({vector_bucket})")
        try:
            s3v = _boto("s3vectors", region)
            index_name = config.get("vector_index", f"{project}-kb-index")
            try:
                s3v.delete_index(vectorBucketName=vector_bucket, indexName=index_name)
            except Exception:
                pass
            s3v.delete_vector_bucket(vectorBucketName=vector_bucket)
            logger.info("Deleted vector bucket: %s", vector_bucket)
        except Exception as e:
            logger.warning("Could not delete vector bucket %s: %s", vector_bucket, e)
    else:
        _step(2, "No vector bucket to delete")

    # Step 3 — Empty and delete S3 docs bucket
    docs_bucket = config.get("docs_bucket")
    if docs_bucket:
        _step(3, f"Deleting S3 docs bucket ({docs_bucket})")
        try:
            s3 = _boto("s3", region)
            # Empty the bucket first (required before deletion)
            paginator = s3.get_paginator("list_object_versions")
            for page in paginator.paginate(Bucket=docs_bucket):
                objects = []
                for v in page.get("Versions", []):
                    objects.append({"Key": v["Key"], "VersionId": v["VersionId"]})
                for m in page.get("DeleteMarkers", []):
                    objects.append({"Key": m["Key"], "VersionId": m["VersionId"]})
                if objects:
                    s3.delete_objects(Bucket=docs_bucket, Delete={"Objects": objects})
            s3.delete_bucket(Bucket=docs_bucket)
            logger.info("Deleted docs bucket: %s", docs_bucket)
        except Exception as e:
            logger.warning("Could not delete docs bucket %s: %s", docs_bucket, e)
    else:
        _step(3, "No docs bucket to delete")

    # Step 4 — Delete IAM role
    _step(4, "Deleting IAM role")
    role_name = f"{project}-bedrock-kb-role"
    try:
        iam = _boto("iam", region)
        # Remove inline policies first
        try:
            iam.delete_role_policy(RoleName=role_name, PolicyName=f"{project}-kb-policy")
        except Exception:
            pass
        iam.delete_role(RoleName=role_name)
        logger.info("Deleted IAM role: %s", role_name)
    except Exception as e:
        logger.warning("Could not delete IAM role %s: %s", role_name, e)

    # Step 5 — Delete config file
    _step(5, "Cleaning up local config")
    cfg_path = Path.home() / ".nimbus-cli.json"
    if cfg_path.exists():
        cfg_path.unlink()
