"""Merkle tree endpoints.

GET /v1/merkle/root           — Current signed Merkle root
GET /v1/merkle/proof/{agent_id} — Inclusion proof for an agent
"""

from __future__ import annotations

import json
from datetime import datetime, timezone

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession

from kredo_drift.server.auth import AuthenticatedAgent, require_api_key
from kredo_drift.server.db.models import Agent, ChainEntry
from kredo_drift.server.db.session import get_db
from kredo_drift.server.deps import get_service_signing_key, get_service_pubkey
from kredo_drift.merkle import build_merkle_tree, get_inclusion_proof

router = APIRouter(tags=["merkle"])

# Cached Merkle tree (rebuilt periodically)
_current_tree = None
_tree_built_at = None


async def _rebuild_tree(db: AsyncSession):
    """Rebuild the global Merkle tree from all agent chain heads."""
    global _current_tree, _tree_built_at

    # Get latest chain entry for each agent (subquery for max sequence)
    subq = (
        select(
            ChainEntry.agent_pubkey,
            func.max(ChainEntry.sequence).label("max_seq"),
        )
        .group_by(ChainEntry.agent_pubkey)
        .subquery()
    )

    result = await db.execute(
        select(ChainEntry.agent_pubkey, ChainEntry.chain_hash)
        .join(subq, (ChainEntry.agent_pubkey == subq.c.agent_pubkey) &
              (ChainEntry.sequence == subq.c.max_seq))
    )

    chain_heads = {row[0]: row[1] for row in result.all()}
    _current_tree = build_merkle_tree(chain_heads)
    _tree_built_at = datetime.now(timezone.utc)


@router.get("/merkle/root")
async def merkle_root(
    db: AsyncSession = Depends(get_db),
):
    """Get the current signed Merkle root. No auth required — public verifiability."""
    if _current_tree is None:
        await _rebuild_tree(db)

    service_sk = get_service_signing_key()
    root_sig = service_sk.sign(_current_tree.root.encode()).signature.hex()

    return {
        "root": _current_tree.root,
        "leaf_count": _current_tree.leaf_count,
        "built_at": _tree_built_at.isoformat() if _tree_built_at else None,
        "service_signature": f"ed25519:{root_sig}",
        "service_pubkey": get_service_pubkey(),
    }


@router.get("/merkle/proof/{agent_id}")
async def merkle_proof(
    agent_id: str,
    auth: AuthenticatedAgent = Depends(require_api_key),
    db: AsyncSession = Depends(get_db),
):
    """Get a Merkle inclusion proof for an agent."""
    if _current_tree is None:
        await _rebuild_tree(db)

    # Find agent pubkey
    agent_result = await db.execute(
        select(Agent).where(Agent.agent_id == agent_id)
    )
    agent = agent_result.scalar_one_or_none()
    if not agent:
        raise HTTPException(status_code=404, detail="Agent not found")

    proof = get_inclusion_proof(_current_tree, agent.pubkey)
    if proof is None:
        raise HTTPException(status_code=404, detail="Agent not in current Merkle tree")

    return {
        "agent_id": agent_id,
        "agent_pubkey": agent.pubkey,
        "root": _current_tree.root,
        "proof": proof,
        "leaf_count": _current_tree.leaf_count,
    }
