# discovery/consul_registry.py
import os
import socket
from contextlib import asynccontextmanager
from typing import Optional

import consul.aio
from fastapi import FastAPI
from pydantic_settings import BaseSettings


class ConsulSettings(BaseSettings):
    consul_host: str = "localhost"
    consul_port: int = 8500
    service_name: str = "order-service"
    service_port: int = 8000
    check_interval: str = "10s"
    deregister_after: str = "1m"

    class Config:
        env_file = ".env"


class ConsulRegistry:
    def __init__(self, settings: ConsulSettings) -> None:
        self.settings = settings
        self.consul = consul.aio.Consul(
            host=settings.consul_host,
            port=settings.consul_port,
        )
        hostname = os.environ.get("HOSTNAME", socket.gethostname())
        self.service_id = f"{settings.service_name}-{hostname}"

    def _get_service_address(self) -> str:
        if os.environ.get("DOCKER"):
            return os.environ.get("HOSTNAME", socket.gethostname())
        try:
            # Get a non-loopback IPv4 address
            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
                sock.connect(("8.8.8.8", 80))
                return sock.getsockname()[0]
        except OSError:
            return "localhost"

    async def register(self) -> None:
        address = self._get_service_address()
        await self.consul.agent.service.register(
            name=self.settings.service_name,
            service_id=self.service_id,
            address=address,
            port=self.settings.service_port,
            tags=["python", "api"],
            check=consul.Check.http(
                url=f"http://{address}:{self.settings.service_port}/health",
                interval=self.settings.check_interval,
                deregister=self.settings.deregister_after,
            ),
        )
        print(f"Service registered: {self.service_id}")

    async def deregister(self) -> None:
        try:
            await self.consul.agent.service.deregister(self.service_id)
            print(f"Service deregistered: {self.service_id}")
        except Exception as exc:  # noqa: BLE001
            print(f"Failed to deregister service: {exc}")


# FastAPI lifespan handles startup/shutdown cleanly
registry: Optional[ConsulRegistry] = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global registry
    settings = ConsulSettings()
    registry = ConsulRegistry(settings)
    await registry.register()
    try:
        yield
    finally:
        # Deregister BEFORE the process exits, avoiding stale entries
        await registry.deregister()


app = FastAPI(lifespan=lifespan)


@app.get("/health")
async def health() -> dict[str, str]:
    return {"status": "ok"}

# discovery/consul_discovery.py
import asyncio
import time
from typing import Awaitable, Callable, Optional

import consul.aio
from pydantic import BaseModel


class ServiceInstance(BaseModel):
    id: str
    address: str
    port: int
    tags: list[str] = []
    status: str = "healthy"


class CacheEntry(BaseModel):
    instances: list[ServiceInstance]
    timestamp: float


class ConsulDiscovery:
    def __init__(self, client: consul.aio.Consul, cache_ttl: float = 5.0) -> None:
        self.consul = client
        self.cache: dict[str, CacheEntry] = {}
        self.cache_ttl = cache_ttl
        self._watch_tasks: dict[str, asyncio.Task] = {}

    async def discover(
        self,
        service_name: str,
        healthy: bool = True,
        use_cache: bool = True,
    ) -> list[ServiceInstance]:
        # Serve from cache when fresh
        if use_cache and service_name in self.cache:
            entry = self.cache[service_name]
            if time.monotonic() - entry.timestamp < self.cache_ttl:
                return entry.instances

        try:
            _, services = await self.consul.health.service(
                service=service_name,
                passing=healthy,
            )
            instances = [
                ServiceInstance(
                    id=entry["Service"]["ID"],
                    address=entry["Service"]["Address"],
                    port=entry["Service"]["Port"],
                    tags=entry["Service"].get("Tags", []),
                    status="healthy"
                    if all(c["Status"] == "passing" for c in entry["Checks"])
                    else "unhealthy",
                )
                for entry in services
            ]
            self.cache[service_name] = CacheEntry(
                instances=instances,
                timestamp=time.monotonic(),
            )
            return instances
        except Exception as exc:  # noqa: BLE001
            print(f"Failed to discover {service_name}: {exc}")
            # Degrade gracefully: return stale cache if available
            if service_name in self.cache:
                return self.cache[service_name].instances
            return []

    async def watch(
        self,
        service_name: str,
        callback: Callable[[list[ServiceInstance]], Awaitable[None]],
    ) -> None:
        if service_name in self._watch_tasks:
            return

        async def _watch_loop() -> None:
            index: Optional[int] = None
            while True:
                try:
                    index, services = await self.consul.health.service(
                        service=service_name,
                        passing=True,
                        index=index,
                        wait="30s",  # long-poll blocking query
                    )
                    instances = [
                        ServiceInstance(
                            id=entry["Service"]["ID"],
                            address=entry["Service"]["Address"],
                            port=entry["Service"]["Port"],
                        )
                        for entry in services
                    ]
                    self.cache[service_name] = CacheEntry(
                        instances=instances,
                        timestamp=time.monotonic(),
                    )
                    await callback(instances)
                except Exception as exc:  # noqa: BLE001
                    print(f"Watch error for {service_name}: {exc}")
                    await asyncio.sleep(5)

        self._watch_tasks[service_name] = asyncio.create_task(_watch_loop())

    async def stop_watch(self, service_name: str) -> None:
        task = self._watch_tasks.pop(service_name, None)
        if task:
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass

