//! FiberLB Controller - Manages LoadBalancer service VIP allocation
//!
//! This controller continuously reconciles tenant-scoped LoadBalancer Services
//! into FiberLB resources and removes stale publications when Services disappear
//! or stop being externally published.

use crate::auth::{authorized_request, issue_controller_token};
use crate::storage::{Storage, TenantRef};
use anyhow::Result;
use fiberlb_api::backend_service_client::BackendServiceClient;
use fiberlb_api::listener_service_client::ListenerServiceClient;
use fiberlb_api::load_balancer_service_client::LoadBalancerServiceClient;
use fiberlb_api::pool_service_client::PoolServiceClient;
use fiberlb_api::{
    CreateBackendRequest, CreateListenerRequest, CreateLoadBalancerRequest, CreatePoolRequest,
    DeleteBackendRequest, DeleteListenerRequest, DeleteLoadBalancerRequest, ListBackendsRequest,
    ListListenersRequest, ListLoadBalancersRequest, ListPoolsRequest, ListenerProtocol,
    LoadBalancer, PoolAlgorithm, PoolProtocol,
};
use k8shost_types::{LoadBalancerIngress, LoadBalancerStatus, Pod, Service, ServiceStatus};
use std::collections::{BTreeSet, HashMap, HashSet};
use std::sync::Arc;
use std::time::Duration;
use tokio::time::sleep;
use tonic::{transport::Channel, Code};
use tracing::{debug, info, warn};

const CONTROLLER_PRINCIPAL_ID: &str = "k8shost-controller";
const LB_ID_ANNOTATION: &str = "fiberlb.ultracloud.io/lb-id";
const POOL_ID_ANNOTATION: &str = "fiberlb.ultracloud.io/pool-id";

/// FiberLB controller for managing LoadBalancer service VIPs
pub struct FiberLbController {
    storage: Arc<Storage>,
    fiberlb_addr: String,
    iam_server_addr: String,
    interval: Duration,
}

impl FiberLbController {
    pub fn new(storage: Arc<Storage>, fiberlb_addr: String, iam_server_addr: String) -> Self {
        Self {
            storage,
            fiberlb_addr,
            iam_server_addr,
            interval: Duration::from_secs(10),
        }
    }

    pub async fn run(self: Arc<Self>) {
        info!(
            "FiberLB controller started (FiberLB at {}, {}s interval)",
            self.fiberlb_addr,
            self.interval.as_secs()
        );

        loop {
            if let Err(error) = self.reconcile_loadbalancers().await {
                warn!(error = %error, "FiberLB controller cycle failed");
            }

            sleep(self.interval).await;
        }
    }

    async fn reconcile_loadbalancers(&self) -> Result<()> {
        let tenants = self.reconciliation_tenants().await?;
        if tenants.is_empty() {
            debug!("No active tenants need FiberLB reconciliation");
            return Ok(());
        }

        for tenant in tenants {
            if let Err(error) = self.reconcile_tenant_loadbalancers(&tenant).await {
                warn!(
                    org_id = %tenant.org_id,
                    project_id = %tenant.project_id,
                    error = %error,
                    "failed to reconcile tenant load balancers"
                );
            }
        }

        Ok(())
    }

