# Copyright (c) 2020 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com
#

import datetime
import sys

from src.api.v4_autolink import (
    disable_connection_auto_link_v4,
    display_organization_auto_link_status,
    display_organization_failure_analysis,
    enable_connection_auto_link_v4,
    get_organization_auto_link_status_v4,
    get_organization_failure_analysis_v4,
    triage_connection_failures_v4,
    update_connection_auto_link_ignore_list_v4,
)
from src.exceptions import CzOpsApiException
from src.utils.console import eprint


def _get_cost_connections_for_org(organization_id: str, debug: bool = False) -> list:
    """Get COST connections for an org, falling back to global endpoint if needed."""
    from src.api.v4_connections import get_connections, get_connections_for_org

    connections = get_connections_for_org(organization_id, connection_type="COST", debug=debug)
    if connections:
        return connections

    # Org GSI may not index all connections — fall back to global COST endpoint
    if debug:
        eprint("DEBUG: Org endpoint returned no COST connections, trying global endpoint")
    all_cost = get_connections(connection_type="COST", enabled_only=False, debug=debug)
    return [c for c in all_cost if c.cz_organization_id == organization_id]


def _get_payer_accounts(organization_id: str, debug: bool = False) -> dict:
    """Get AWS payer accounts for an organization using connections API."""
    try:
        connections = _get_cost_connections_for_org(organization_id, debug=debug)

        if debug:
            eprint(f"DEBUG: Found {len(connections)} COST connections for org {organization_id}")

        payer_accounts = {}
        for conn in connections:
            conn_dict = conn.model_dump() if hasattr(conn, "model_dump") else conn
            if conn_dict.get("provider") != "AWS":
                continue
            conn_id = conn_dict["connection_id"]
            payer_accounts[conn_id] = conn_dict
            if debug:
                conn_name = conn_dict.get("connection_name", conn_id)
                eprint(f"DEBUG: Found payer account: {conn_id} ({conn_name})")

        return payer_accounts

    except Exception as e:
        eprint(f"ERROR: Cannot retrieve payer accounts: {e}")
        return {}


def _get_auto_link_enabled_connections(organization_id: str) -> dict:
    """Get connections with auto-link enabled for an organization."""
    try:
        auto_link_status = get_organization_auto_link_status_v4(organization_id)
        return auto_link_status.auto_link_enabled_connections
    except CzOpsApiException as e:
        eprint(f"ERROR: Failed to get organization auto-link status: {e}")
        return {}


def enable_auto_link(organization_id: str, role_name: str, arguments: dict | None = None) -> bool:
    """Enable auto-link for payer accounts in an organization."""
    if not role_name:
        eprint("ERROR: You must provide a role name to use for auto-link")
        return False

    debug = arguments.get("--debug", False) if arguments else False

    # Get payer accounts for the organization
    payer_accounts = _get_payer_accounts(organization_id, debug=debug)

    if not payer_accounts:
        eprint(f"No AWS payer accounts found for organization {organization_id}")
        return False

    print(f"Enabling auto-link for organization {organization_id}")
    success_count = 0
    failed_accounts = []

    for connection_pk in payer_accounts:
        # Extract AWS account ID from connection PK (e.g., "095300923743#COST" -> "095300923743")
        aws_account_id = connection_pk.split("#")[0]

        try:
            print(f" - Configuring payer account {aws_account_id} to use auto-link role '{role_name}'")

            result = enable_connection_auto_link_v4(aws_account_id, role_name, force=False)

            if result.get("success") and result.get("auto_link_enabled"):
                print(f" - Auto-link enabled for {aws_account_id}")
                success_count += 1
            else:
                print(f" - Warning: {result}")
                failed_accounts.append(aws_account_id)

        except CzOpsApiException as error:
            eprint(f"ERROR: Enabling auto-link for {aws_account_id} failed: {error}")
            failed_accounts.append(aws_account_id)
        except Exception as error:
            eprint(f"ERROR: Unexpected error enabling auto-link for {aws_account_id}: {error}")
            failed_accounts.append(aws_account_id)

    # Report results
    if success_count > 0:
        next_hour = (datetime.datetime.now() + datetime.timedelta(hours=1)).astimezone().strftime("%H:00 %Z")
        print(
            f"\nAuto-link enabled for {success_count}/{len(payer_accounts)} accounts in organization {organization_id}"
        )
        print(f"Next auto-link account discovery will run at {next_hour}")

        if failed_accounts:
            print(f"Failed accounts: {', '.join(failed_accounts)}")

        return True
    else:
        eprint(f"Failed to enable auto-link for any accounts in organization {organization_id}")
        return False