# discovery/service_client.py
import asyncio
import random
import time
from dataclasses import dataclass, field
from typing import Any, Literal

import httpx

from .consul_discovery import ConsulDiscovery, ServiceInstance

Strategy = Literal["round-robin", "random", "least-connections"]


@dataclass
class CircuitBreaker:
    state: str = "CLOSED"   # CLOSED, OPEN, HALF_OPEN
    failures: int = 0
    opened_at: float = 0.0


@dataclass
class ServiceClient:
    discovery: ConsulDiscovery
    service_name: str
    strategy: Strategy = "round-robin"
    timeout: float = 5.0
    instances: list[ServiceInstance] = field(default_factory=list)
    _index: int = 0
    _breakers: dict[str, CircuitBreaker] = field(default_factory=dict)
    _http: httpx.AsyncClient = field(default_factory=lambda: httpx.AsyncClient(timeout=5.0))

    async def start(self) -> None:
        self.instances = await self.discovery.discover(self.service_name)
        await self.discovery.watch(self.service_name, self._on_instances_changed)

    async def _on_instances_changed(self, instances: list[ServiceInstance]) -> None:
        print(f"Service {self.service_name} updated: {len(instances)} instances")
        self.instances = instances

    def _is_circuit_open(self, instance_id: str) -> bool:
        cb = self._breakers.get(instance_id)
        if not cb:
            return False
        if cb.state == "OPEN" and time.monotonic() - cb.opened_at > 30:
            cb.state = "HALF_OPEN"
        return cb.state == "OPEN"

    def _record_success(self, instance_id: str) -> None:
        cb = self._breakers.setdefault(instance_id, CircuitBreaker())
        cb.failures = 0
        cb.state = "CLOSED"

    def _record_failure(self, instance_id: str) -> None:
        cb = self._breakers.setdefault(instance_id, CircuitBreaker())
        cb.failures += 1
        if cb.failures >= 3:
            cb.state = "OPEN"
            cb.opened_at = time.monotonic()

    def _select(self) -> ServiceInstance:
        candidates = [
            inst for inst in self.instances if not self._is_circuit_open(inst.id)
        ]
        if not candidates:
            raise RuntimeError(f"No healthy instances of {self.service_name}")

        if self.strategy == "round-robin":
            inst = candidates[self._index % len(candidates)]
            self._index += 1
            return inst
        if self.strategy == "random":
            return random.choice(candidates)
        return candidates[0]

    async def request(
        self,
        path: str,
        method: str = "GET",
        **kwargs: Any,
    ) -> httpx.Response:
        instance = self._select()
        url = f"http://{instance.address}:{instance.port}{path}"
        try:
            response = await self._http.request(method, url, **kwargs)
            self._record_success(instance.id)
            return response
        except Exception:
            self._record_failure(instance.id)
            raise


# Usage inside a FastAPI service
async def main() -> None:
    import consul.aio
    client = consul.aio.Consul(host="consul", port=8500)
    discovery = ConsulDiscovery(client)
    payment = ServiceClient(discovery, "payment-service", strategy="round-robin")
    await payment.start()

    response = await payment.request(
        "/api/payments",
        method="POST",
        json={"amount": 100},
    )
    print(response.status_code, response.json())


if __name__ == "__main__":
    asyncio.run(main())

# discovery/dns_discovery.py
import time
from dataclasses import dataclass, field

import dns.asyncresolver


@dataclass
class DnsInstance:
    address: str
    port: int
    priority: int = 0
    weight: int = 0


@dataclass
class DNSDiscovery:
    suffix: str = ".service.consul"
    dns_server: str = "127.0.0.1"
    dns_port: int = 8600
    cache_ttl: float = 10.0
    default_port: int = 8000
    _cache: dict[str, tuple[float, list[DnsInstance]]] = field(default_factory=dict)

    def __post_init__(self) -> None:
        self._resolver = dns.asyncresolver.Resolver(configure=False)
        self._resolver.nameservers = [self.dns_server]
        self._resolver.port = self.dns_port

    async def discover(self, service_name: str) -> list[DnsInstance]:
        hostname = f"{service_name}{self.suffix}"

        cached = self._cache.get(hostname)
        if cached and time.monotonic() - cached[0] < self.cache_ttl:
            return cached[1]

        try:
            answer = await self._resolver.resolve(hostname, "A")
            instances = [
                DnsInstance(address=rdata.address, port=self.default_port)
                for rdata in answer
            ]
            self._cache[hostname] = (time.monotonic(), instances)
            return instances
        except dns.resolver.NXDOMAIN:
            print(f"Service not found: {service_name}")
            return []

    async def discover_with_port(self, service_name: str) -> list[DnsInstance]:
        # SRV records return address + port in a single lookup
        hostname = f"_{service_name}._tcp{self.suffix}"
        try:
            answer = await self._resolver.resolve(hostname, "SRV")
            return [
                DnsInstance(
                    address=str(rdata.target).rstrip("."),
                    port=rdata.port,
                    priority=rdata.priority,
                    weight=rdata.weight,
                )
                for rdata in answer
            ]
        except Exception as exc:  # noqa: BLE001
            print(f"SRV lookup failed: {exc}")
            return await self.discover(service_name)


