#!/usr/bin/env python3
"""
Drift detection for all stacks and StackSets managed by this repo.

Compares the local templates and parameters against what is deployed
in AWS, reporting any differences.
"""

import json
import os
import sys
from pathlib import Path
from typing import Any, Optional

import boto3
import click
from botocore.exceptions import ClientError
from dotenv import load_dotenv

from .stackset_registry import get_stackset_registry

# Import config_loader directly to avoid triggering ai_org.__init__
# which requires the package to be installed.
sys.path.insert(0, str(Path(__file__).parent.parent))
import importlib.util

_spec = importlib.util.spec_from_file_location(
    "config_loader",
    Path(__file__).parent.parent / "src" / "ai_org" / "utils" / "config_loader.py",
)
_mod = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_mod)  # type: ignore[union-attr]
load_config_value = _mod.load_config_value

load_dotenv()

# Configuration
GITHUB_ORG = load_config_value("GH_ACCOUNT")
GITHUB_REPO = load_config_value("GH_REPO")
GITHUB_ORG_SECONDARY = os.getenv("GH_ACCOUNT_SECONDARY", "")
NOTIFICATIONS_EMAIL = load_config_value("NOTIFICATIONS_EMAIL")
BUDGETS_MONTHLY_DEFAULT = load_config_value("BUDGETS_MONTHLY_DEFAULT")
BUDGETS_ANOMALY_THRESHOLD = load_config_value("BUDGETS_ANOMALY_THRESHOLD")

HOME_REGION = "us-east-1"
PROJECT_ROOT = Path(__file__).parent.parent
STACKSETS_DIR = PROJECT_ROOT / "stacksets"


def _normalise_template(body: str) -> str:
    """Strip trailing whitespace and ensure consistent newline ending."""
    return "\n".join(line.rstrip() for line in body.splitlines()).strip() + "\n"


def _params_to_dict(cf_params: list[dict[str, str]]) -> dict[str, str]:
    """Convert CloudFormation parameter list to a plain dict."""
    return {p["ParameterKey"]: p["ParameterValue"] for p in cf_params}


