# Copyright (c) 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 csv
import json
import sys
from typing import Any

from rich.console import Console
from rich.pretty import pprint
from rich.progress import Progress, SpinnerColumn, TextColumn
from rich.table import Table, box

from src.api.v4_connections import get_connection, get_connections_for_org
from src.api.v4_healthcheck import get_connection_healthcheck
from src.api.v4_organizations import get_organizations
from src.constants import HEALTHCHECK_SUPPORTED_PROVIDERS
from src.utils.console import eprint

console = Console(width=None)
CONSOLE_WIDTH = console.width or 150
TABLE_WIDTH = CONSOLE_WIDTH - 4


# ------------------------------------------------------------------------------------------------
# DATA MODELS AND PROCESSING
# ------------------------------------------------------------------------------------------------


class HealthcheckData:
    """Class for processing and enriching healthcheck data."""

    @staticmethod
    def fetch_healthcheck(connection_id: str) -> dict[str, Any]:
        """Fetch healthcheck data for a connection using v4 API single connection endpoint."""
        try:
            response = get_connection_healthcheck(connection_id, include_details=True, debug=False)
            return response.model_dump()
        except Exception as e:
            eprint(f"[red]Error fetching healthcheck for {connection_id}: {e}[/]")
            return {}

    @staticmethod
    def enrich_connection(connection: dict[str, Any]) -> dict[str, Any]:
        """Enrich connection data with healthcheck information."""
        if not connection:
            return {}

        connection_id = connection.get("connection_id")
        if not connection_id:
            return connection

        healthcheck_result = HealthcheckData.fetch_healthcheck(connection_id)
        connection["healthcheck_result"] = healthcheck_result
        return connection

    @staticmethod
    def process_connection_batch(
        connections: dict[str, dict[str, Any]],
        progress_callback=None,
        progress_task=None,
        progress_obj=None,
    ) -> dict[str, dict[str, Any]]:
        """Process connections using single-shot V4 API calls."""
        results = {}

        for org_id, org_connections in connections.items():
            results[org_id] = {}
            for conn_id, connection in org_connections.items():
                connection_id = connection.get("connection_id")
                if connection_id:
                    connection_name = connection.get("connection_name", conn_id)
                    if progress_obj and progress_task is not None:
                        progress_obj.update(progress_task, description=f"Checking {connection_name}...")

                    try:
                        is_debug = connection.get("__debug", False)
                        if is_debug:
                            eprint(f"[cyan]DEBUG: Using single connection healthcheck API for {connection_id}[/]")
                        healthcheck_response = get_connection_healthcheck(
                            connection_id, include_details=True, debug=is_debug
                        )

                        healthcheck_dict = healthcheck_response.model_dump()
                        connection["healthcheck_result"] = healthcheck_dict

                        connection["is_healthy"] = healthcheck_dict.get("status") == "healthy"
                        connection["health_status"] = healthcheck_dict.get("status", "unknown")
                        connection["health_score"] = healthcheck_dict.get("overall_score")

                    except Exception as e:
                        eprint(f"[yellow]Healthcheck failed for {connection_id}: {e}[/]")
                        connection["healthcheck_result"] = {"error": str(e), "status": "unknown"}

                results[org_id][conn_id] = connection

                if progress_callback:
                    progress_callback()

        return results


# ------------------------------------------------------------------------------------------------
# HEALTHCHECK MANAGER
# ------------------------------------------------------------------------------------------------


