package com.cameleer.server.app.controller;

import com.cameleer.server.app.config.AgentRegistryConfig;
import com.cameleer.server.app.dto.AgentInstanceResponse;
import com.cameleer.server.app.dto.AgentRefreshRequest;
import com.cameleer.server.app.dto.AgentRefreshResponse;
import com.cameleer.server.app.dto.AgentRegistrationRequest;
import com.cameleer.server.app.dto.AgentRegistrationResponse;
import com.cameleer.server.app.dto.ErrorResponse;
import com.cameleer.common.model.HeartbeatRequest;
import com.cameleer.server.app.security.BootstrapTokenValidator;
import com.cameleer.server.app.security.JwtAuthenticationFilter;
import com.cameleer.server.core.admin.AuditCategory;
import com.cameleer.server.core.admin.AuditResult;
import com.cameleer.server.core.admin.AuditService;
import com.cameleer.server.core.agent.AgentEventService;
import com.cameleer.server.core.agent.AgentInfo;
import com.cameleer.server.core.agent.AgentRegistryService;
import com.cameleer.server.core.agent.AgentState;
import com.cameleer.server.core.agent.RouteStateRegistry;
import com.cameleer.server.core.security.Ed25519SigningService;
import com.cameleer.server.core.storage.RouteCatalogStore;
import com.cameleer.server.core.security.InvalidTokenException;
import com.cameleer.server.core.security.JwtService;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.media.Content;
import io.swagger.v3.oas.annotations.media.Schema;
import io.swagger.v3.oas.annotations.responses.ApiResponse;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.servlet.http.HttpServletRequest;
import org.slf4j.Logger;
import org.springframework.web.servlet.support.ServletUriComponentsBuilder;
import org.slf4j.LoggerFactory;
import org.springframework.http.ResponseEntity;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.sql.Timestamp;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Agent registration, heartbeat, listing, and token refresh endpoints.
 */
@RestController
@RequestMapping("/api/v1/agents")
@Tag(name = "Agent Management", description = "Agent registration and lifecycle endpoints")
public class AgentRegistrationController {

    private static final Logger log = LoggerFactory.getLogger(AgentRegistrationController.class);
    private static final String BEARER_PREFIX = "Bearer ";

    private final AgentRegistryService registryService;
    private final AgentRegistryConfig config;
    private final BootstrapTokenValidator bootstrapTokenValidator;
    private final JwtService jwtService;
    private final Ed25519SigningService ed25519SigningService;
    private final AgentEventService agentEventService;
    private final AuditService auditService;
    private final JdbcTemplate jdbc;
    private final RouteStateRegistry routeStateRegistry;
    private final RouteCatalogStore routeCatalogStore;

    public AgentRegistrationController(AgentRegistryService registryService,
                                        AgentRegistryConfig config,
                                        BootstrapTokenValidator bootstrapTokenValidator,
                                        JwtService jwtService,
                                        Ed25519SigningService ed25519SigningService,
                                        AgentEventService agentEventService,
                                        AuditService auditService,
                                        @org.springframework.beans.factory.annotation.Qualifier("clickHouseJdbcTemplate") JdbcTemplate jdbc,
                                        RouteStateRegistry routeStateRegistry,
                                        RouteCatalogStore routeCatalogStore) {
        this.registryService = registryService;
        this.config = config;
        this.bootstrapTokenValidator = bootstrapTokenValidator;
        this.jwtService = jwtService;
        this.ed25519SigningService = ed25519SigningService;
        this.agentEventService = agentEventService;
        this.auditService = auditService;
        this.jdbc = jdbc;
        this.routeStateRegistry = routeStateRegistry;
        this.routeCatalogStore = routeCatalogStore;
    }

