import math
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Sequence

from pydantic import BaseModel
from pydantic import Field

from service_capacity_modeling.interface import AccessConsistency
from service_capacity_modeling.interface import AccessPattern
from service_capacity_modeling.interface import CapacityDesires
from service_capacity_modeling.interface import CapacityPlan
from service_capacity_modeling.interface import CapacityRegretParameters
from service_capacity_modeling.interface import CapacityRequirement
from service_capacity_modeling.interface import certain_float
from service_capacity_modeling.interface import certain_int
from service_capacity_modeling.interface import ClusterCapacity
from service_capacity_modeling.interface import Clusters
from service_capacity_modeling.interface import Consistency
from service_capacity_modeling.interface import DataShape
from service_capacity_modeling.interface import Drive
from service_capacity_modeling.interface import FixedInterval
from service_capacity_modeling.interface import GlobalConsistency
from service_capacity_modeling.interface import Instance
from service_capacity_modeling.interface import Interval
from service_capacity_modeling.interface import QueryPattern
from service_capacity_modeling.interface import RegionClusterCapacity
from service_capacity_modeling.interface import RegionContext
from service_capacity_modeling.interface import Requirements
from service_capacity_modeling.interface import ServiceCapacity
from service_capacity_modeling.models import CapacityModel
from service_capacity_modeling.models import CostAwareModel
from service_capacity_modeling.models.common import cluster_infra_cost
from service_capacity_modeling.models.common import compute_stateless_region
from service_capacity_modeling.models.common import network_services
from service_capacity_modeling.models.common import normalize_cores
from service_capacity_modeling.models.common import simple_network_mbps
from service_capacity_modeling.models.common import sqrt_staffed_cores


def _estimate_java_app_requirement(
    instance: Instance,
    desires: CapacityDesires,
    failover: bool = True,
    jvm_memory_overhead: float = 1.2,
) -> CapacityRequirement:
    needed_cores = sqrt_staffed_cores(desires)
    needed_network_mbps = simple_network_mbps(desires)

    if failover:
        # For failover provision at 40% utilization
        needed_cores = int(math.ceil(needed_cores * (1 / 0.4)))
        needed_network_mbps = int(math.ceil(needed_network_mbps * (1 / 0.4)))

    # Adjust cores to clock frequency differences against measurement
    needed_cores = normalize_cores(
        core_count=needed_cores,
        target_shape=instance,
        reference_shape=desires.reference_shape,
    )

    # Assume a Java application that can allocate about 1 GiB/s to heap
    # per 2 GiB of heap with some amount of overhead on the network traffic.
    # So for example if we have 512 MiB of network traffic there is some
    # overhead associated with that...

    # TODO: we should probably have the memory bandwidth attached to
    # the instance type, e.g. Intel CPUs and AMD CPUs have different
    # per core memory bandwidth.
    mem_allocation_mbps = needed_network_mbps * jvm_memory_overhead
    heap_allocation_gibps = (mem_allocation_mbps / 8) / 1024
    network_heap = heap_allocation_gibps * 2

    needed_memory_gib = network_heap

    return CapacityRequirement(
        requirement_type="java-app",
        reference_shape=desires.reference_shape,
        cpu_cores=certain_int(needed_cores),
        mem_gib=certain_float(needed_memory_gib),
        network_mbps=certain_float(needed_network_mbps),
        context={
            "network_heap_gib": network_heap,
            "reserved_mem": desires.data_shape.reserved_instance_app_mem_gib,
        },
    )


def _estimate_java_app_region(  # pylint: disable=too-many-positional-arguments
    instance: Instance,
    drive: Drive,
    desires: CapacityDesires,
    context: RegionContext,
    root_disk_gib: int = 10,
    failover: bool = True,
    jvm_memory_overhead: float = 2,
) -> Optional[CapacityPlan]:
    if drive.name != "gp2":
        return None

    zones_per_region = context.zones_in_region
    requirement = _estimate_java_app_requirement(
        instance, desires, failover, jvm_memory_overhead
    )

    drive = drive.model_copy()
    drive.size_gib = root_disk_gib
    attached_drives = (drive,)

    cluster: RegionClusterCapacity = compute_stateless_region(
        instance=instance,
        needed_cores=int(requirement.cpu_cores.mid),
        needed_memory_gib=requirement.mem_gib.mid,
        needed_network_mbps=requirement.network_mbps.mid,
        num_zones=zones_per_region,
    )
    cluster.cluster_type = NflxJavaAppCapacityModel.cluster_type
    cluster.attached_drives = attached_drives

    # Generally don't want giant clusters
    # Especially not above 1000 because some load balancers struggle
    # with such large clusters

    if cluster.count <= 256:
        costs = NflxJavaAppCapacityModel.cluster_costs(
            service_type=NflxJavaAppCapacityModel.service_name,
            regional_clusters=[cluster],
        )
        services = NflxJavaAppCapacityModel.service_costs(
            service_type=NflxJavaAppCapacityModel.service_name,
            context=context,
            desires=desires,
            extra_model_arguments={},
        )
        for s in services:
            costs[s.service_type] = s.annual_cost

        return CapacityPlan(
            requirements=Requirements(regional=[requirement]),
            candidate_clusters=Clusters(
                annual_costs=costs,
                regional=[cluster],
                zonal=[],
                services=services,
            ),
        )
    return None


class NflxJavaAppArguments(BaseModel):
    failover: bool = Field(
        default=True, description="If this app participates in failover"
    )
    jvm_memory_overhead: float = Field(
        default=1.2,
        description="How much overhead does the heap have per read byte",
    )
    root_disk_gib: int = Field(
        default=10, description="How many GiB of root volume to attach"
    )