class HealthcheckManager:
    """Manages the overall healthcheck process."""

    @staticmethod
    def fetch_connections(
        arguments: dict[str, Any], organization_data: dict[str, Any] = None
    ) -> dict[str, dict[str, dict[str, Any]]]:
        """Fetch connections based on arguments and organization data."""
        connections_to_check = {}

        from src.utils.ui import PanelFactory

        console.print(
            PanelFactory.create_header_panel(
                "[bold]CloudZero Connection Healthcheck[/bold]\nChecking connection health, permissions, and configuration",
                title="Connection Healthcheck",
                width=CONSOLE_WIDTH,
                border_style="green",
            )
        )

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

        # Case 1: Specific connection ID provided
        if connection_id := arguments.get("<ID-or-NAME>"):
            try:
                connection_v4 = get_connection(connection_id, debug=is_debug)
                if connection_v4:
                    connection_response = {
                        "cz_organization_ids": [connection_v4.cz_organization_id],
                        "data": {connection_v4.cz_organization_id: {connection_id: connection_v4.model_dump()}},
                    }
                else:
                    connection_response = {}
            except Exception:
                connection_response = {}

            # If no connection found, try to interpret input as an organization ID or name
            if not connection_response.get("cz_organization_ids"):
                if is_debug:
                    eprint(f"Debug - Connection response: {connection_response}")
                    eprint(f"Debug - No connection found with ID: {connection_id}. Trying as organization ID/name.")

                try:
                    from src.utils.core import is_valid_uuid

                    if is_valid_uuid(connection_id):
                        org_response = get_organizations(
                            organization_query=connection_id, debug=is_debug, enabled_only=False
                        )
                        if org_response and connection_id in org_response:
                            org_data = {connection_id: org_response[connection_id].model_dump()}
                        else:
                            org_data = None
                    else:
                        org_response = get_organizations(
                            organization_query=connection_id, debug=is_debug, enabled_only=False
                        )
                        if org_response:
                            org_data = {}
                            for org_id, org_details in org_response.items():
                                org_data[org_id] = org_details.model_dump()
                        else:
                            org_data = None
                except Exception as e:
                    if is_debug:
                        eprint(f"[yellow]Debug: Failed to load organization data: {e}[/]")
                        eprint("[yellow]Debug: Continuing healthcheck without organization data[/]")
                    org_data = None

                if org_data:
                    org_id = list(org_data.keys())[0]
                    org_name = org_data[org_id].get("cz_organization_name") or org_data[org_id].get(
                        "organization_name", "Unknown"
                    )

                    try:
                        connections_list = get_connections_for_org(org_id, debug=is_debug)

                        org_connections_data = {}
                        for conn in connections_list:
                            types_key = ",".join(conn.connection_types) if conn.connection_types else "UNKNOWN"
                            composite_key = f"{conn.connection_id}|{types_key}"
                            org_connections_data[composite_key] = conn.model_dump()
                    except Exception as e:
                        console.print(f"[red]Error fetching connections for organization: {e}[/]")
                        if is_debug:
                            console.print_exception()
                        org_connections_data = {}

                    # Filter only AWS payer connections
                    payer_connections = {}
                    for conn_id, conn in org_connections_data.items():
                        provider = conn.get("provider", "")
                        conn_types = conn.get("connection_types", [])
                        is_cost = "COST" in conn_types

                        if provider in HEALTHCHECK_SUPPORTED_PROVIDERS and is_cost:
                            payer_connections[conn_id] = conn

                    payer_count = len(payer_connections)

                    if payer_count == 0:
                        console.print(
                            f"[yellow]Organization '{org_name}' found, but it has no AWS payer connections that support healthchecking.[/]"
                        )
                        console.print("[yellow]Only AWS payer accounts can be healthchecked at this time.[/]")
                        sys.exit(1)
                    elif payer_count == 1:
                        payer_id = list(payer_connections.keys())[0]
                        payer_conn = payer_connections[payer_id]
                        payer_name = payer_conn.get("connection_name", "Unknown")
                        console.print(
                            f"[green]Found single AWS payer connection '{payer_name}' for organization '{org_name}'.[/]"
                        )
                        connection_response = {"cz_organization_ids": [org_id], "data": {org_id: payer_connections}}
                    else:
                        console.print(
                            f"[yellow]Organization '{org_name}' has multiple payer connections. Please specify which one to check:[/]"
                        )

                        payer_table = Table(
                            title=f"Payer Connections for {org_name}",
                            box=box.ROUNDED,
                            header_style="bright_white",
                            title_style="cyan",
                        )

                        payer_table.add_column("Connection ID", style="yellow")
                        payer_table.add_column("Connection Name", style="bright_white")
                        payer_table.add_column("Provider", style="cyan")

                        for conn_id, conn in payer_connections.items():
                            conn_name = conn.get("connection_name", "Unknown")
                            provider = conn.get("provider") or conn.get("cloud_provider", "Unknown")
                            payer_table.add_row(conn_id, conn_name, provider)

                        console.print(payer_table)
                        console.print("[yellow]Run:[/] [cyan]czops connection healthcheck [yellow]<CONNECTION_ID>[/]")
                        sys.exit(1)
                else:
                    console.print(
                        f"[red]Error: Could not find connection '{connection_id}' or any organization with that name/ID[/]"
                    )
                    console.print(
                        f"[yellow]Hint: If '{connection_id}' is a valid AWS account ID, verify the connection exists and you have permissions to access it[/]"
                    )
                    sys.exit(1)

            org_id = connection_response.get("cz_organization_ids")[0]
            connections_data = connection_response.get("data", {})

            for org_connections in connections_data.values():
                for conn_id, conn_data in org_connections.items():
                    if is_debug:
                        conn_data["__debug"] = True

            connections_to_check[org_id] = connections_data.get(org_id, {})

            if not organization_data:
                try:
                    org_response = get_organizations(
                        organization_query=org_id, debug=arguments.get("--debug"), enabled_only=False
                    )
                    organization_data = (
                        {org_id: org_response[org_id].model_dump()} if org_response and org_id in org_response else {}
                    )
                except Exception as e:
                    if arguments.get("--debug"):
                        eprint(f"[yellow]Debug: Failed to load organization data: {e}[/]")
                        eprint("[yellow]Debug: Continuing with healthcheck without organization data[/]")
                    organization_data = {}

            if organization_data and org_id in organization_data:
                org_name = organization_data[org_id].get("cz_organization_name", "")
                for conn_data in connections_to_check.get(org_id, {}).values():
                    if "cz_organization_name" not in conn_data:
                        conn_data["cz_organization_name"] = org_name

        # Case 2: Use existing organization data
        elif organization_data:
            for org_id in organization_data:
                try:
                    connections_list = get_connections_for_org(org_id, debug=is_debug)
                    connections_data = {org_id: {conn["connection_id"]: conn for conn in connections_list}}
                except Exception as e:
                    console.print(f"[red]Error fetching connections for organization {org_id}: {e}[/]")
                    if is_debug:
                        console.print_exception()
                    connections_data = {}

                if org_id in connections_data:
                    for conn_id, conn_data in connections_data[org_id].items():
                        if is_debug:
                            conn_data["__debug"] = True

                connections_to_check[org_id] = connections_data.get(org_id, {})

                org_name = organization_data[org_id].get("cz_organization_name", "")
                for conn_data in connections_to_check.get(org_id, {}).values():
                    if "cz_organization_name" not in conn_data:
                        conn_data["cz_organization_name"] = org_name

        # Case 3: No data provided
        else:
            console.print("[red]Error: No organization or connection provided[/]")
            sys.exit(1)

        return connections_to_check

    @staticmethod
    def process_connections(
        connections: dict[str, dict[str, dict[str, Any]]], arguments: dict[str, Any]
    ) -> dict[str, dict[str, dict[str, Any]]]:
        """Process connections by adding healthcheck data."""
        connection_count = sum(len(connections_data) for connections_data in connections.values())
        processed_connections = {}

        with Progress(
            SpinnerColumn(),
            TextColumn("[cyan]Running healthcheck... [bright_white]{task.completed}/{task.total}[/]"),
            TextColumn("[dim]|[/]"),
            TextColumn("[yellow]{task.description}[/]"),
            console=console,
        ) as progress:
            task = progress.add_task("Checking connections...", total=connection_count)

            def update_progress():
                progress.advance(task)

            processed_connections = HealthcheckData.process_connection_batch(
                connections,
                progress_callback=update_progress,
                progress_task=task,
                progress_obj=progress,
            )

        return processed_connections

    @staticmethod
    def display_connections(connections: dict[str, dict[str, dict[str, Any]]], arguments: dict[str, Any]) -> None:
        """Display healthcheck information for connections."""
        export_file = arguments.get("-o") or arguments.get("--output")

        if export_file:
            if export_file.endswith(".json"):
                HealthcheckManager.export_json(connections, export_file)
            elif export_file.endswith(".csv"):
                HealthcheckManager.export_csv(connections, export_file)
            else:
                console.print("[red]Unsupported export format. Use .json or .csv extension.[/]")
                sys.exit(1)
            return

        if arguments.get("--raw"):
            print(json.dumps(connections, indent=2, default=str))
            return

        if arguments.get("--json"):
            pprint(connections)
            return

        from src.features.connections.healthcheck_display import display_healthcheck

        for _org_id, org_connections in connections.items():
            for _conn_id, connection in org_connections.items():
                healthcheck_result = connection.get("healthcheck_result", {})

                if healthcheck_result:
                    display_healthcheck(healthcheck_result, connection)
                else:
                    console.print("[yellow]No healthcheck data available for this connection[/]")

    @staticmethod
    def export_json(connections: dict[str, dict[str, dict[str, Any]]], filename: str) -> None:
        """Export healthcheck results to JSON file."""
        try:
            with open(filename, "w") as f:
                json.dump(connections, f, indent=2, default=str)
            console.print(f"[green]Healthcheck results exported to {filename}[/]")
        except Exception as e:
            console.print(f"[red]Error exporting to JSON: {e}[/]")
            sys.exit(1)

    @staticmethod
    def export_csv(connections: dict[str, dict[str, dict[str, Any]]], filename: str) -> None:
        """Export healthcheck results to CSV file."""
        try:
            rows = []
            for org_id, org_connections in connections.items():
                for conn_id, connection in org_connections.items():
                    healthcheck = connection.get("healthcheck_result", {})
                    metadata = healthcheck.get("metadata", {})
                    summary = healthcheck.get("summary", {})

                    row = {
                        "organization_id": metadata.get("cz_organization_id", org_id),
                        "organization_name": metadata.get("cz_organization_name")
                        or connection.get("cz_organization_name", "Unknown"),
                        "connection_id": healthcheck.get("connection_id", conn_id),
                        "connection_name": metadata.get("connection_name")
                        or connection.get("connection_name", "Unknown"),
                        "cloud_provider": metadata.get("cloud_provider", "Unknown"),
                        "connection_types": ",".join(metadata.get("connection_types", [])),
                        "status": healthcheck.get("status", "unknown"),
                        "overall_score": healthcheck.get("overall_score", ""),
                        "last_updated": metadata.get("last_updated", ""),
                    }

                    # Extract per-check status from checks array
                    checks = healthcheck.get("checks", [])
                    for check in checks:
                        if not isinstance(check, dict):
                            continue
                        check_name = check.get("name", "")
                        row[f"check_{check_name}_status"] = check.get("status", "unknown")
                        row[f"check_{check_name}_message"] = check.get("message", "")

                    # Extract summary counts
                    row["total_checks"] = summary.get("total_checks", 0)
                    row["healthy_count"] = summary.get("healthy_count", 0)
                    row["unhealthy_count"] = summary.get("unhealthy_count", 0)

                    rows.append(row)

            if not rows:
                console.print("[yellow]No data to export[/]")
                return

            # Collect all possible field names across rows
            all_fields = []
            seen = set()
            for row in rows:
                for key in row:
                    if key not in seen:
                        all_fields.append(key)
                        seen.add(key)

            with open(filename, "w", newline="") as f:
                writer = csv.DictWriter(f, fieldnames=all_fields)
                writer.writeheader()
                writer.writerows(rows)

            console.print(f"[green]Healthcheck results exported to {filename}[/]")
        except Exception as e:
            console.print(f"[red]Error exporting to CSV: {e}[/]")
            sys.exit(1)


