# Copyright (c) 2022 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 json
from collections import Counter
from typing import Any

from rich import box
from rich.console import Console
from rich.panel import Panel
from rich.table import Column, Table
from rich.text import Text
from toolz import get_in

from src.constants import CONNECTION_PROPERTIES, CONNECTION_TYPES, LINKS_TO_IGNORE, NO_ACCOUNT_NAME
from src.utils.connections import build_cur_s3_path

# Initialize Rich console with full terminal width
console = Console(width=None)


def format_account_stage(stage: str) -> str:
    """
    Format account stage with appropriate colors.

    Args:
        stage: The account stage to format.

    Returns:
        str: Formatted account stage with Rich formatting tags.
    """
    if not stage:
        return stage

    stage_lower = stage.lower()

    # Define stage mappings with fuzzy matching keywords
    stage_formats = {
        "closed": {"style": "[red]", "keywords": ["closed", "lost", "closed lost"]},
        "churned": {"style": "[hot_pink]", "keywords": ["churn", "churned customer", "former"]},
        "customer": {"style": "[bold cornflower_blue]", "keywords": ["customer", "client"]},
        "prospect": {"style": "[bold green]", "keywords": ["prospect", "lead"]},
        "pipeline": {"style": "[aquamarine1]", "keywords": ["pipeline", "in progress"]},
        "competitor": {"style": "[grey30]", "keywords": ["competitor", "lost to competitor"]},
    }

    # Find matching stage format
    for format_info in stage_formats.values():
        if any(keyword in stage_lower for keyword in format_info["keywords"]):
            return f"{format_info['style']}{stage}[/]"

    return stage


