"""Bedrock RAG remediation — core logic shared by CLI, SDK, and backend."""

from __future__ import annotations

import json
import logging
from typing import Any, Optional

import boto3

from nimbus_core.models import FindingInput, RemediationResult

logger = logging.getLogger(__name__)

# Default model/account values (can be overridden via NimbusClient)
DEFAULT_MODEL_ID = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
DEFAULT_REGION = "us-east-1"

EMBEDDING_MODEL_ARN = (
    "arn:aws:bedrock:{region}::foundation-model/amazon.titan-embed-text-v2:0"
)


def normalize_finding(finding: FindingInput) -> dict[str, Any]:
    """Normalize a FindingInput into the dict format expected by the Bedrock prompt.

    Also accepts a dict with the same keys (for backend compatibility).
    """
    if isinstance(finding, dict):
        # Backend passes dicts via the Finding pydantic model — adapt
        return _normalize_from_dict(finding)

    finding_type = "CONFIGURATION_CHECK"
    if finding.details:
        raw_type = finding.details.get("type", "")
        if raw_type:
            finding_type = raw_type

    resource_type = "AWS_RESOURCE"
    if finding.details and finding.details.get("resource_type"):
        resource_type = finding.details["resource_type"]
    else:
        service_type_map = {
            "ec2": "AWS_EC2_INSTANCE",
            "s3": "AWS_S3_BUCKET",
            "rds": "AWS_RDS_DB_INSTANCE",
            "lambda": "AWS_LAMBDA_FUNCTION",
            "iam": "AWS_IAM_RESOURCE",
            "ecr": "AWS_ECR_CONTAINER_IMAGE",
            "ecs": "AWS_ECS_SERVICE",
            "eks": "AWS_EKS_CLUSTER",
            "kms": "AWS_KMS_KEY",
            "sqs": "AWS_SQS_QUEUE",
            "vpc": "AWS_EC2_VPC",
        }
        resource_type = service_type_map.get(finding.service, "AWS_RESOURCE")

    network_components: list[dict] = []
    if finding.details and finding.details.get("related_resources"):
        for res in finding.details["related_resources"]:
            if isinstance(res, str):
                network_components.append({"id": res, "type": "RELATED_RESOURCE"})

    vulnerability_id = "N/A"
    if finding.source in ("inspector", "inspector_cis"):
        if finding.title and finding.title.startswith("CVE-"):
            vulnerability_id = finding.title.split(" ")[0]
    if finding.details:
        check_id = finding.details.get("check_id")
        if check_id and vulnerability_id == "N/A":
            vulnerability_id = check_id

    return {
        "title": finding.title,
        "severity": finding.severity,
        "finding_type": finding_type,
        "resource_id": finding.resource_id or finding.resource_arn or "N/A",
        "resource_type": resource_type,
        "description": _build_enriched_description(finding),
        "network_components": network_components,
        "vulnerabilities": [],
        "vulnerability_id": vulnerability_id,
    }


def _normalize_from_dict(finding: dict) -> dict[str, Any]:
    """Normalize from a raw dict (backend Finding.model_dump() compatibility)."""
    source = finding.get("source", "direct_check")
    service = finding.get("service", "general")
    details = finding.get("details") or {}

    finding_type = "CONFIGURATION_CHECK"
    if details.get("type"):
        finding_type = details["type"]

    resource_type = "AWS_RESOURCE"
    if details.get("resource_type"):
        resource_type = details["resource_type"]
    else:
        service_type_map = {
            "ec2": "AWS_EC2_INSTANCE", "s3": "AWS_S3_BUCKET",
            "rds": "AWS_RDS_DB_INSTANCE", "lambda": "AWS_LAMBDA_FUNCTION",
            "iam": "AWS_IAM_RESOURCE", "ecr": "AWS_ECR_CONTAINER_IMAGE",
            "ecs": "AWS_ECS_SERVICE", "eks": "AWS_EKS_CLUSTER",
            "kms": "AWS_KMS_KEY", "sqs": "AWS_SQS_QUEUE", "vpc": "AWS_EC2_VPC",
        }
        resource_type = service_type_map.get(service, "AWS_RESOURCE")

    network_components: list[dict] = []
    if details.get("related_resources"):
        for res in details["related_resources"]:
            if isinstance(res, str):
                network_components.append({"id": res, "type": "RELATED_RESOURCE"})

    vulnerability_id = "N/A"
    title = finding.get("title", "")
    if source in ("inspector", "inspector_cis") and title.startswith("CVE-"):
        vulnerability_id = title.split(" ")[0]
    if details.get("check_id") and vulnerability_id == "N/A":
        vulnerability_id = details["check_id"]

    # Build enriched description from dict
    fi = FindingInput(
        title=title,
        description=finding.get("description", ""),
        severity=finding.get("severity", "HIGH"),
        source=source,
        service=service,
        resource_arn=finding.get("resource_arn"),
        resource_id=finding.get("resource_id"),
        region=finding.get("region"),
        recommendation=finding.get("recommendation"),
        details=details,
    )

    return {
        "title": title,
        "severity": finding.get("severity", "HIGH"),
        "finding_type": finding_type,
        "resource_id": finding.get("resource_id") or finding.get("resource_arn") or "N/A",
        "resource_type": resource_type,
        "description": _build_enriched_description(fi),
        "network_components": network_components,
        "vulnerabilities": [],
        "vulnerability_id": vulnerability_id,
    }


