"""
Remediation report generation from scan results.

This module provides utilities to generate detailed remediation reports
from compliance scan results, with step-by-step instructions and AWS CLI commands.
"""

import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, Field

from complio.config.settings import get_settings
from complio.core.runner import ScanResults
from complio.tests_library.base import Finding, TestResult
from complio.utils.json_encoder import ComplioJSONEncoder
from complio.utils.logger import get_logger

logger = get_logger(__name__)


# ISO 27001 control mapping for tests
ISO_CONTROL_MAPPING = {
    "s3_encryption": "A.10.1.1",
    "ec2_security_groups": "A.13.1.1",
    "iam_password_policy": "A.9.4.3",
    "cloudtrail_logging": "A.12.4.1",
    "ebs_encryption": "A.8.24",
    "rds_encryption": "A.8.24",
    "secrets_manager_encryption": "A.8.24",
    "s3_public_access_block": "A.8.11",
    "cloudtrail_log_validation": "A.8.15",
    "cloudtrail_encryption": "A.8.15",
    "vpc_flow_logs": "A.8.16",
    "nacl_security": "A.8.20",
    "redshift_encryption": "A.8.24",
    "efs_encryption": "A.8.24",
    "dynamodb_encryption": "A.8.24",
    "elasticache_encryption": "A.8.24",
    "kms_key_rotation": "A.8.24",
    "access_key_rotation": "A.8.5",
    "mfa_enforcement": "A.8.5",
    "root_account_protection": "A.8.2",
    "s3_versioning": "A.8.13",
    "backup_encryption": "A.8.24",
    "cloudwatch_retention": "A.8.15",
    "sns_encryption": "A.8.24",
    "cloudwatch_logs_encryption": "A.8.24",
    "vpn_security": "A.8.22",
    "nacl_configuration": "A.8.20",
    "alb_nlb_security": "A.8.22",
    "cloudfront_https": "A.8.24",
    "transit_gateway_security": "A.8.22",
    "vpc_endpoints_security": "A.8.22",
    "network_firewall": "A.8.20",
    "direct_connect_security": "A.8.22",
    "cloudwatch_alarms": "A.8.16",
    "config_enabled": "A.8.16",
    "waf_configuration": "A.8.20",
    "api_gateway_security": "A.8.22",
    "guardduty_enabled": "A.8.16",
    "security_hub_enabled": "A.8.16",
    "eventbridge_rules": "A.8.16",
}


class RemediationResource(BaseModel):
    """Resource affected by a finding."""
    id: str = Field(..., description="Resource identifier")
    type: str = Field(..., description="Resource type")


class RemediationSteps(BaseModel):
    """Detailed remediation instructions."""
    summary: str = Field(..., description="Brief remediation summary")
    steps: List[str] = Field(..., description="Step-by-step instructions")
    aws_cli: Optional[str] = Field(None, description="AWS CLI command if applicable")
    aws_console_path: Optional[str] = Field(None, description="Path in AWS Console")
    estimated_effort: str = Field(..., description="Estimated time to remediate")
    documentation_url: Optional[str] = Field(None, description="AWS documentation link")


class RemediationItem(BaseModel):
    """Single remediation item."""
    priority: int = Field(..., description="Priority (1=highest)")
    severity: str = Field(..., description="Finding severity")
    title: str = Field(..., description="Finding title")
    test_id: str = Field(..., description="Test identifier")
    iso_control: str = Field(..., description="ISO 27001 control")
    resources: List[RemediationResource] = Field(..., description="Affected resources")
    description: str = Field(..., description="Detailed description")
    risk: str = Field(..., description="Risk if not remediated")
    remediation: RemediationSteps = Field(..., description="Remediation instructions")


class RemediationReport(BaseModel):
    """Complete remediation report."""
    scan_id: str = Field(..., description="Scan identifier")
    generated_at: str = Field(..., description="Report generation timestamp")
    region: str = Field(..., description="AWS region")
    account_id: str = Field(..., description="AWS account ID (masked)")
    compliance_score: float = Field(..., description="Overall compliance score")
    total_findings: int = Field(..., description="Total number of findings")
    summary: Dict[str, int] = Field(..., description="Findings count by severity")
    remediations: List[RemediationItem] = Field(..., description="Remediation items")