class DriftDetector:
    """Detects drift between local definitions and deployed AWS resources."""

    def __init__(self, profile: Optional[str] = None):
        session = boto3.Session(profile_name=profile) if profile else boto3.Session()
        self.cf = session.client("cloudformation", region_name=HOME_REGION)
        self.org = session.client("organizations")
        self.sts = session.client("sts")
        self.iam = session.client("iam")

        identity = self.sts.get_caller_identity()
        self.account_id = identity["Account"]
        self.drifts: list[dict[str, Any]] = []

    # ------------------------------------------------------------------
    # org-pipeline-role (regular CF stack on management account)
    # ------------------------------------------------------------------
    def check_pipeline_role_stack(self) -> None:
        """Check the org-pipeline-role stack for template and parameter drift."""
        stack_name = "org-pipeline-role"
        template_path = STACKSETS_DIR / "pipeline-role.yaml"
        click.echo(f"\n--- {stack_name} (CF stack) ---")

        if not template_path.exists():
            self._record(stack_name, "MISSING_LOCAL", "Local template not found")
            return

        # Check stack exists
        try:
            stack = self.cf.describe_stacks(StackName=stack_name)["Stacks"][0]
        except ClientError as e:
            if "does not exist" in str(e):
                self._record(stack_name, "NOT_DEPLOYED", "Stack does not exist in AWS")
                return
            raise

        # Compare parameters
        deployed_params = _params_to_dict(stack.get("Parameters", []))
        expected_params = {"GitHubOrg": GITHUB_ORG, "GitHubRepo": GITHUB_REPO}

        for key, expected in expected_params.items():
            deployed = deployed_params.get(key)
            if deployed != expected:
                self._record(
                    stack_name,
                    "PARAMETER_DRIFT",
                    f"{key}: deployed={deployed!r}, expected={expected!r}",
                )

        # Compare template body
        try:
            deployed_body = self.cf.get_template(StackName=stack_name, TemplateStage="Original")[
                "TemplateBody"
            ]
            if isinstance(deployed_body, dict):
                # CloudFormation sometimes returns parsed JSON
                deployed_body = json.dumps(deployed_body, indent=2)
        except ClientError:
            self._record(stack_name, "TEMPLATE_UNREADABLE", "Cannot read deployed template")
            return

        local_body = template_path.read_text()
        if _normalise_template(deployed_body) != _normalise_template(local_body):
            self._record(stack_name, "TEMPLATE_DRIFT", "Deployed template differs from local file")

        if not any(d["resource"] == stack_name for d in self.drifts):
            click.echo("  OK (no drift)")

    # ------------------------------------------------------------------
    # StackSets from the registry
    # ------------------------------------------------------------------
    def check_stacksets(self) -> None:
        """Check all StackSets in the registry for drift."""
        budgets_monthly = int(BUDGETS_MONTHLY_DEFAULT) if BUDGETS_MONTHLY_DEFAULT else 1000
        budgets_anomaly = int(BUDGETS_ANOMALY_THRESHOLD) if BUDGETS_ANOMALY_THRESHOLD else 100

        registry = get_stackset_registry(
            github_org=GITHUB_ORG or "",
            github_repo=GITHUB_REPO or "",
            notifications_email=NOTIFICATIONS_EMAIL or "",
            budgets_monthly_default=budgets_monthly,
            budgets_anomaly_threshold=budgets_anomaly,
            github_org_secondary=GITHUB_ORG_SECONDARY,
        )

        for name, config in sorted(registry.items(), key=lambda x: x[1].get("priority", 999)):
            self._check_one_stackset(name, config)

    def _check_one_stackset(self, name: str, config: dict[str, Any]) -> None:
        click.echo(f"\n--- {name} (StackSet, {config['deployment_mode']}) ---")

        template_path = STACKSETS_DIR / config["template"]
        if not template_path.exists():
            self._record(name, "MISSING_LOCAL", f"Template not found: {template_path}")
            return

        # Check StackSet exists
        try:
            ss = self.cf.describe_stack_set(StackSetName=name)["StackSet"]
        except ClientError as e:
            if "StackSetNotFoundException" in str(e):
                self._record(name, "NOT_DEPLOYED", "StackSet does not exist in AWS")
                return
            raise

        # Compare parameters
        deployed_params = _params_to_dict(ss.get("Parameters", []))
        expected_params = config.get("parameters") or {}

        for key, expected in expected_params.items():
            deployed = deployed_params.get(key)
            if deployed != expected:
                self._record(
                    name,
                    "PARAMETER_DRIFT",
                    f"{key}: deployed={deployed!r}, expected={expected!r}",
                )

        # Flag unexpected parameters deployed but not in local config
        # Skip for manual-deploy StackSets -- parameters are provided at deploy time
        if config["deployment_mode"] != "manual":
            for key in deployed_params:
                if key not in expected_params:
                    self._record(
                        name,
                        "EXTRA_PARAMETER",
                        f"{key}={deployed_params[key]!r} deployed but not in registry",
                    )

        # Compare template body
        deployed_body = ss.get("TemplateBody", "")
        if isinstance(deployed_body, dict):
            deployed_body = json.dumps(deployed_body, indent=2)

        local_body = template_path.read_text()
        if _normalise_template(deployed_body) != _normalise_template(local_body):
            self._record(name, "TEMPLATE_DRIFT", "Deployed template differs from local file")

        # Check auto-deployment setting
        auto_deploy = ss.get("AutoDeployment", {})
        expected_auto = config["deployment_mode"] == "auto"
        deployed_auto = auto_deploy.get("Enabled", False)
        if deployed_auto != expected_auto:
            self._record(
                name,
                "AUTO_DEPLOY_DRIFT",
                f"AutoDeployment: deployed={deployed_auto}, expected={expected_auto}",
            )

        # Check instances for auto-deploy StackSets
        if config["deployment_mode"] == "auto" and config.get("target_ous"):
            instances = self.cf.list_stack_instances(StackSetName=name, CallAs="SELF").get(
                "Summaries", []
            )
            if not instances:
                self._record(name, "NO_INSTANCES", "Auto-deploy StackSet has no instances")
            else:
                # Check for failed instances
                for inst in instances:
                    status = inst.get("StackInstanceStatus", {}).get("DetailedStatus", "")
                    if inst.get("Status") == "OUTDATED" or "FAILED" in status.upper():
                        acct = inst.get("Account", "?")
                        self._record(
                            name,
                            "INSTANCE_DRIFT",
                            f"Account {acct}: status={inst.get('Status')}, "
                            f"reason={inst.get('StatusReason', 'n/a')[:80]}",
                        )

        if not any(d["resource"] == name for d in self.drifts):
            click.echo("  OK (no drift)")

    # ------------------------------------------------------------------
    # Service Control Policies
    # ------------------------------------------------------------------
    def check_scps(self) -> None:
        """Check SCP content drift."""
        scp_path = STACKSETS_DIR / "scps" / "workloads-baseline.json"
        name = "workloads-baseline"
        click.echo(f"\n--- {name} (SCP) ---")

        if not scp_path.exists():
            self._record(name, "MISSING_LOCAL", "Local SCP file not found")
            return

        local_content = json.loads(scp_path.read_text())

        # Find the policy in AWS
        policies = self.org.list_policies(Filter="SERVICE_CONTROL_POLICY")["Policies"]
        policy_id = None
        for p in policies:
            if p["Name"] == name:
                policy_id = p["Id"]
                break

        if not policy_id:
            self._record(name, "NOT_DEPLOYED", "SCP not found in AWS Organizations")
            return

        # Compare content
        deployed = self.org.describe_policy(PolicyId=policy_id)["Policy"]
        deployed_content = json.loads(deployed["Content"])

        if json.dumps(local_content, sort_keys=True) != json.dumps(
            deployed_content, sort_keys=True
        ):
            self._record(name, "CONTENT_DRIFT", "Deployed SCP content differs from local file")

        # Check attachment
        targets = self.org.list_targets_for_policy(PolicyId=policy_id)["Targets"]
        if not targets:
            self._record(name, "NOT_ATTACHED", "SCP exists but is not attached to any OU")

        if not any(d["resource"] == name for d in self.drifts):
            click.echo("  OK (no drift)")

    # ------------------------------------------------------------------
    # OIDC trust policy sanity check
    # ------------------------------------------------------------------
    def check_oidc_trust(self) -> None:
        """Verify OrgPipelineRole trusts the current repo."""
        click.echo("\n--- OrgPipelineRole OIDC trust ---")
        try:
            role = self.iam.get_role(RoleName="OrgPipelineRole")["Role"]
        except ClientError:
            self._record("OrgPipelineRole", "NOT_FOUND", "IAM role does not exist")
            return

        trust = role["AssumeRolePolicyDocument"]
        expected_sub = f"repo:{GITHUB_ORG}/{GITHUB_REPO}:ref:refs/heads/main"

        found = False
        for stmt in trust.get("Statement", []):
            subs = (
                stmt.get("Condition", {})
                .get("StringLike", {})
                .get("token.actions.githubusercontent.com:sub", [])
            )
            if isinstance(subs, str):
                subs = [subs]
            for sub in subs:
                if sub == expected_sub:
                    found = True
                    break

        if not found:
            deployed_subs = []
            for stmt in trust.get("Statement", []):
                deployed_subs.extend(
                    stmt.get("Condition", {})
                    .get("StringLike", {})
                    .get("token.actions.githubusercontent.com:sub", [])
                )
            self._record(
                "OrgPipelineRole",
                "TRUST_DRIFT",
                f"Expected trust for {expected_sub!r}, found: {deployed_subs}",
            )
        else:
            click.echo("  OK (trusts current repo)")

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------
    def _record(self, resource: str, drift_type: str, detail: str) -> None:
        self.drifts.append({"resource": resource, "type": drift_type, "detail": detail})
        click.echo(f"  DRIFT [{drift_type}] {detail}")

    def summary(self) -> int:
        """Print summary and return exit code (0 = clean, 1 = drift found)."""
        click.echo("\n" + "=" * 60)
        if not self.drifts:
            click.echo("  NO DRIFT DETECTED -- local and deployed state match")
            click.echo("=" * 60)
            return 0

        click.echo(f"  DRIFT DETECTED -- {len(self.drifts)} issue(s) found")
        click.echo("=" * 60)
        for d in self.drifts:
            click.echo(f"  [{d['type']:20s}] {d['resource']}: {d['detail']}")
        click.echo("")
        return 1


@click.command()
@click.option("--profile", default=os.getenv("AWS_PROFILE"), help="AWS profile to use")
@click.option("--json-output", is_flag=True, help="Output results as JSON")
def main(profile: str, json_output: bool) -> None:
    """Detect drift between local definitions and deployed AWS resources."""
    if not GITHUB_ORG or not GITHUB_REPO:
        click.echo(
            "GH_ACCOUNT and GH_REPO are required. "
            "Set via environment, .env, or ~/.augint/.env.augint-org",
            err=True,
        )
        sys.exit(1)

    detector = DriftDetector(profile=profile)

    click.echo("=" * 60)
    click.echo("  DRIFT DETECTION")
    click.echo("=" * 60)
    click.echo(f"  Repo identity: {GITHUB_ORG}/{GITHUB_REPO}")
    click.echo(f"  Account: {detector.account_id}")

    detector.check_oidc_trust()
    detector.check_pipeline_role_stack()
    detector.check_stacksets()
    detector.check_scps()

    if json_output:
        click.echo(json.dumps(detector.drifts, indent=2))

    exit_code = detector.summary()
    sys.exit(exit_code)


if __name__ == "__main__":
    main()