    async fn reconcile_tenant_loadbalancers(&self, tenant: &TenantRef) -> Result<()> {
        let services = self
            .storage
            .list_services(&tenant.org_id, &tenant.project_id, None)
            .await?;

        for service in services.iter().filter(|service| {
            !service_is_load_balancer(service) && has_load_balancer_state(service)
        }) {
            self.clear_service_load_balancer_state(
                &tenant.org_id,
                &tenant.project_id,
                service_namespace(service),
                &service.metadata.name,
            )
            .await?;
        }

        let desired_services = services
            .into_iter()
            .filter(service_is_load_balancer)
            .collect::<Vec<_>>();
        if !desired_services.is_empty() {
            self.storage
                .remember_controller_tenant("fiberlb", tenant)
                .await?;
        }
        let desired_names = desired_services
            .iter()
            .map(service_lb_name)
            .collect::<HashSet<_>>();

        let auth_token = issue_controller_token(
            &self.iam_server_addr,
            CONTROLLER_PRINCIPAL_ID,
            &tenant.org_id,
            &tenant.project_id,
        )
        .await?;

        let mut lb_client =
            match LoadBalancerServiceClient::connect(self.fiberlb_addr.clone()).await {
                Ok(client) => client,
                Err(error) => {
                    warn!(
                        address = %self.fiberlb_addr,
                        error = %error,
                        "failed to connect to FiberLB load balancer service"
                    );
                    return Ok(());
                }
            };
        let mut pool_client = PoolServiceClient::connect(self.fiberlb_addr.clone()).await?;
        let mut listener_client = ListenerServiceClient::connect(self.fiberlb_addr.clone()).await?;
        let mut backend_client = BackendServiceClient::connect(self.fiberlb_addr.clone()).await?;

        let current_load_balancers = self
            .list_load_balancers(&mut lb_client, tenant, &auth_token)
            .await?;
        for load_balancer in current_load_balancers
            .iter()
            .filter(|load_balancer| is_managed_load_balancer(load_balancer))
        {
            if !desired_names.contains(&load_balancer.name) {
                delete_load_balancer(&mut lb_client, &auth_token, &load_balancer.id).await?;
            }
        }

        for service in desired_services {
            self.reconcile_service_load_balancer(
                tenant,
                &mut lb_client,
                &mut pool_client,
                &mut listener_client,
                &mut backend_client,
                &auth_token,
                &current_load_balancers,
                &service,
            )
            .await?;
        }

        if desired_names.is_empty() {
            self.storage
                .forget_controller_tenant("fiberlb", tenant)
                .await?;
        }

        Ok(())
    }

    async fn reconciliation_tenants(&self) -> Result<Vec<TenantRef>> {
        let mut tenants = self
            .storage
            .list_active_tenants()
            .await?
            .into_iter()
            .collect::<BTreeSet<_>>();

        for tenant in self.storage.list_controller_tenants("fiberlb").await? {
            tenants.insert(tenant);
        }

        Ok(tenants.into_iter().collect())
    }

    async fn reconcile_service_load_balancer(
        &self,
        tenant: &TenantRef,
        lb_client: &mut LoadBalancerServiceClient<Channel>,
        pool_client: &mut PoolServiceClient<Channel>,
        listener_client: &mut ListenerServiceClient<Channel>,
        backend_client: &mut BackendServiceClient<Channel>,
        auth_token: &str,
        current_load_balancers: &[LoadBalancer],
        service: &Service,
    ) -> Result<()> {
        let namespace = service_namespace(service).to_string();
        let lb_name = service_lb_name(service);
        let description = service_lb_description(service);
        let annotated_lb_id = service.metadata.annotations.get(LB_ID_ANNOTATION);
        let existing_load_balancer = current_load_balancers.iter().find(|load_balancer| {
            annotated_lb_id.map(String::as_str) == Some(load_balancer.id.as_str())
                || load_balancer.name == lb_name
        });

        let load_balancer = ensure_load_balancer(
            lb_client,
            auth_token,
            tenant,
            &lb_name,
            &description,
            existing_load_balancer,
        )
        .await?;
        let pool = ensure_pool(
            pool_client,
            auth_token,
            &load_balancer.id,
            &format!("{lb_name}-pool"),
        )
        .await?;

        reconcile_listeners(
            listener_client,
            auth_token,
            &load_balancer.id,
            &pool.id,
            service,
            &lb_name,
        )
        .await?;

        let pods = self
            .storage
            .list_pods(
                &tenant.org_id,
                &tenant.project_id,
                Some(&namespace),
                if service.spec.selector.is_empty() {
                    None
                } else {
                    Some(&service.spec.selector)
                },
            )
            .await?;

        reconcile_backends(
            backend_client,
            auth_token,
            &pool.id,
            service,
            &lb_name,
            &pods,
        )
        .await?;

        self.reconcile_service_load_balancer_state(
            &tenant.org_id,
            &tenant.project_id,
            &namespace,
            &service.metadata.name,
            &load_balancer.id,
            &pool.id,
            non_empty_value(&load_balancer.vip_address),
        )
        .await?;

        Ok(())
    }

    async fn reconcile_service_load_balancer_state(
        &self,
        org_id: &str,
        project_id: &str,
        namespace: &str,
        name: &str,
        lb_id: &str,
        pool_id: &str,
        vip: Option<&str>,
    ) -> Result<()> {
        let Some(mut service) = self
            .storage
            .get_service(org_id, project_id, namespace, name)
            .await?
        else {
            return Ok(());
        };

        let mut changed = false;
        changed |= set_annotation(&mut service, LB_ID_ANNOTATION, lb_id);
        changed |= set_annotation(&mut service, POOL_ID_ANNOTATION, pool_id);
        changed |= set_load_balancer_status(&mut service, vip);

        if !changed {
            return Ok(());
        }

        bump_resource_version(&mut service);
        self.storage.put_service(&service).await?;
        Ok(())
    }