def _get_remediation_details(test_id: str, finding: Finding) -> RemediationSteps:
    """Get detailed remediation steps for a specific test and finding.

    Args:
        test_id: Test identifier
        finding: Finding object

    Returns:
        RemediationSteps with detailed instructions
    """
    # This is a comprehensive mapping of remediation steps for each test
    # In a production system, this would be in a database or configuration file

    remediation_db = {
        "root_account_protection": RemediationSteps(
            summary="Enable MFA for the root account",
            steps=[
                "Sign in to AWS Console as root user",
                "Click on account name → Security credentials",
                "Under 'Multi-factor authentication (MFA)', click 'Assign MFA device'",
                "Choose virtual MFA device (recommended) or hardware MFA device",
                "Install an authenticator app (Google Authenticator, Authy, etc.)",
                "Scan the QR code and enter two consecutive MFA codes"
            ],
            aws_cli=None,
            aws_console_path="IAM Console → Security credentials → MFA",
            estimated_effort="5 minutes",
            documentation_url="https://docs.aws.amazon.com/IAM/latest/UserGuide/id_root-user.html#id_root-user_manage_mfa"
        ),
        "s3_public_access_block": RemediationSteps(
            summary="Enable public access block for S3 buckets",
            steps=[
                "Go to AWS S3 console",
                "Select the bucket",
                "Go to 'Permissions' tab",
                "Under 'Block public access', click 'Edit'",
                "Enable all four settings",
                "Click 'Save changes'"
            ],
            aws_cli="aws s3api put-public-access-block --bucket <BUCKET_NAME> --public-access-block-configuration BlockPublicAcls=true,IgnorePublicAcls=true,BlockPublicPolicy=true,RestrictPublicBuckets=true",
            aws_console_path="S3 Console → Bucket → Permissions → Block public access",
            estimated_effort="2 minutes per bucket",
            documentation_url="https://docs.aws.amazon.com/AmazonS3/latest/userguide/access-control-block-public-access.html"
        ),
        "vpc_flow_logs": RemediationSteps(
            summary="Enable VPC Flow Logs",
            steps=[
                "Go to VPC console",
                "Select the VPC",
                "Click Actions → Create flow log",
                "Choose 'All' for filter (Accept and Reject traffic)",
                "Select destination (CloudWatch Logs or S3)",
                "Choose or create IAM role with appropriate permissions",
                "Click 'Create flow log'"
            ],
            aws_cli="aws ec2 create-flow-logs --resource-type VPC --resource-ids <VPC_ID> --traffic-type ALL --log-destination-type cloud-watch-logs --log-group-name /aws/vpc/flowlogs --deliver-logs-permission-arn <IAM_ROLE_ARN>",
            aws_console_path="VPC Console → VPC → Actions → Create flow log",
            estimated_effort="5 minutes per VPC",
            documentation_url="https://docs.aws.amazon.com/vpc/latest/userguide/flow-logs.html"
        ),
        "config_enabled": RemediationSteps(
            summary="Enable AWS Config",
            steps=[
                "Go to AWS Config console",
                "Click 'Get started'",
                "Choose 'Record all resources supported in this region'",
                "Select or create S3 bucket for config storage",
                "Select or create SNS topic for notifications",
                "Choose or create IAM role",
                "Click 'Next' and 'Confirm'"
            ],
            aws_cli="aws configservice put-configuration-recorder --configuration-recorder name=default,roleARN=<ROLE_ARN> && aws configservice put-delivery-channel --delivery-channel name=default,s3BucketName=<BUCKET_NAME> && aws configservice start-configuration-recorder --configuration-recorder-name default",
            aws_console_path="AWS Config Console → Get started",
            estimated_effort="10 minutes",
            documentation_url="https://docs.aws.amazon.com/config/latest/developerguide/gs-console.html"
        ),
        "security_hub_enabled": RemediationSteps(
            summary="Enable AWS Security Hub",
            steps=[
                "Go to AWS Security Hub console",
                "Click 'Enable Security Hub'",
                "Choose security standards to enable (CIS, PCI-DSS, etc.)",
                "Review and enable standards",
                "Configure integrations with other AWS services"
            ],
            aws_cli="aws securityhub enable-security-hub",
            aws_console_path="Security Hub Console → Enable Security Hub",
            estimated_effort="5 minutes",
            documentation_url="https://docs.aws.amazon.com/securityhub/latest/userguide/securityhub-enable.html"
        ),
        "s3_versioning": RemediationSteps(
            summary="Enable S3 bucket versioning",
            steps=[
                "Go to AWS S3 console",
                "Select the bucket",
                "Go to 'Properties' tab",
                "Under 'Bucket Versioning', click 'Edit'",
                "Select 'Enable'",
                "Click 'Save changes'"
            ],
            aws_cli="aws s3api put-bucket-versioning --bucket <BUCKET_NAME> --versioning-configuration Status=Enabled",
            aws_console_path="S3 Console → Bucket → Properties → Bucket Versioning",
            estimated_effort="2 minutes per bucket",
            documentation_url="https://docs.aws.amazon.com/AmazonS3/latest/userguide/Versioning.html"
        ),
        "cloudtrail_encryption": RemediationSteps(
            summary="Enable CloudTrail log encryption with KMS",
            steps=[
                "Go to CloudTrail console",
                "Select the trail",
                "Click 'Edit'",
                "Under 'Log file SSE-KMS encryption', select 'Enabled'",
                "Choose existing KMS key or create new one",
                "Update trail policy if needed",
                "Click 'Save changes'"
            ],
            aws_cli="aws cloudtrail update-trail --name <TRAIL_NAME> --kms-key-id <KMS_KEY_ID>",
            aws_console_path="CloudTrail Console → Trails → Trail → Edit",
            estimated_effort="5 minutes",
            documentation_url="https://docs.aws.amazon.com/awscloudtrail/latest/userguide/encrypting-cloudtrail-log-files-with-aws-kms.html"
        ),
        "cloudwatch_logs_encryption": RemediationSteps(
            summary="Enable CloudWatch Logs encryption",
            steps=[
                "Go to CloudWatch console",
                "Select 'Log groups'",
                "Select the log group",
                "Click Actions → Edit",
                "Under 'KMS key ARN', select or enter KMS key",
                "Click 'Save'"
            ],
            aws_cli="aws logs associate-kms-key --log-group-name <LOG_GROUP_NAME> --kms-key-id <KMS_KEY_ID>",
            aws_console_path="CloudWatch Console → Log groups → Log group → Actions → Edit",
            estimated_effort="3 minutes per log group",
            documentation_url="https://docs.aws.amazon.com/AmazonCloudWatch/latest/logs/encrypt-log-data-kms.html"
        ),
        "nacl_security": RemediationSteps(
            summary="Restrict Network ACL rules to specific IP ranges",
            steps=[
                "Go to VPC console",
                "Select 'Network ACLs'",
                "Select the NACL",
                "Review inbound and outbound rules",
                "Edit overly permissive rules (0.0.0.0/0)",
                "Replace with specific IP ranges or security groups",
                "Click 'Save rules'"
            ],
            aws_cli="aws ec2 replace-network-acl-entry --network-acl-id <NACL_ID> --rule-number <RULE_NUMBER> --protocol <PROTOCOL> --rule-action allow --cidr-block <SPECIFIC_CIDR>",
            aws_console_path="VPC Console → Network ACLs → NACL → Inbound/Outbound rules",
            estimated_effort="10 minutes per NACL",
            documentation_url="https://docs.aws.amazon.com/vpc/latest/userguide/vpc-network-acls.html"
        ),
        "eventbridge_rules": RemediationSteps(
            summary="Configure EventBridge security monitoring rules",
            steps=[
                "Go to EventBridge console",
                "Click 'Create rule'",
                "Define event pattern for security events (root login, policy changes, etc.)",
                "Add target (SNS topic, Lambda function, CloudWatch Logs)",
                "Configure event pattern for specific security scenarios",
                "Enable the rule"
            ],
            aws_cli="aws events put-rule --name <RULE_NAME> --event-pattern '{\"source\":[\"aws.iam\"],\"detail-type\":[\"AWS API Call via CloudTrail\"]}' && aws events put-targets --rule <RULE_NAME> --targets Id=1,Arn=<SNS_TOPIC_ARN>",
            aws_console_path="EventBridge Console → Rules → Create rule",
            estimated_effort="15 minutes",
            documentation_url="https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-create-rule.html"
        ),
    }

    # Get remediation or use generic from finding.remediation
    if test_id in remediation_db:
        return remediation_db[test_id]
    else:
        # Generic remediation from finding
        return RemediationSteps(
            summary=finding.remediation,
            steps=[finding.remediation],
            aws_cli=None,
            aws_console_path=None,
            estimated_effort="Unknown",
            documentation_url=None
        )