def disable_auto_link(organization_id: str, arguments: dict | None = None) -> bool:
    """Disable auto-link for all enabled accounts in an organization."""
    auto_link_accounts = _get_auto_link_enabled_connections(organization_id)

    if not auto_link_accounts:
        eprint(f"No auto-link enabled accounts found for organization {organization_id}")
        return False

    print(f"Disabling auto-link for organization {organization_id}")
    success_count = 0
    failed_accounts = []

    for connection_pk in auto_link_accounts:
        # Extract AWS account ID from connection PK (e.g., "095300923743#COST" -> "095300923743")
        aws_account_id = connection_pk.split("#")[0]

        try:
            print(f" - Disabling auto-link on account {aws_account_id}")

            result = disable_connection_auto_link_v4(aws_account_id, force=False)

            if result.get("success") and not result.get("auto_link_enabled"):
                print(f" - Auto-link disabled for {aws_account_id}")
                success_count += 1
            else:
                print(f" - Warning: {result.get('message')}")
                failed_accounts.append(aws_account_id)

        except CzOpsApiException as error:
            eprint(f"ERROR: Disabling auto-link for {aws_account_id} failed: {error}")
            failed_accounts.append(aws_account_id)
        except Exception as error:
            eprint(f"ERROR: Unexpected error disabling auto-link for {aws_account_id}: {error}")
            failed_accounts.append(aws_account_id)

    # Report results
    if success_count > 0:
        print(
            f"\nAuto-link disabled for {success_count}/{len(auto_link_accounts)} accounts in organization {organization_id}"
        )

        if failed_accounts:
            print(f"Failed accounts: {', '.join(failed_accounts)}")

        return True
    else:
        eprint(f"Failed to disable auto-link for any accounts in organization {organization_id}")
        return False


def status_auto_link(organization_id: str, arguments: dict | None = None, org_name: str | None = None) -> bool:
    """Get the auto-link status for an organization."""
    debug = arguments.get("--debug", False) if arguments else False
    show_all = arguments.get("--show-all", False) if arguments else False

    try:
        auto_link_status = get_organization_auto_link_status_v4(organization_id, debug=debug)
        if org_name and not auto_link_status.cz_organization_name:
            auto_link_status.cz_organization_name = org_name

        # Fetch org details and all COST connections for enriched display
        org_details = None
        cost_connections = []
        try:
            from src.api.v4_organizations import get_organizations

            orgs = get_organizations(organization_id, debug=debug)
            if orgs and organization_id in orgs:
                org_details = orgs[organization_id]

            cost_conn_list = _get_cost_connections_for_org(organization_id, debug=debug)
            cost_connections = [c.model_dump() if hasattr(c, "model_dump") else c for c in cost_conn_list]
        except Exception as e:
            eprint(f"Warning: Could not fetch org/connection details: {e}")

        display_organization_auto_link_status(
            auto_link_status, show_all=show_all, org_details=org_details, cost_connections=cost_connections
        )
        return True
    except CzOpsApiException as e:
        eprint(f"ERROR: Failed to get auto-link status: {e}")
        return False


def ignore_auto_link(organization_id: str, ignore_list: str, arguments: dict | None = None) -> bool:
    """Add AWS account IDs to the auto-link ignore list."""
    if not ignore_list:
        eprint("ERROR: No ignore list provided. Please specify accounts to ignore with --ignore-list")
        return False

    # Parse the ignore list from comma-separated string to list
    parsed_ignore_list = [acc.strip() for acc in ignore_list.replace("'", "").split(",") if acc.strip()]

    if not parsed_ignore_list:
        eprint("ERROR: No valid account IDs found in ignore list")
        return False

    auto_link_accounts = _get_auto_link_enabled_connections(organization_id)

    if not auto_link_accounts:
        eprint(f"No auto-link enabled accounts found for organization {organization_id}")
        return False

    print(f"Adding the following accounts to the auto-link ignore list for organization {organization_id}:")
    for account_id in sorted(parsed_ignore_list):
        print(f"  {account_id}")
    print("")

    success_count = 0
    failed_accounts = []

    for connection_pk in auto_link_accounts:
        # Extract AWS account ID from connection PK (e.g., "095300923743#COST" -> "095300923743")
        aws_account_id = connection_pk.split("#")[0]

        try:
            print(f" - Updating ignore list for account {aws_account_id}")

            result = update_connection_auto_link_ignore_list_v4(aws_account_id, parsed_ignore_list, action="add")

            if result.get("success"):
                print(f" - Successfully updated ignore list for {aws_account_id}")
                success_count += 1
            else:
                print(f" - Warning: {result.get('message')}")
                failed_accounts.append(aws_account_id)

        except CzOpsApiException as error:
            eprint(f"ERROR: Updating ignore list for {aws_account_id} failed: {error}")
            failed_accounts.append(aws_account_id)
        except Exception as error:
            eprint(f"ERROR: Unexpected error updating ignore list for {aws_account_id}: {error}")
            failed_accounts.append(aws_account_id)

    # Report results
    if success_count > 0:
        print(
            f"\nAuto-link ignore list updated for {success_count}/{len(auto_link_accounts)} accounts in organization {organization_id}"
        )

        if failed_accounts:
            print(f"Failed accounts: {', '.join(failed_accounts)}")

        return True
    else:
        eprint(f"Failed to update auto-link ignore list for any accounts in organization {organization_id}")
        return False