class HealthcheckPrinter:
    """Prints formatted healthcheck results for CloudZero connections."""

    def __init__(self, connection: dict[str, Any]):
        self.connection = connection
        self.healthcheck_result = connection.get("healthcheck_result") or {}

    def validate_permissions(self):
        """Validate basic permissions before proceeding."""
        if self.healthcheck_result.get("permission_errors") and self.healthcheck_result["permission_errors"].get(
            "cz:healthcheck_accounts"
        ):
            console.print(
                Panel(
                    f"[red]ERROR: {self.healthcheck_result.get('permission_errors', {})}\n\n"
                    "Account is likely disconnected, verify account connection with customer and try again",
                    title="Permission Error",
                    border_style="red",
                )
            )
            return False
        return True

    def print_basic_info(self):
        """Print basic connection information."""
        conn_details = self.connection.get("connection_details") or {}
        primary_role = conn_details.get("role_arn")
        account_type = "Cost Account" if "COST" in self.connection.get("connection_types", []) else "Resource Account"

        info_table = Table(box=box.ROUNDED, show_header=False)
        info_table.add_column("Property", style="cyan")
        info_table.add_column("Value", style="green")

        info_rows = [
            ("Connection Date", self.connection.get("created", "unknown")),
            ("Provisioning method", self.connection.get("provisioning_process")),
            ("Last update", self.connection["last_seen"]),
            ("Account Type", account_type),
            ("Primary Role", primary_role),
            ("Primary Region", self.connection.get("connection_region")),
            (
                "Active Regions",
                self.connection.get("active_regions") or "UNKNOWN (this may fix itself, but is problematic)",
            ),
        ]

        for row in info_rows:
            info_table.add_row(*row)

        console.print(Panel(info_table, title="Basic Information", border_style="blue"))

    def print_spend_info(self):
        """Print spending information."""
        last_30_days_total_billed_spend = round(sum(self.healthcheck_result.get("30_day_spend", {}).values()), 2) or 0
        account_count = len(self.healthcheck_result.get("30_day_spend", {}).values()) or "UNKNOWN"

        spend_table = Table(box=box.ROUNDED, show_header=False)
        spend_table.add_column("Metric", style="cyan")
        spend_table.add_column("Value", style="green")

        spend_table.add_row("Total 30 day spend", f"${last_30_days_total_billed_spend:,.2f} (AWS ONLY)")
        spend_table.add_row("Total Accounts", str(account_count))

        console.print(Panel(spend_table, title="Spend Information", border_style="blue"))

    def print_health_status(self):
        """Print connection health status."""
        health_table = Table(
            "Type", "Connected", "Healthy", title="Connection Health Status", box=box.ROUNDED, title_style="bold cyan"
        )

        for link, link_data in self.healthcheck_result["link_health"].items():
            if link in LINKS_TO_IGNORE:
                continue

            connected = "🟢 Yes" if link_data else "🔴 No"
            healthy = "🟢 Yes" if link_data and link_data.get("is_healthy") else "🔴 No"

            health_table.add_row(link, connected, healthy)

        console.print(health_table)

    def print_payer_account_status(self):
        """Print AWS payer account status if applicable."""
        if "COST" not in self.connection.get("connection_types", []):
            return

        cur_stats = self.healthcheck_result.get("cur_stats") or self.connection.get("cur_metadata", {}).get("cur_stats")

        if not cur_stats:
            console.print(
                "[yellow]CUR Stats are UN-ACCESSIBLE (Permissions are old, bad or were intentionally limited)"
            )
            return

        # CUR Location information
        cur_location = Panel(
            f"s3://{self.connection['master_payer_billing_bucket_name']}/{self.connection['master_payer_billing_bucket_path']}",
            title="CUR Location",
            border_style="cyan",
        )
        console.print(cur_location)

        # CUR Statistics
        stats_table = Table(box=box.ROUNDED, show_header=False)
        stats_table.add_column("Metric", style="cyan")
        stats_table.add_column("Value", style="green")

        stats_rows = [
            ("Avg this month", f"Every {cur_stats['avg_cur_drop_time_current']} hours"),
            ("Avg last month", f"Every {cur_stats['avg_cur_drop_time_previous']} hours"),
            ("Last CUR Drop", cur_stats["cur_drop_datetime"]),
            ("Last Manifest Update", cur_stats["manifest_datetime"]),
            ("Last CZ Ingest", cur_stats["cur_ingest_datetime"]),
            (
                "CUR File Count",
                f"{cur_stats['cur_report_count']} ({self._get_complexity_label(cur_stats['cur_report_count'])})",
            ),
        ]

        for row in stats_rows:
            stats_table.add_row(*row)

        console.print(Panel(stats_table, title="CUR Statistics", border_style="blue"))

        self._print_cur_reports()

    def _get_complexity_label(self, count: int) -> str:
        """Get complexity label based on CUR report count."""
        if count <= 20:
            return "[green]easy[/]"
        elif count <= 499:
            return "[yellow]Breaking a sweat[/]"
        elif count <= 1000:
            return "[orange1]Getting Spicy[/]"
        return "[red]Yikes[/]"

    def _print_cur_reports(self):
        """Print CUR reports information."""
        reports_table = Table(
            "Name",
            "ECS Data",
            "Compatible",
            "Reachable",
            "Last CUR Drop",
            "Location",
            title="AWS Cost and Usage Reports",
            box=box.ROUNDED,
        )

        for report_name, report in self.connection["cur_metadata"].items():
            is_connected = report_name == self.connection.get("master_payer_billing_report_name")
            name = f"{report_name} (CONNECTED)" if is_connected else report_name

            ecs_enabled = "🟢" if "SPLIT_COST_ALLOCATION_DATA" in report.get("AdditionalSchemaElements", []) else "🔴"
            compatible = self._get_compatibility_status(report_name, report)
            reachable = (
                "🟢"
                if self.healthcheck_result.get("cur_health", {}).get(report_name, {}).get("bucket_healthy")
                else "🔴"
            )

            last_update = self._format_datetime(
                self.healthcheck_result.get("cur_stats", {}).get("cur_drop_datetime", "")
            )
            location = build_cur_s3_path(report.get("S3Bucket", ""), report.get("S3Prefix"), report.get("ReportName"))

            reports_table.add_row(
                Text(name, style="green" if is_connected else "white"),
                ecs_enabled,
                compatible,
                reachable,
                last_update,
                location,
            )

        console.print(reports_table)

    def _get_compatibility_status(self, report_name: str, report: dict[str, Any]) -> str:
        """
        Determine compatibility status for a report.

        Uses backend health check results exclusively - backend owns the definition
        of what makes a CUR viable through DynamoDB-stored validation rules.
        """
        cur_health = self.healthcheck_result.get("cur_health", {})
        cur_report_health = cur_health.get(report_name, {})
        is_healthy = cur_report_health.get("healthy", False)
        has_warnings = cur_report_health.get("has_warnings", False)

        # Backend determines compatibility through validation rules
        if is_healthy and not has_warnings:
            return "🟢"
        elif is_healthy and has_warnings:
            return "🟡"
        return "🔴"

    def _format_datetime(self, dt_str: str) -> str:
        """Format datetime string."""
        try:
            from datetime import datetime

            dt = datetime.fromisoformat(dt_str.replace("Z", "+00:00"))
            return dt.strftime("%Y-%m-%d %H:%M:%S")
        except (ValueError, AttributeError):
            return dt_str

    def print_all(self, show_details: bool = False):
        """Print all healthcheck information."""
        if not self.validate_permissions():
            return

        self.print_basic_info()
        self.print_spend_info()

        if self.connection.get("permission_errors"):
            console.print(
                Panel(str(self.connection["permission_errors"]), title="Permission Warnings", border_style="yellow")
            )

        if self.connection.get("internal_connection_notes"):
            console.print(
                Panel(
                    self.connection["internal_connection_notes"], title="Internal Connection Notes", border_style="red"
                )
            )

        self.print_health_status()

        if "COST" in self.connection.get("connection_types", []):
            self.print_payer_account_status()

        if show_details:
            console.print(
                Panel(
                    json.dumps(self.healthcheck_result.get("role_metadata"), indent=2),
                    title="IAM Roles",
                    border_style="blue",
                )
            )