    async fn clear_service_load_balancer_state(
        &self,
        org_id: &str,
        project_id: &str,
        namespace: &str,
        name: &str,
    ) -> Result<()> {
        let Some(mut service) = self
            .storage
            .get_service(org_id, project_id, namespace, name)
            .await?
        else {
            return Ok(());
        };

        let mut changed = false;
        changed |= service
            .metadata
            .annotations
            .remove(LB_ID_ANNOTATION)
            .is_some();
        changed |= service
            .metadata
            .annotations
            .remove(POOL_ID_ANNOTATION)
            .is_some();

        let had_load_balancer = service
            .status
            .as_ref()
            .and_then(|status| status.load_balancer.as_ref())
            .is_some();
        if had_load_balancer {
            if let Some(status) = service.status.as_mut() {
                status.load_balancer = None;
            }
            if matches!(service.status.as_ref(), Some(status) if status.load_balancer.is_none()) {
                service.status = None;
            }
            changed = true;
        }

        if !changed {
            return Ok(());
        }

        bump_resource_version(&mut service);
        self.storage.put_service(&service).await?;
        Ok(())
    }

    async fn list_load_balancers(
        &self,
        client: &mut LoadBalancerServiceClient<Channel>,
        tenant: &TenantRef,
        auth_token: &str,
    ) -> Result<Vec<LoadBalancer>> {
        let mut load_balancers = Vec::new();
        let mut page_token = String::new();

        loop {
            let response = client
                .list_load_balancers(authorized_request(
                    ListLoadBalancersRequest {
                        org_id: tenant.org_id.clone(),
                        project_id: tenant.project_id.clone(),
                        page_size: 256,
                        page_token: page_token.clone(),
                    },
                    auth_token,
                ))
                .await?
                .into_inner();

            load_balancers.extend(response.loadbalancers);
            if response.next_page_token.is_empty() {
                break;
            }
            page_token = response.next_page_token;
        }

        Ok(load_balancers)
    }
}

#[derive(Clone)]
struct DesiredBackend {
    name: String,
    address: String,
    port: u32,
}

async fn ensure_load_balancer(
    client: &mut LoadBalancerServiceClient<Channel>,
    auth_token: &str,
    tenant: &TenantRef,
    name: &str,
    description: &str,
    existing: Option<&LoadBalancer>,
) -> Result<LoadBalancer> {
    if let Some(load_balancer) = existing {
        return Ok(load_balancer.clone());
    }

    Ok(client
        .create_load_balancer(authorized_request(
            CreateLoadBalancerRequest {
                name: name.to_string(),
                org_id: tenant.org_id.clone(),
                project_id: tenant.project_id.clone(),
                description: description.to_string(),
                vip_address: String::new(),
            },
            auth_token,
        ))
        .await?
        .into_inner()
        .loadbalancer
        .ok_or_else(|| anyhow::anyhow!("FiberLB returned empty CreateLoadBalancer response"))?)
}

async fn ensure_pool(
    client: &mut PoolServiceClient<Channel>,
    auth_token: &str,
    load_balancer_id: &str,
    name: &str,
) -> Result<fiberlb_api::Pool> {
    let mut page_token = String::new();

    loop {
        let response = client
            .list_pools(authorized_request(
                ListPoolsRequest {
                    loadbalancer_id: load_balancer_id.to_string(),
                    page_size: 256,
                    page_token: page_token.clone(),
                },
                auth_token,
            ))
            .await?
            .into_inner();

        if let Some(pool) = response.pools.into_iter().find(|pool| pool.name == name) {
            return Ok(pool);
        }

        if response.next_page_token.is_empty() {
            break;
        }
        page_token = response.next_page_token;
    }

    Ok(client
        .create_pool(authorized_request(
            CreatePoolRequest {
                name: name.to_string(),
                loadbalancer_id: load_balancer_id.to_string(),
                algorithm: PoolAlgorithm::RoundRobin as i32,
                protocol: PoolProtocol::Tcp as i32,
                session_persistence: None,
            },
            auth_token,
        ))
        .await?
        .into_inner()
        .pool
        .ok_or_else(|| anyhow::anyhow!("FiberLB returned empty CreatePool response"))?)
}