# Usage
async def main() -> None:
    discovery = DNSDiscovery(dns_server="consul", dns_port=8600)
    payment_instances = await discovery.discover("payment")
    print(payment_instances)

# In Kubernetes, just use the service name!
import os
import httpx
from pydantic_settings import BaseSettings


class Settings(BaseSettings):
    payment_url: str = "http://payment"
    # Full DNS name form also works:
    # payment_url: str = "http://payment.default.svc.cluster.local"


settings = Settings()

async def charge(amount: int) -> dict:
    async with httpx.AsyncClient(timeout=5.0) as client:
        response = await client.post(
            f"{settings.payment_url}/api/payments",
            json={"amount": amount},
        )
        response.raise_for_status()
        return response.json()


# Cross-namespace calls use fully qualified DNS
ORDER_URL = "http://order.other-namespace.svc.cluster.local"

# discovery/k8s_discovery.py
from dataclasses import dataclass

import dns.asyncresolver


@dataclass
class PodEndpoint:
    address: str
    port: int


class K8sDiscovery:
    def __init__(self, port: int = 8000) -> None:
        self._resolver = dns.asyncresolver.Resolver()
        self.default_port = port

    async def discover_pods(
        self,
        service_name: str,
        namespace: str = "default",
    ) -> list[PodEndpoint]:
        # Headless service DNS returns one A record per pod
        hostname = f"{service_name}-headless.{namespace}.svc.cluster.local"
        try:
            answer = await self._resolver.resolve(hostname, "A")
            return [
                PodEndpoint(address=rdata.address, port=self.default_port)
                for rdata in answer
            ]
        except Exception as exc:  # noqa: BLE001
            print(f"Failed to discover {service_name}: {exc}")
            return []


# Usage: get raw pod IPs for custom load balancing
async def main() -> None:
    discovery = K8sDiscovery(port=8000)
    pods = await discovery.discover_pods("payment")
    print(pods)

# discovery/k8s_client.py
import asyncio
from dataclasses import dataclass
from typing import Awaitable, Callable

from kubernetes_asyncio import client, config, watch


@dataclass
class Endpoint:
    ip: str
    port: int
    node_name: str | None
    ready: bool


class K8sServiceDiscovery:
    def __init__(self) -> None:
        self._core_api: client.CoreV1Api | None = None

    async def init(self, in_cluster: bool = True) -> None:
        if in_cluster:
            config.load_incluster_config()
        else:
            await config.load_kube_config()
        self._core_api = client.CoreV1Api()

    async def get_endpoints(
        self,
        service_name: str,
        namespace: str = "default",
    ) -> list[Endpoint]:
        assert self._core_api is not None
        try:
            ep = await self._core_api.read_namespaced_endpoints(
                service_name, namespace
            )
            results: list[Endpoint] = []
            for subset in ep.subsets or []:
                ports = subset.ports or []
                for address in subset.addresses or []:
                    for port in ports:
                        results.append(
                            Endpoint(
                                ip=address.ip,
                                port=port.port,
                                node_name=address.node_name,
                                ready=True,
                            )
                        )
                for address in subset.not_ready_addresses or []:
                    for port in ports:
                        results.append(
                            Endpoint(
                                ip=address.ip,
                                port=port.port,
                                node_name=address.node_name,
                                ready=False,
                            )
                        )
            return results
        except Exception as exc:  # noqa: BLE001
            print(f"Failed to get endpoints: {exc}")
            return []

    async def watch_endpoints(
        self,
        service_name: str,
        namespace: str,
        callback: Callable[[str, object], Awaitable[None]],
    ) -> None:
        assert self._core_api is not None
        async with watch.Watch().stream(
            self._core_api.list_namespaced_endpoints,
            namespace=namespace,
            field_selector=f"metadata.name={service_name}",
        ) as stream:
            async for event in stream:
                event_type = event["type"]
                obj = event["object"]
                print(f"Endpoint {event_type}: {obj.metadata.name}")
                await callback(event_type, obj)

    async def watch_endpoints_with_retry(
        self,
        service_name: str,
        namespace: str,
        callback: Callable[[str, object], Awaitable[None]],
    ) -> None:
        while True:
            try:
                await self.watch_endpoints(service_name, namespace, callback)
            except Exception as exc:  # noqa: BLE001
                print(f"Watch error: {exc}")
                await asyncio.sleep(5)  # Reconnect after transient failures