def _build_enriched_description(finding: FindingInput) -> str:
    """Build an enriched description with full resource context."""
    parts = [finding.description]

    region = finding.region
    if not region and finding.resource_arn:
        arn_parts = finding.resource_arn.split(":")
        if len(arn_parts) >= 4 and arn_parts[3]:
            region = arn_parts[3]

    account_id = None
    if finding.details:
        account_id = finding.details.get("account_id")
    if not account_id and finding.resource_arn:
        arn_parts = finding.resource_arn.split(":")
        if len(arn_parts) >= 5 and arn_parts[4]:
            account_id = arn_parts[4]

    resource_name_from_arn = None
    if finding.resource_arn:
        arn_parts = finding.resource_arn.split(":")
        if len(arn_parts) >= 6:
            resource_name_from_arn = ":".join(arn_parts[5:])

    context_lines: list[str] = []
    if finding.resource_arn:
        context_lines.append(f"Resource ARN: {finding.resource_arn}")
    if resource_name_from_arn:
        context_lines.append(f"Resource Name/Path: {resource_name_from_arn}")
    if finding.resource_id and finding.resource_id != finding.resource_arn:
        context_lines.append(f"Resource ID: {finding.resource_id}")
    if region:
        context_lines.append(f"AWS Region: {region} (use this exact value in all CLI commands, do NOT use <your-region>)")
    if account_id:
        context_lines.append(f"AWS Account ID: {account_id}")
    if finding.service:
        context_lines.append(f"Service: {finding.service}")
    if finding.details:
        related = finding.details.get("related_resources", [])
        if related:
            for i, rel in enumerate(related):
                if isinstance(rel, str) and rel:
                    context_lines.append(f"Related Resource {i + 1}: {rel}")
        if finding.details.get("check_id"):
            context_lines.append(f"Check ID: {finding.details['check_id']}")
        if finding.details.get("resource_type"):
            context_lines.append(f"Resource Type: {finding.details['resource_type']}")
    if finding.recommendation:
        context_lines.append(f"Original Recommendation: {finding.recommendation}")

    if context_lines:
        parts.append(
            "\n\n--- RESOURCE CONTEXT (use these exact values in remediation commands) ---\n"
            + "\n".join(f"- {line}" for line in context_lines)
        )

    return "".join(parts)


def get_remediation(
    finding_json: dict[str, Any],
    *,
    kb_id: str,
    model_arn: str,
    bedrock_region: str,
) -> str:
    """Send normalized finding JSON to Bedrock RAG and return remediation text."""
    client = boto3.client("bedrock-agent-runtime", region_name=bedrock_region)

    search_query = finding_json.get("title", "Cloud Security Finding")

    prompt_template = f"""You are an AWS Cloud Security Expert. 
Your task is to provide a highly specific remediation plan for the security finding provided in JSON format below.

### KNOWLEDGE BASE CONTEXT:
$search_results$

### INPUT FINDING (FULL JSON):
{json.dumps(finding_json, indent=2)}

### INSTRUCTIONS:
1. Carefully analyze the INPUT FINDING JSON above.
2. Search the KNOWLEDGE BASE CONTEXT for a matching security control or patching procedure.
3. **REMEDIATION STRATEGY**: 
   - If a specific fix exists in the Knowledge Base, use it and cite "Source: Knowledge Base".
   - If the Knowledge Base does not have the specific fix, use your internal expertise to provide the best-practice AWS remediation and cite "Source: Best Practice Recommendation".
4. **TAILORING**: Replace any generic placeholders with the actual 'resource_id', 'vulnerability_id', or 'network_components' found in the JSON.
5. Provide the output in clear, numbered steps.

### TAILORED REMEDIATION:"""

    try:
        response = client.retrieve_and_generate(
            input={"text": search_query},
            retrieveAndGenerateConfiguration={
                "type": "KNOWLEDGE_BASE",
                "knowledgeBaseConfiguration": {
                    "knowledgeBaseId": kb_id,
                    "modelArn": model_arn,
                    "generationConfiguration": {
                        "promptTemplate": {
                            "textPromptTemplate": prompt_template
                        }
                    },
                    "retrievalConfiguration": {
                        "vectorSearchConfiguration": {
                            "numberOfResults": 5
                        }
                    },
                },
            },
        )
        return response["output"]["text"]
    except Exception as e:
        logger.error("bedrock_error: %s", e)
        raise RuntimeError(f"Bedrock API error: {e}") from e