    @PostMapping("/register")
    @Operation(summary = "Register an agent",
            description = "Registers a new agent or re-registers an existing one. "
                    + "Requires bootstrap token in Authorization header.")
    @ApiResponse(responseCode = "200", description = "Agent registered successfully")
    @ApiResponse(responseCode = "400", description = "Invalid registration payload",
            content = @Content(schema = @Schema(implementation = ErrorResponse.class)))
    @ApiResponse(responseCode = "401", description = "Missing or invalid bootstrap token")
    public ResponseEntity<AgentRegistrationResponse> register(
            @RequestBody AgentRegistrationRequest request,
            HttpServletRequest httpRequest) {
        // Validate bootstrap token
        String authHeader = httpRequest.getHeader("Authorization");
        String bootstrapToken = null;
        if (authHeader != null && authHeader.startsWith(BEARER_PREFIX)) {
            bootstrapToken = authHeader.substring(BEARER_PREFIX.length());
        }
        if (bootstrapToken == null || !bootstrapTokenValidator.validate(bootstrapToken)) {
            return ResponseEntity.status(401).build();
        }

        if (request.instanceId() == null || request.instanceId().isBlank()) {
            return ResponseEntity.badRequest().build();
        }
        if (request.environmentId() == null || request.environmentId().isBlank()) {
            String remote = httpRequest.getRemoteAddr();
            log.warn("Agent registration rejected (no environmentId): instanceId={} remote={}",
                    request.instanceId(), remote);
            return ResponseEntity.badRequest().build();
        }

        String application = request.applicationId() != null ? request.applicationId() : "default";
        String environmentId = request.environmentId();
        List<String> routeIds = request.routeIds() != null ? request.routeIds() : List.of();
        var capabilities = request.capabilities() != null ? request.capabilities() : Collections.<String, Object>emptyMap();

        boolean reRegistration = registryService.findById(request.instanceId()) != null;

        AgentInfo agent = registryService.register(
                request.instanceId(), request.instanceId(), application, environmentId,
                request.version(), routeIds, capabilities);

        // Persist routes in catalog for server-restart recovery
        if (!routeIds.isEmpty()) {
            routeCatalogStore.upsert(application, environmentId, routeIds);
        }

        if (reRegistration) {
            log.info("Agent re-registered: {} (application={}, routes={}, capabilities={})",
                    request.instanceId(), application, routeIds.size(), capabilities.keySet());
            agentEventService.recordEvent(request.instanceId(), application, environmentId, "RE_REGISTERED",
                    "Agent re-registered with " + routeIds.size() + " routes");
        } else {
            log.info("Agent registered: {} (application={}, routes={})",
                    request.instanceId(), application, routeIds.size());
            agentEventService.recordEvent(request.instanceId(), application, environmentId, "REGISTERED",
                    "Agent registered: " + request.instanceId());
        }

        auditService.log(request.instanceId(), reRegistration ? "agent_reregister" : "agent_register",
                AuditCategory.AGENT, request.instanceId(),
                Map.of("application", application, "routeCount", routeIds.size(),
                        "reRegistration", reRegistration),
                AuditResult.SUCCESS, httpRequest);

        // Issue JWT tokens with AGENT role + environment
        List<String> roles = List.of("AGENT");
        String accessToken = jwtService.createAccessToken(request.instanceId(), application, environmentId, roles);
        String refreshToken = jwtService.createRefreshToken(request.instanceId(), application, environmentId, roles);

        String sseEndpoint = ServletUriComponentsBuilder.fromCurrentContextPath()
                .path("/api/v1/agents/{id}/events")
                .buildAndExpand(agent.instanceId())
                .toUriString();

        return ResponseEntity.ok(new AgentRegistrationResponse(
                agent.instanceId(),
                sseEndpoint,
                config.getHeartbeatIntervalMs(),
                ed25519SigningService.getPublicKeyBase64(),
                accessToken,
                refreshToken
        ));
    }

    @PostMapping("/{id}/refresh")
    @Operation(summary = "Refresh access token",
            description = "Issues a new access JWT from a valid refresh token")
    @ApiResponse(responseCode = "200", description = "New access token issued")
    @ApiResponse(responseCode = "401", description = "Invalid or expired refresh token")
    @ApiResponse(responseCode = "404", description = "Agent not found")
    public ResponseEntity<AgentRefreshResponse> refresh(@PathVariable String id,
                                                         @RequestBody AgentRefreshRequest request,
                                                         HttpServletRequest httpRequest) {
        if (request.refreshToken() == null || request.refreshToken().isBlank()) {
            return ResponseEntity.status(401).build();
        }

        // Validate refresh token
        JwtService.JwtValidationResult result;
        try {
            result = jwtService.validateRefreshToken(request.refreshToken());
        } catch (InvalidTokenException e) {
            log.debug("Refresh token validation failed: {}", e.getMessage());
            return ResponseEntity.status(401).build();
        }

        String agentId = result.subject();

        // Verify agent ID in path matches token
        if (!id.equals(agentId)) {
            log.debug("Refresh token agent ID mismatch: path={}, token={}", id, agentId);
            return ResponseEntity.status(401).build();
        }

        // Preserve roles and application from refresh token
        List<String> roles = result.roles().isEmpty()
                ? List.of("AGENT") : result.roles();
        String application = result.application() != null ? result.application() : "default";
        String environment = result.environment();

        // Refresh-token env claim is required — agents without env shouldn't have gotten a token in the first place.
        if (environment == null || environment.isBlank()) {
            log.warn("Refresh token has no environment claim: agentId={}", agentId);
            return ResponseEntity.status(401).build();
        }

        // Prefer registry if agent is still registered (covers env edits on re-registration)
        AgentInfo agent = registryService.findById(agentId);
        if (agent != null) {
            application = agent.applicationId();
            environment = agent.environmentId();
        }

        String newAccessToken = jwtService.createAccessToken(agentId, application, environment, roles);
        String newRefreshToken = jwtService.createRefreshToken(agentId, application, environment, roles);

        auditService.log(agentId, "agent_token_refresh", AuditCategory.AUTH, agentId,
                null, AuditResult.SUCCESS, httpRequest);

        return ResponseEntity.ok(new AgentRefreshResponse(newAccessToken, newRefreshToken));
    }