class NflxJavaAppCapacityModel(CapacityModel, CostAwareModel):
    service_name = "nflx-java-app"
    cluster_type = "nflx-java-app"

    @staticmethod
    def cluster_costs(
        service_type: str,
        zonal_clusters: Sequence[ClusterCapacity] = (),
        regional_clusters: Sequence[ClusterCapacity] = (),
    ) -> Dict[str, float]:
        return cluster_infra_cost(
            service_type,
            zonal_clusters,
            regional_clusters,
            cluster_type=NflxJavaAppCapacityModel.cluster_type,
        )

    @staticmethod
    def service_costs(
        service_type: str,
        context: RegionContext,
        desires: CapacityDesires,
        extra_model_arguments: Dict[str, Any],
    ) -> List[ServiceCapacity]:
        # TODO(matthewho): Currently returns empty because RegionContext is
        # created without services. Need to determine if stateless apps should
        # have cross-zone costs (copies_per_region=2 implies 1 cross-AZ hop).
        _ = (context, extra_model_arguments)
        return network_services(
            service_type,
            RegionContext(num_regions=1),
            desires,
            copies_per_region=2,
        )

    @staticmethod
    def capacity_plan(
        instance: Instance,
        drive: Drive,
        context: RegionContext,
        desires: CapacityDesires,
        extra_model_arguments: Dict[str, Any],
    ) -> Optional[CapacityPlan]:
        failover: bool = extra_model_arguments.get("failover", True)
        jvm_memory_overhead: float = extra_model_arguments.get(
            "jvm_memory_overhead", 1.2
        )
        root_disk_gib: int = extra_model_arguments.get("root_disk_gib", 10)

        return _estimate_java_app_region(
            instance=instance,
            drive=drive,
            desires=desires,
            failover=failover,
            root_disk_gib=root_disk_gib,
            jvm_memory_overhead=jvm_memory_overhead,
            context=context,
        )

    @staticmethod
    def description() -> str:
        return "Netflix Streaming Java App Model"

    @staticmethod
    def extra_model_arguments_schema() -> Dict[str, Any]:
        return NflxJavaAppArguments.model_json_schema()

    @staticmethod
    def regret(
        regret_params: CapacityRegretParameters,
        optimal_plan: CapacityPlan,
        proposed_plan: CapacityPlan,
    ) -> Dict[str, float]:
        regret = super(NflxJavaAppCapacityModel, NflxJavaAppCapacityModel).regret(
            regret_params, optimal_plan, proposed_plan
        )
        regret["disk_space"] = 0
        return regret

    @staticmethod
    def default_desires(
        user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
    ) -> CapacityDesires:
        if user_desires.query_pattern.access_pattern == AccessPattern.latency:
            return CapacityDesires(
                query_pattern=QueryPattern(
                    access_pattern=AccessPattern.latency,
                    access_consistency=GlobalConsistency(
                        same_region=Consistency(
                            target_consistency=AccessConsistency.read_your_writes,
                        ),
                        cross_region=Consistency(
                            target_consistency=AccessConsistency.never,
                        ),
                    ),
                    estimated_mean_read_size_bytes=Interval(
                        low=128, mid=1024, high=65536, confidence=0.95
                    ),
                    estimated_mean_write_size_bytes=Interval(
                        low=64, mid=128, high=1024, confidence=0.95
                    ),
                    estimated_mean_read_latency_ms=Interval(
                        low=0.2, mid=1, high=2, confidence=0.98
                    ),
                    estimated_mean_write_latency_ms=Interval(
                        low=0.2, mid=1, high=2, confidence=0.98
                    ),
                    # "Single digit milliseconds SLO"
                    read_latency_slo_ms=FixedInterval(
                        minimum_value=0.5,
                        maximum_value=10,
                        low=1,
                        mid=2,
                        high=5,
                        confidence=0.98,
                    ),
                    write_latency_slo_ms=FixedInterval(
                        low=1, mid=2, high=5, confidence=0.98
                    ),
                ),
                data_shape=DataShape(
                    # Assume 4 GiB heaps
                    reserved_instance_app_mem_gib=4
                ),
            )
        else:
            return CapacityDesires(
                query_pattern=QueryPattern(
                    access_pattern=AccessPattern.latency,
                    access_consistency=GlobalConsistency(
                        same_region=Consistency(
                            target_consistency=AccessConsistency.read_your_writes,
                        ),
                        cross_region=Consistency(
                            target_consistency=AccessConsistency.never,
                        ),
                    ),
                    estimated_mean_read_size_bytes=Interval(
                        low=128, mid=1024, high=65536, confidence=0.95
                    ),
                    estimated_mean_write_size_bytes=Interval(
                        low=64, mid=128, high=1024, confidence=0.95
                    ),
                    # Throughput ops can be slower
                    estimated_mean_read_latency_ms=Interval(
                        low=0.2, mid=4, high=8, confidence=0.98
                    ),
                    estimated_mean_write_latency_ms=Interval(
                        low=0.2, mid=1, high=5, confidence=0.98
                    ),
                    # "Tens of millisecond SLO"
                    read_latency_slo_ms=FixedInterval(
                        minimum_value=0.5,
                        maximum_value=100,
                        low=1,
                        mid=5,
                        high=40,
                        confidence=0.98,
                    ),
                    write_latency_slo_ms=FixedInterval(
                        minimum_value=0.5,
                        maximum_value=100,
                        low=1,
                        mid=5,
                        high=40,
                        confidence=0.98,
                    ),
                ),
                data_shape=DataShape(
                    # Assume 4 GiB heaps
                    reserved_instance_app_mem_gib=4
                ),
            )


nflx_java_app_capacity_model = NflxJavaAppCapacityModel()