def failures_auto_link(organization_id: str, arguments: dict | None = None, org_name: str | None = None) -> bool:
    """Get failure analysis for an organization's auto-link connections."""
    debug = arguments.get("--debug", False) if arguments else False
    show_all = arguments.get("--show-all", False) if arguments else False

    try:
        analysis = get_organization_failure_analysis_v4(organization_id, debug=debug)
        display_organization_failure_analysis(analysis, org_name=org_name, show_all=show_all)
        return True
    except CzOpsApiException as e:
        eprint(f"ERROR: Failed to get failure analysis: {e}")
        return False


def triage_auto_link(organization_id: str, arguments: dict) -> bool:
    """Triage auto-link failures across all enabled connections in an organization."""
    from src.schemas.autolink import TriageFailuresRequest

    debug = arguments.get("--debug", False)

    # Determine action from flags
    if arguments.get("--ignore"):
        action = "ignore"
    elif arguments.get("--retry"):
        action = "retry"
    elif arguments.get("--retry-all"):
        action = "retry_all"
    else:
        eprint("ERROR: You must specify one of --ignore, --retry, or --retry-all")
        return False

    # Parse account IDs if provided
    account_ids = None
    if arguments.get("--accounts"):
        account_ids = [a.strip() for a in arguments["--accounts"].split(",") if a.strip()]

    # Build and validate the request
    try:
        request = TriageFailuresRequest(action=action, account_ids=account_ids)
    except Exception as e:
        eprint(f"ERROR: Invalid triage request: {e}")
        return False

    # Get all auto-link enabled connections
    auto_link_connections = _get_auto_link_enabled_connections(organization_id)
    if not auto_link_connections:
        eprint(f"No auto-link enabled connections found for organization {organization_id}")
        return False

    print(f"Triaging failures ({action}) for organization {organization_id}")
    success_count = 0
    failed_connections = []

    for connection_pk in auto_link_connections:
        aws_account_id = connection_pk.split("#")[0]

        try:
            if debug:
                eprint(f"DEBUG: Triaging connection {aws_account_id} with action={action}")

            result = triage_connection_failures_v4(aws_account_id, request, debug=debug)

            if result.success:
                print(f" - {aws_account_id}: {result.message}")
                success_count += 1
            else:
                print(f" - {aws_account_id}: FAILED - {result.message}")
                failed_connections.append(aws_account_id)

        except CzOpsApiException as error:
            eprint(f"ERROR: Triage for {aws_account_id} failed: {error}")
            failed_connections.append(aws_account_id)

    # Report results
    total = len(auto_link_connections)
    if success_count > 0:
        print(f"\nTriage completed for {success_count}/{total} connections")
        if failed_connections:
            print(f"Failed connections: {', '.join(failed_connections)}")
        return True
    else:
        eprint(f"Failed to triage any connections in organization {organization_id}")
        return False


def functions(arguments, organization_data):
    """Entry point for auto-link commands, passing through arguments from CLI."""
    for cz_organization_id, details in organization_data.items():
        if arguments["enable"]:
            result = enable_auto_link(cz_organization_id, arguments["<ROLE-NAME>"], arguments)
        elif arguments["disable"]:
            result = disable_auto_link(cz_organization_id, arguments)
        elif arguments["status"]:
            org_name = getattr(details, "cz_organization_name", None)
            result = status_auto_link(cz_organization_id, arguments, org_name=org_name)
        elif arguments["ignore"]:
            result = ignore_auto_link(cz_organization_id, arguments["--ignore-list"], arguments)
        elif arguments["failures"]:
            org_name = getattr(details, "cz_organization_name", None)
            result = failures_auto_link(cz_organization_id, arguments, org_name=org_name)
        elif arguments["triage"]:
            result = triage_auto_link(cz_organization_id, arguments)
        else:
            eprint("Unsupported")
            sys.exit(-1)

        return result
    return None