def pprint_healthcheck_results(connection: dict[str, Any], show_details: bool = False):
    """Print formatted healthcheck results."""
    printer = HealthcheckPrinter(connection)
    printer.print_all(show_details)


def pprint_connection_list(organization_data, connection_data, arguments):
    """Print formatted list of CloudZero connections."""
    metrics = Counter()
    cloud_providers_list = set()
    console = Console(width=150, highlight=False)

    for cz_org_id, connections_dict in connection_data.items():
        connections = connections_dict.values()
        org_name = organization_data.get(cz_org_id, {}).get("organization_name", "").capitalize() or cz_org_id
        org_enabled = organization_data.get(cz_org_id, {}).get("organization_enabled", True)
        org_stage = organization_data.get(cz_org_id, {}).get("organization_stage", "unknown")
        metrics[f"stage-{org_stage}"] += 1
        if not org_enabled:
            metrics["disabled_orgs"] += 1
        else:
            metrics["disabled_orgs"] = 0
            metrics["enabled_orgs"] += 1

        org_title = f" {org_name:25} | {cz_org_id} | {org_stage}"

        title_grid = Table(
            Column(width=25, justify="left", style="underline bright_blue"),
            Column(width=36, justify="left", style="underline"),
            Column(justify="left", style="underline"),
            width=150,
            show_header=False,
            box=None,
            padding=0,
            expand=True,
            show_lines=True,
        )
        title_grid.add_row(org_name, cz_org_id, org_stage)
        print(title_grid)

        connection_table = Table(
            Column("Connected", width=17),
            Column("Type", width=12),
            Column("Properties", width=18),
            Column("Method", justify="left", width=20),
            Column("Provider", justify="left", width=15),
            Column("Account ID", justify="left", width=40),
            Column("Account Name", width=60),
            title=org_title,
            title_justify="left",
            title_style="bold green",
            show_header=True,
            box=box.ROUNDED,
            header_style="bright_white",
        )

        for connection in connections:
            metrics[connection.get("provider")] += 1
            # add cloud provider to account type list
            cloud_providers_list.add(connection.get("provider", "unknown"))
            c_type, c_props = get_connection_status(connection)
            formatted_c_type = "".join([x["display"] for x in c_type])
            formatted_c_props = ",".join([x["display"] for x in c_props])

            if connection.get("is_connected"):
                account_status = (
                    "[bright_yellow]Sorta[/]" if connection.get("permission_errors") else "[bright_green]Yes[/]"
                )
            else:
                account_status = "[bright_red]No[/]"
            provisioning_process = connection.get("provisioning_process") or "legacy"
            csp_name = connection.get("provider")
            connection_id = connection.get("connection_id")
            connection_name = connection.get("connection_name", NO_ACCOUNT_NAME)
            connection_table.add_row(
                *[
                    account_status,
                    formatted_c_type,
                    formatted_c_props,
                    provisioning_process,
                    csp_name,
                    connection_id,
                    connection_name,
                ]
            )
            if connection.get("healthcheck_result"):
                pprint_healthcheck_results(connection, arguments.get("--detailed"))
        if not connections:
            connection_table.add_row("No connections found")
        print(connection_table)
        print("")

    # Create a lookup of code to description for connection types
    account_key = {}
    for _k, v in {**CONNECTION_PROPERTIES, **CONNECTION_TYPES}.items():
        account_key[f"[cyan]{v['code']}[/cyan]"] = v["description"]

    account_types = ", ".join(cloud_providers_list)
    print(
        f"Analyzed {metrics['enabled_orgs']} enabled organization(s) with {metrics.get('total_connections', 0)} connections ({metrics.get('malformed', 0)} malformed), "
        f"across {len(cloud_providers_list)} types ({account_types})."
    )

    # Display key for connection codes
    key_text = ", ".join([f"{k}: {v}" for k, v in account_key.items()])
    console.print(f"Key: {key_text}", width=150, overflow="fold")

    if metrics["enabled_orgs"] > 1:
        # print a table of all recorded org_stages
        org_stage_table = Table(
            title="\nOrganization Stage Analysis",
            title_justify="left",
            title_style="bold green",
            show_header=True,
            box=box.ROUNDED,
            header_style="bright_white",
            width=50,
        )
        org_stage_table.add_column("Stage", justify="left")
        org_stage_table.add_column("Count", justify="left")
        for stage, count in metrics.items():
            if stage.startswith("stage-"):
                clean_stage = stage.replace("stage-", "")
                org_stage_table.add_row(clean_stage, str(count))
        console.print(org_stage_table)