def _extract_risk_description(test_id: str, finding: Finding) -> str:
    """Extract risk description for a finding.

    Args:
        test_id: Test identifier
        finding: Finding object

    Returns:
        Risk description string
    """
    risk_db = {
        "root_account_protection": "If root credentials are stolen, attackers gain complete control of the AWS account.",
        "s3_public_access_block": "Data could be accidentally exposed to the internet.",
        "vpc_flow_logs": "Network traffic is not logged, making it difficult to detect and investigate security incidents.",
        "config_enabled": "Configuration changes are not tracked, making compliance auditing impossible.",
        "security_hub_enabled": "Security findings from multiple services are not centralized, making threat detection difficult.",
        "s3_versioning": "Accidental deletions or modifications cannot be recovered.",
        "cloudtrail_encryption": "CloudTrail logs could be read by unauthorized users if S3 bucket is compromised.",
        "cloudwatch_logs_encryption": "Log data could be exposed if storage is compromised.",
        "nacl_security": "Unrestricted network access allows potential unauthorized access to resources.",
        "eventbridge_rules": "Security events may not be detected or responded to in a timely manner.",
    }

    return risk_db.get(test_id, "Non-compliance with ISO 27001 requirements.")


def generate_remediation_report(
    scan_id: str,
    results: ScanResults,
    region: str,
    account_id: str,
    output_path: Optional[Path] = None,
) -> Path:
    """Generate remediation report JSON from scan results.

    Args:
        scan_id: Scan identifier
        results: Scan results
        region: AWS region
        account_id: AWS account ID (masked)
        output_path: Optional output path (defaults to ~/.complio/reports/)

    Returns:
        Path to generated JSON file

    Example:
        >>> from complio.core.remediation import generate_remediation_report
        >>> path = generate_remediation_report(scan_id, results, "us-east-1", "****1234")
        >>> print(f"Report: {path}")
    """
    settings = get_settings()

    # Determine output path
    if output_path is None:
        settings.report_dir.mkdir(parents=True, exist_ok=True)
        output_path = settings.report_dir / f"{scan_id}_remediation.json"

    # Collect all findings
    all_findings: List[tuple[str, TestResult, Finding]] = []
    for test_result in results.test_results:
        for finding in test_result.findings:
            all_findings.append((test_result.test_id, test_result, finding))

    # Count by severity
    severity_counts = {"critical": 0, "high": 0, "medium": 0, "low": 0}
    for _, _, finding in all_findings:
        if finding.severity in severity_counts:
            severity_counts[finding.severity] += 1

    # Group findings by (test_id, title) to consolidate resources
    from collections import defaultdict
    grouped_findings: Dict[tuple, Dict] = defaultdict(lambda: {
        "test_id": "",
        "test_result": None,
        "finding": None,
        "resources": []
    })

    for test_id, test_result, finding in all_findings:
        key = (test_id, finding.title)
        group = grouped_findings[key]
        group["test_id"] = test_id
        group["test_result"] = test_result
        group["finding"] = finding
        group["resources"].append({
            "id": finding.resource_id,
            "type": finding.resource_type
        })

    # Create remediation items
    remediation_items = []
    for (test_id, title), data in grouped_findings.items():
        finding = data["finding"]
        test_result = data["test_result"]

        # Get detailed remediation
        remediation_steps = _get_remediation_details(test_id, finding)

        # Get ISO control
        iso_control = ISO_CONTROL_MAPPING.get(test_id, "N/A")

        # Get risk description
        risk = _extract_risk_description(test_id, finding)

        remediation_item = RemediationItem(
            priority=0,  # Will be set later based on severity
            severity=finding.severity,
            title=finding.title,
            test_id=test_id,
            iso_control=iso_control,
            resources=[RemediationResource(**r) for r in data["resources"]],
            description=finding.description,
            risk=risk,
            remediation=remediation_steps
        )

        remediation_items.append(remediation_item)

    # Sort by severity (critical > high > medium > low) and assign priorities
    severity_order = {"critical": 0, "high": 1, "medium": 2, "low": 3}
    remediation_items.sort(key=lambda x: (severity_order.get(x.severity, 99), x.title))

    for idx, item in enumerate(remediation_items, 1):
        item.priority = idx

    # Create report
    report = RemediationReport(
        scan_id=scan_id,
        generated_at=datetime.now(timezone.utc).isoformat(),
        region=region,
        account_id=account_id,
        compliance_score=results.overall_score,
        total_findings=len(all_findings),
        summary=severity_counts,
        remediations=remediation_items
    )

    # Write JSON
    with open(output_path, 'w') as f:
        json.dump(report.model_dump(), f, cls=ComplioJSONEncoder, indent=2)

    logger.info("remediation_report_generated", path=str(output_path), findings=len(all_findings))

    return output_path