# ------------------------------------------------------------------------------------------------
# COMMAND FUNCTION
# ------------------------------------------------------------------------------------------------


def functions(arguments: dict[str, Any], organization_data: dict[str, Any] = None) -> None:
    """Main function for the healthcheck command."""
    global CONSOLE_WIDTH, TABLE_WIDTH
    current_width = console.width
    if current_width:
        CONSOLE_WIDTH = current_width
        TABLE_WIDTH = CONSOLE_WIDTH - 4

    if organization_data is None:
        organization_data = {}

    try:
        connections = HealthcheckManager.fetch_connections(arguments, organization_data)

        if arguments.get("--debug"):
            eprint("\nConnections structure after fetch:")
            for org_id, org_connections in connections.items():
                eprint(f"Organization ID: {org_id}")
                if isinstance(org_connections, dict):
                    eprint(f"  Connection count: {len(org_connections)}")
                    for conn_id, conn_data in org_connections.items():
                        eprint(f"  Connection ID: {conn_id}")
                        if isinstance(conn_data, dict) and "connection_id" in conn_data:
                            eprint(f"    connection_id: {conn_data['connection_id']}")
                        else:
                            eprint("    Missing connection_id field")
                else:
                    eprint(f"  Org connections not a dictionary: {type(org_connections)}")

        has_connections = False
        for org_id, org_connections in connections.items():
            if org_connections:
                has_connections = True
                break

        if not has_connections:
            from src.utils.console import print_error_image

            print_error_image(
                context=f"No valid connections found for ID: {arguments.get('<ID-or-NAME>')}\n\n"
                f"Please verify the connection ID and try again.\n"
                f"If the issue persists, please report this to #cz-tools.",
                error_type="healthcheck",
            )
            sys.exit(1)

        processed_connections = HealthcheckManager.process_connections(connections, arguments)
        HealthcheckManager.display_connections(processed_connections, arguments)

    except Exception as error:
        from src.utils.console import print_error_image

        print_error_image(
            context=f"An error occurred during healthcheck: {str(error)}\n\n"
            f"Try running your command again with --debug for more information.\n"
            f"If the issue persists, please report this to #cz-tools.",
            error_type="healthcheck",
        )
        if arguments.get("--debug"):
            console.print_exception()
        sys.exit(1)