def get_connection_status(account):
    """
    Get connection status for a CloudZero account.

    Args:
        account: The account connection data

    Returns:
        tuple: (connection_type_status, connection_prop_status)
    """
    connection_type_status = []
    connection_prop_status = []
    try:
        if account.get("malformed"):
            connection_type_status.append(
                {
                    "code": "BAD",
                    "name": "BAD",
                    "healthy": False,
                    "display": ":warning:[bright_red]BAD[/bright_red]",
                    "description": "Malformed connection record",
                }
            )
        else:
            for _k, v in CONNECTION_TYPES.items():
                if v["type_key"] in account.get("connection_types", []):
                    if account["cloud_provider"] == "AWS":
                        healthy = get_in(["link_health", v["health_key"], "is_healthy"], account)
                    else:
                        healthy = account.get("is_healthy")
                    connection_type_status.append(
                        {
                            "code": v["code"],
                            "name": v["name"],
                            "healthy": healthy,
                            "display": healthy and f"[bright_green]{v['code']}[/]" or f"[bright_red]{v['code']}[/]",
                            "description": v["description"],
                        }
                    )

            for _k, v in CONNECTION_PROPERTIES.items():
                if account.get(v["enabled_key"]):
                    connection_prop_status.append(
                        {
                            "code": v["code"],
                            "name": v["name"],
                            "healthy": True,
                            "display": f"[deep_sky_blue1]{v['code']}[/]",
                            "description": v["description"],
                        }
                    )
                elif v.get("provisioning_keywords"):
                    if v.get("provisioning_keywords", [])[0] == "!":
                        if account.get("provisioning_process") not in v["provisioning_keywords"][1:]:
                            connection_prop_status.append(
                                {
                                    "code": v["code"],
                                    "name": v["name"],
                                    "healthy": True,
                                    "display": f"[deep_sky_blue1]{v['code']}[/]",
                                    "description": v["description"],
                                }
                            )

    except Exception as error:
        print(
            f"Error processing status for '{account.get('connection_name')} ({account.get('connection_id')})' (error:{error})"
        )
        return [[{"display": "[bright_red]???????[/]", "code": "ERR", "name": "ERROR"}], []]

    return connection_type_status, connection_prop_status