async fn reconcile_listeners(
    client: &mut ListenerServiceClient<Channel>,
    auth_token: &str,
    load_balancer_id: &str,
    pool_id: &str,
    service: &Service,
    lb_name: &str,
) -> Result<()> {
    let desired = service
        .spec
        .ports
        .iter()
        .map(|port| {
            (
                format!(
                    "{}-listener-{}",
                    lb_name,
                    port.name.as_deref().unwrap_or(&port.port.to_string())
                ),
                port.port as u32,
            )
        })
        .collect::<HashMap<_, _>>();
    let mut satisfied = HashSet::new();

    for listener in list_listeners(client, auth_token, load_balancer_id).await? {
        match desired.get(&listener.name) {
            Some(port)
                if listener.port == *port
                    && listener.protocol == ListenerProtocol::Tcp as i32
                    && listener.default_pool_id == pool_id
                    && !satisfied.contains(&listener.name) =>
            {
                satisfied.insert(listener.name.clone());
            }
            _ => delete_listener(client, auth_token, &listener.id).await?,
        }
    }

    for (name, port) in desired {
        if satisfied.contains(&name) {
            continue;
        }

        client
            .create_listener(authorized_request(
                CreateListenerRequest {
                    name,
                    loadbalancer_id: load_balancer_id.to_string(),
                    protocol: ListenerProtocol::Tcp as i32,
                    port,
                    default_pool_id: pool_id.to_string(),
                    tls_config: None,
                    connection_limit: 0,
                },
                auth_token,
            ))
            .await?;
    }

    Ok(())
}

async fn reconcile_backends(
    client: &mut BackendServiceClient<Channel>,
    auth_token: &str,
    pool_id: &str,
    service: &Service,
    lb_name: &str,
    pods: &[Pod],
) -> Result<()> {
    let desired = desired_backends(service, lb_name, pods)
        .into_iter()
        .map(|backend| (backend.name.clone(), backend))
        .collect::<HashMap<_, _>>();
    let mut satisfied = HashSet::new();

    for backend in list_backends(client, auth_token, pool_id).await? {
        match desired.get(&backend.name) {
            Some(target)
                if backend.address == target.address
                    && backend.port == target.port
                    && !satisfied.contains(&backend.name) =>
            {
                satisfied.insert(backend.name.clone());
            }
            _ => delete_backend(client, auth_token, &backend.id).await?,
        }
    }

    for backend in desired.into_values() {
        if satisfied.contains(&backend.name) {
            continue;
        }

        client
            .create_backend(authorized_request(
                CreateBackendRequest {
                    name: backend.name,
                    pool_id: pool_id.to_string(),
                    address: backend.address,
                    port: backend.port,
                    weight: 1,
                },
                auth_token,
            ))
            .await?;
    }

    Ok(())
}

async fn list_listeners(
    client: &mut ListenerServiceClient<Channel>,
    auth_token: &str,
    load_balancer_id: &str,
) -> Result<Vec<fiberlb_api::Listener>> {
    let mut listeners = Vec::new();
    let mut page_token = String::new();

    loop {
        let response = client
            .list_listeners(authorized_request(
                ListListenersRequest {
                    loadbalancer_id: load_balancer_id.to_string(),
                    page_size: 256,
                    page_token: page_token.clone(),
                },
                auth_token,
            ))
            .await?
            .into_inner();

        listeners.extend(response.listeners);
        if response.next_page_token.is_empty() {
            break;
        }
        page_token = response.next_page_token;
    }

    Ok(listeners)
}

async fn list_backends(
    client: &mut BackendServiceClient<Channel>,
    auth_token: &str,
    pool_id: &str,
) -> Result<Vec<fiberlb_api::Backend>> {
    let mut backends = Vec::new();
    let mut page_token = String::new();

    loop {
        let response = client
            .list_backends(authorized_request(
                ListBackendsRequest {
                    pool_id: pool_id.to_string(),
                    page_size: 256,
                    page_token: page_token.clone(),
                },
                auth_token,
            ))
            .await?
            .into_inner();

        backends.extend(response.backends);
        if response.next_page_token.is_empty() {
            break;
        }
        page_token = response.next_page_token;
    }

    Ok(backends)
}

