"""Key management endpoints.

POST /v1/keys/rotate  — Rotate agent key (signed by old key)
POST /v1/keys/recover — Recover via owner key
"""

from __future__ import annotations

from datetime import datetime, timezone

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

from kredo_drift.server.auth import (
    AuthenticatedAgent, require_api_key,
    generate_api_key, hash_api_key, verify_request_signature,
)
from kredo_drift.server.db.models import Agent, AuditEvent
from kredo_drift.server.db.session import get_db
from kredo_drift.canonical import canonical_json

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


class KeyRotateRequest(BaseModel):
    new_pubkey: str  # "ed25519:<hex>"
    signature: str  # old key signs: canonical({old_pubkey, new_pubkey, timestamp})


class KeyRecoverRequest(BaseModel):
    agent_id: str
    owner_pubkey: str
    new_agent_pubkey: str
    owner_signature: str  # owner signs: canonical({agent_id, new_agent_pubkey})


@router.post("/keys/rotate")
async def rotate_key(
    req: KeyRotateRequest,
    auth: AuthenticatedAgent = Depends(require_api_key),
    db: AsyncSession = Depends(get_db),
):
    """Rotate an agent's Ed25519 key. Must be signed by the current key."""
    # Validate new pubkey format
    if not req.new_pubkey.startswith("ed25519:") or len(req.new_pubkey) != 72:
        raise HTTPException(status_code=400, detail="Invalid pubkey format")

    if req.new_pubkey == auth.pubkey:
        raise HTTPException(status_code=400, detail="New key must differ from current key")

    # Verify signature by old key over rotation payload
    rotation_payload = canonical_json({
        "old_pubkey": auth.pubkey,
        "new_pubkey": req.new_pubkey,
    })
    if not verify_request_signature(rotation_payload, req.signature, auth.pubkey):
        raise HTTPException(status_code=401, detail="Invalid rotation signature")

    # Check new key isn't already in use
    existing = await db.execute(
        select(Agent).where(Agent.pubkey == req.new_pubkey)
    )
    if existing.scalar_one_or_none():
        raise HTTPException(status_code=409, detail="New pubkey already in use")

    # Create new agent record with new key, generate new API key
    old_agent = auth.agent
    new_api_key = generate_api_key()

    # WHY deactivate + create new row (not UPDATE): Agent rows are keyed by
    # pubkey. Creating a new row with the new pubkey preserves the old row
    # (and its FK references from chain_entries, baselines, etc.) intact.
    # The agent_id column links old and new rows for continuity.  This is
    # also compatible with append-only DB triggers that reject UPDATEs on
    # the agents table in strict deployments.
    old_agent.active = False

    new_agent = Agent(
        pubkey=req.new_pubkey,
        agent_id=old_agent.agent_id,
        name=old_agent.name,
        api_key_hash=hash_api_key(new_api_key),
        tier=old_agent.tier,
        owner_pubkey=old_agent.owner_pubkey,
    )
    db.add(new_agent)

    # WHY chain entry under OLD pubkey: The rotation event is the last thing
    # the old key signs. Recording it under the old pubkey's chain creates an
    # unbroken audit trail: the old chain ends with "I rotated to <new_key>",
    # and the new chain can reference this entry for provenance.
    from kredo_drift.server.chain_store import append_to_chain_async
    import hashlib as _hashlib
    rotation_hash = _hashlib.sha256(f"{auth.pubkey}:{req.new_pubkey}".encode()).hexdigest()
    await append_to_chain_async(auth.pubkey, "key_rotation", rotation_hash, db)

    db.add(AuditEvent(
        event_type="key_rotation",
        agent_pubkey=req.new_pubkey,
        detail={
            "old_pubkey": auth.pubkey,
            "new_pubkey": req.new_pubkey,
            "agent_id": old_agent.agent_id,
        },
    ))

    return {
        "status": "rotated",
        "agent_id": old_agent.agent_id,
        "new_pubkey": req.new_pubkey,
        "new_api_key": new_api_key,
    }


@router.post("/keys/recover")
async def recover_key(
    req: KeyRecoverRequest,
    db: AsyncSession = Depends(get_db),
):
    """Recover an agent's key using the linked owner key.

    This is for when the agent's key is lost. The owner can authorize
    a new key for the agent.
    """
    # Find the agent
    agent_result = await db.execute(
        select(Agent).where(
            Agent.agent_id == req.agent_id,
            Agent.owner_pubkey == req.owner_pubkey,
        )
    )
    agent = agent_result.scalar_one_or_none()
    if not agent:
        raise HTTPException(
            status_code=404,
            detail="Agent not found or owner pubkey doesn't match",
        )

    # Verify owner signature
    recovery_payload = canonical_json({
        "agent_id": req.agent_id,
        "new_agent_pubkey": req.new_agent_pubkey,
    })
    if not verify_request_signature(recovery_payload, req.owner_signature, req.owner_pubkey):
        raise HTTPException(status_code=401, detail="Invalid owner signature")

    # Deactivate old key
    agent.active = False

    # Create new agent record
    new_api_key = generate_api_key()
    new_agent = Agent(
        pubkey=req.new_agent_pubkey,
        agent_id=agent.agent_id,
        name=agent.name,
        api_key_hash=hash_api_key(new_api_key),
        tier=agent.tier,
        owner_pubkey=req.owner_pubkey,
    )
    db.add(new_agent)

    # Append recovery event to chain (under old pubkey)
    from kredo_drift.server.chain_store import append_to_chain_async
    import hashlib as _hashlib
    recovery_content = _hashlib.sha256(
        f"{agent.pubkey}:{req.new_agent_pubkey}:recovery".encode()
    ).hexdigest()
    await append_to_chain_async(agent.pubkey, "recovery", recovery_content, db)

    db.add(AuditEvent(
        event_type="key_recovery",
        agent_pubkey=req.new_agent_pubkey,
        detail={
            "agent_id": req.agent_id,
            "old_pubkey": agent.pubkey,
            "recovered_by": req.owner_pubkey,
        },
    ))

    return {
        "status": "recovered",
        "agent_id": req.agent_id,
        "new_pubkey": req.new_agent_pubkey,
        "new_api_key": new_api_key,
    }