    @PostMapping("/{id}/heartbeat")
    @Operation(summary = "Agent heartbeat ping",
            description = "Updates the agent's last heartbeat timestamp. Auto-registers the agent if not in registry (e.g. after server restart).")
    @ApiResponse(responseCode = "200", description = "Heartbeat accepted")
    public ResponseEntity<Void> heartbeat(@PathVariable String id,
                                           @RequestBody(required = false) HeartbeatRequest request,
                                           HttpServletRequest httpRequest) {
        Map<String, Object> capabilities = request != null ? request.getCapabilities() : null;
        String heartbeatEnv = request != null ? request.getEnvironmentId() : null;
        List<String> routeIds = request != null && request.getRouteStates() != null
                ? List.copyOf(request.getRouteStates().keySet()) : null;
        boolean found = registryService.heartbeat(id, routeIds, capabilities);
        if (!found) {
            // Auto-heal: re-register agent from heartbeat body + JWT claims after server restart
            var jwtResult = (JwtService.JwtValidationResult) httpRequest.getAttribute(
                    JwtAuthenticationFilter.JWT_RESULT_ATTR);
            if (jwtResult != null) {
                String application = jwtResult.application() != null ? jwtResult.application() : "default";
                // Env: prefer heartbeat body (current), else JWT claim. No silent default.
                String env = (heartbeatEnv != null && !heartbeatEnv.isBlank())
                        ? heartbeatEnv
                        : jwtResult.environment();
                if (env == null || env.isBlank()) {
                    log.warn("Heartbeat auto-heal rejected (no environment on JWT or body): agentId={}", id);
                    return ResponseEntity.status(400).build();
                }
                Map<String, Object> caps = capabilities != null ? capabilities : Map.of();
                List<String> healRouteIds = routeIds != null ? routeIds : List.of();
                registryService.register(id, id, application, env, "unknown",
                        healRouteIds, caps);
                registryService.heartbeat(id);
                log.info("Auto-registered agent {} (app={}, env={}, routes={}) from heartbeat after server restart",
                        id, application, env, healRouteIds.size());
            } else {
                return ResponseEntity.notFound().build();
            }
        }

        if (routeIds != null && !routeIds.isEmpty()) {
            AgentInfo agent = registryService.findById(id);
            if (agent != null) {
                // Update route states from heartbeat
                if (request != null && request.getRouteStates() != null) {
                    for (var entry : request.getRouteStates().entrySet()) {
                        RouteStateRegistry.RouteState state = parseRouteState(entry.getValue());
                        if (state != null) {
                            routeStateRegistry.setState(agent.applicationId(), entry.getKey(), state);
                        }
                    }
                }
                // Persist routes in catalog for server-restart recovery
                routeCatalogStore.upsert(agent.applicationId(), agent.environmentId(), routeIds);
            }
        }

        return ResponseEntity.ok().build();
    }

    private RouteStateRegistry.RouteState parseRouteState(String state) {
        if (state == null) return null;
        return switch (state) {
            case "Started" -> RouteStateRegistry.RouteState.STARTED;
            case "Stopped" -> RouteStateRegistry.RouteState.STOPPED;
            case "Suspended" -> RouteStateRegistry.RouteState.SUSPENDED;
            default -> null;
        };
    }

    @PostMapping("/{id}/deregister")
    @Operation(summary = "Deregister agent",
            description = "Removes the agent from the registry. Called by agents during graceful shutdown.")
    @ApiResponse(responseCode = "200", description = "Agent deregistered")
    @ApiResponse(responseCode = "404", description = "Agent not registered")
    public ResponseEntity<Void> deregister(@PathVariable String id, HttpServletRequest httpRequest) {
        AgentInfo agent = registryService.findById(id);
        if (agent == null) {
            return ResponseEntity.notFound().build();
        }
        String applicationId = agent.applicationId();
        String environment = agent.environmentId();
        registryService.deregister(id);
        agentEventService.recordEvent(id, applicationId, environment, "DEREGISTERED", "Agent deregistered");
        auditService.log(id, "agent_deregister", AuditCategory.AGENT, id, null, AuditResult.SUCCESS, httpRequest);
        return ResponseEntity.ok().build();
    }

    // Agent list moved to AgentListController at /api/v1/environments/{envSlug}/agents.
    // Agent register/refresh/heartbeat/deregister remain here at /api/v1/agents/** —
    // these are JWT-authoritative and intentionally flat (env from token).
}