fn desired_backends(service: &Service, lb_name: &str, pods: &[Pod]) -> Vec<DesiredBackend> {
    let mut desired = Vec::new();

    for pod in pods {
        let Some(pod_ip) = pod
            .status
            .as_ref()
            .and_then(|status| status.pod_ip.as_ref())
        else {
            continue;
        };

        for service_port in &service.spec.ports {
            let backend_port = service_port.target_port.unwrap_or(service_port.port);
            desired.push(DesiredBackend {
                name: format!("{}-backend-{}-{}", lb_name, pod.metadata.name, backend_port),
                address: pod_ip.clone(),
                port: backend_port as u32,
            });
        }
    }

    desired
}

fn service_is_load_balancer(service: &Service) -> bool {
    service.spec.r#type.as_deref() == Some("LoadBalancer")
}

fn has_load_balancer_state(service: &Service) -> bool {
    service.metadata.annotations.contains_key(LB_ID_ANNOTATION)
        || service
            .metadata
            .annotations
            .contains_key(POOL_ID_ANNOTATION)
        || service
            .status
            .as_ref()
            .and_then(|status| status.load_balancer.as_ref())
            .is_some()
}

fn service_namespace(service: &Service) -> &str {
    service.metadata.namespace.as_deref().unwrap_or("default")
}

fn service_lb_name(service: &Service) -> String {
    format!("{}.{}", service.metadata.name, service_namespace(service))
}

fn service_lb_description(service: &Service) -> String {
    format!(
        "k8s service {}/{}",
        service_namespace(service),
        service.metadata.name
    )
}

fn is_managed_load_balancer(load_balancer: &LoadBalancer) -> bool {
    load_balancer.description.starts_with("k8s service ")
}

fn non_empty_value(value: &str) -> Option<&str> {
    if value.is_empty() {
        None
    } else {
        Some(value)
    }
}

fn set_annotation(service: &mut Service, key: &str, value: &str) -> bool {
    match service.metadata.annotations.get(key) {
        Some(current) if current == value => false,
        _ => {
            service
                .metadata
                .annotations
                .insert(key.to_string(), value.to_string());
            true
        }
    }
}

fn set_load_balancer_status(service: &mut Service, vip: Option<&str>) -> bool {
    let desired = vip.map(|vip| ServiceStatus {
        load_balancer: Some(LoadBalancerStatus {
            ingress: vec![LoadBalancerIngress {
                ip: Some(vip.to_string()),
                hostname: None,
            }],
        }),
    });

    let current_vip = service
        .status
        .as_ref()
        .and_then(|status| status.load_balancer.as_ref())
        .and_then(|load_balancer| load_balancer.ingress.first())
        .and_then(|ingress| ingress.ip.as_deref());
    if current_vip == vip {
        return false;
    }

    service.status = desired;
    true
}

fn bump_resource_version(service: &mut Service) {
    let current = service
        .metadata
        .resource_version
        .as_deref()
        .and_then(|value| value.parse::<u64>().ok())
        .unwrap_or(0);
    service.metadata.resource_version = Some((current + 1).to_string());
}

async fn delete_load_balancer(
    client: &mut LoadBalancerServiceClient<Channel>,
    auth_token: &str,
    lb_id: &str,
) -> Result<()> {
    match client
        .delete_load_balancer(authorized_request(
            DeleteLoadBalancerRequest {
                id: lb_id.to_string(),
            },
            auth_token,
        ))
        .await
    {
        Ok(_) => Ok(()),
        Err(status) if status.code() == Code::NotFound => Ok(()),
        Err(status) => Err(status.into()),
    }
}

async fn delete_listener(
    client: &mut ListenerServiceClient<Channel>,
    auth_token: &str,
    listener_id: &str,
) -> Result<()> {
    match client
        .delete_listener(authorized_request(
            DeleteListenerRequest {
                id: listener_id.to_string(),
            },
            auth_token,
        ))
        .await
    {
        Ok(_) => Ok(()),
        Err(status) if status.code() == Code::NotFound => Ok(()),
        Err(status) => Err(status.into()),
    }
}

async fn delete_backend(
    client: &mut BackendServiceClient<Channel>,
    auth_token: &str,
    backend_id: &str,
) -> Result<()> {
    match client
        .delete_backend(authorized_request(
            DeleteBackendRequest {
                id: backend_id.to_string(),
            },
            auth_token,
        ))
        .await
    {
        Ok(_) => Ok(()),
        Err(status) if status.code() == Code::NotFound => Ok(()),
        Err(status) => Err(status.into()),
    }
}