def load_scan_from_history(scan_id: str) -> tuple[ScanResults, str, str]:
    """Load scan results from history.

    Args:
        scan_id: Scan identifier

    Returns:
        Tuple of (ScanResults, region, account_id)

    Raises:
        FileNotFoundError: If scan not found in history
    """
    settings = get_settings()
    history_dir = settings.config_dir / "history"

    # Try to find scan file
    scan_file = history_dir / f"{scan_id}.json"

    if not scan_file.exists():
        raise FileNotFoundError(f"Scan '{scan_id}' not found in {history_dir}")

    # Load JSON
    try:
        with open(scan_file, 'r') as f:
            data = json.load(f)
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON in scan file: {e}")

    # Extract region and account_id from metadata
    region = data.get("region", "unknown")
    account_id = data.get("account_id", "unknown")

    # Parse scan results with better error handling
    try:
        from complio.tests_library.base import TestResult

        results_data = data["results"]

        # Reconstruct TestResult objects from dicts
        # test_results are stored as dicts, need to convert back to TestResult objects
        test_results = [TestResult(**tr) for tr in results_data["test_results"]]

        # Create ScanResults with reconstructed test_results
        results = ScanResults(
            test_results=test_results,
            total_tests=results_data.get("total_tests", 0),
            passed_tests=results_data.get("passed_tests", 0),
            failed_tests=results_data.get("failed_tests", 0),
            error_tests=results_data.get("error_tests", 0),
            overall_score=results_data.get("overall_score", 0.0),
            execution_time=results_data.get("execution_time", 0.0),
            region=results_data.get("region", ""),
            timestamp=results_data.get("timestamp", 0.0),
        )

    except KeyError as e:
        # Provide helpful error message about missing keys
        available_keys = list(data.keys())
        raise KeyError(
            f"Key {e} not found in scan data. "
            f"Available keys: {', '.join(available_keys)}. "
            f"This may be due to an old scan format. Try running a new scan."
        )
    except Exception as e:
        raise ValueError(f"Failed to parse scan results: {e}")

    return results, region, account_id
