"""Baseline capture endpoints.

POST /v1/baseline/init    — Request assessment prompts
POST /v1/baseline/submit  — Submit encrypted responses
POST /v1/baseline/ack     — Agent countersigns baseline hash
GET  /v1/baseline/{agent_id} — Baseline metadata
"""

from __future__ import annotations

import hashlib
import json
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
from kredo_drift.server.db.models import Agent, Baseline, AuditEvent
from kredo_drift.server.db.session import get_db
from kredo_drift.server.deps import get_vectorizer, get_service_signing_key, get_service_pubkey
from kredo_drift.assessment import AssessmentProtocol
from kredo_drift.canonical import canonical_json
from kredo_drift.chain import append_to_chain
from kredo_drift.scoring import compute_drift_score
from kredo_drift.models import Dimension, PromptResponse

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


class BaselineInitRequest(BaseModel):
    sequence: int = 0


class BaselineInitResponse(BaseModel):
    prompts: list[dict]  # [{id, text, dimension, is_decoy}, ...]
    prompt_selection_hash: str


class BaselineSubmitRequest(BaseModel):
    responses: list[dict]  # [{prompt_id, dimension, response_text, is_decoy}, ...]
    agent_signature: str  # Ed25519 signature over canonical response payload


class BaselineSubmitResponse(BaseModel):
    content_hash: str
    prompt_count: int
    service_signature: str
    created_at: str


class BaselineAckRequest(BaseModel):
    content_hash: str
    agent_signature: str  # countersignature over the content_hash


@router.post("/baseline/init")
async def baseline_init(
    req: BaselineInitRequest,
    auth: AuthenticatedAgent = Depends(require_api_key),
) -> BaselineInitResponse:
    """Request assessment prompts for baseline capture."""
    protocol = AssessmentProtocol()
    prompt_set = protocol.select_prompts(auth.agent_id, sequence=req.sequence)

    prompts = [
        {
            "id": p.id,
            "text": p.text,
            "dimension": p.dimension,
            "is_decoy": p.is_decoy,
        }
        for p in prompt_set.prompts
    ]

    selection_hash = hashlib.sha256(
        json.dumps(sorted([p["id"] for p in prompts])).encode()
    ).hexdigest()

    return BaselineInitResponse(
        prompts=prompts,
        prompt_selection_hash=selection_hash,
    )


@router.post("/baseline/submit")
async def baseline_submit(
    req: BaselineSubmitRequest,
    auth: AuthenticatedAgent = Depends(require_api_key),
    db: AsyncSession = Depends(get_db),
) -> BaselineSubmitResponse:
    """Submit baseline responses. Server vectorizes, stores vector only, purges raw text."""
    from kredo_drift.server.auth import verify_request_signature

    # Verify agent signature over responses
    response_payload = canonical_json({"responses": req.responses})
    if not verify_request_signature(response_payload, req.agent_signature, auth.pubkey):
        raise HTTPException(status_code=401, detail="Invalid response signature")

    vectorizer = get_vectorizer()

    # Build PromptResponse objects
    prompt_responses = []
    for r in req.responses:
        dim_enum = (
            Dimension(r["dimension"])
            if not r.get("is_decoy", False)
            else Dimension.PERSONALITY
        )
        prompt_responses.append(PromptResponse(
            prompt_id=r["prompt_id"],
            dimension=dim_enum,
            prompt_text=r.get("prompt_text", ""),
            response_text=r["response_text"],
            is_decoy=r.get("is_decoy", False),
        ))

    # Filter out decoys for vectorization
    scored = [pr for pr in prompt_responses if not pr.is_decoy]

    # Vectorize all scored responses
    texts = [pr.response_text for pr in scored]
    vectors = vectorizer.encode_batch(texts)

    # Compute per-dimension vectors
    dimension_vectors = {}
    for dim in Dimension:
        dim_responses = [
            (i, pr) for i, pr in enumerate(scored) if pr.dimension == dim
        ]
        if dim_responses:
            dim_vecs = [vectors[i] for i, _ in dim_responses]
            avg = vectorizer.average_vectors(dim_vecs)
            dimension_vectors[dim.value] = avg.tolist()

    # Overall average vector
    overall_vector = vectorizer.average_vectors(list(vectors)).tolist()

    # Response hashes (for audit, raw text is NOT stored)
    response_hashes = [pr.response_hash() for pr in prompt_responses]
    prompt_ids = [pr.prompt_id for pr in prompt_responses]

    # Build content hash
    content_data = canonical_json({
        "agent_pubkey": auth.pubkey,
        "model_name": vectorizer.model_name,
        "vector": overall_vector,
        "dimension_vectors": dimension_vectors,
        "response_hashes": sorted(response_hashes),
        "prompt_ids": sorted(prompt_ids),
    })
    content_hash = hashlib.sha256(content_data).hexdigest()

    # Service signs the baseline
    service_sk = get_service_signing_key()
    service_sig = service_sk.sign(content_hash.encode()).signature.hex()

    # Store baseline (vector only, no raw text)
    baseline_record = Baseline(
        content_hash=content_hash,
        agent_pubkey=auth.pubkey,
        encrypted_vector=json.dumps(overall_vector),
        dimension_vectors=dimension_vectors,
        response_hashes=response_hashes,
        prompt_ids=prompt_ids,
        prompt_count=len(prompt_responses),
        model_name=vectorizer.model_name,
        agent_signature=req.agent_signature,
        service_signature=f"ed25519:{service_sig}",
    )
    db.add(baseline_record)

    # Append to hash chain
    from kredo_drift.server.chain_store import append_to_chain_async
    chain_entry = await append_to_chain_async(auth.pubkey, "baseline", content_hash, db)
    baseline_record.chain_hash = chain_entry.chain_hash

    db.add(AuditEvent(
        event_type="baseline_submit",
        agent_pubkey=auth.pubkey,
        detail={"content_hash": content_hash, "prompt_count": len(prompt_responses)},
    ))

    return BaselineSubmitResponse(
        content_hash=content_hash,
        prompt_count=len(prompt_responses),
        service_signature=f"ed25519:{service_sig}",
        created_at=datetime.now(timezone.utc).isoformat(),
    )


@router.post("/baseline/ack")
async def baseline_ack(
    req: BaselineAckRequest,
    auth: AuthenticatedAgent = Depends(require_api_key),
    db: AsyncSession = Depends(get_db),
):
    """Agent countersigns the baseline hash, completing dual-signature."""
    from kredo_drift.server.auth import verify_request_signature

    # Verify countersignature
    if not verify_request_signature(
        req.content_hash.encode(), req.agent_signature, auth.pubkey
    ):
        raise HTTPException(status_code=401, detail="Invalid countersignature")

    # Update baseline with countersignature
    result = await db.execute(
        select(Baseline).where(
            Baseline.content_hash == req.content_hash,
            Baseline.agent_pubkey == auth.pubkey,
        )
    )
    bl = result.scalar_one_or_none()
    if not bl:
        raise HTTPException(status_code=404, detail="Baseline not found")

    # Note: For append-only tables in production, we'd store the ack as a
    # separate chain entry rather than updating. For the initial implementation,
    # the ack signature was already set during submit.

    db.add(AuditEvent(
        event_type="baseline_ack",
        agent_pubkey=auth.pubkey,
        detail={"content_hash": req.content_hash},
    ))

    return {"status": "acknowledged", "content_hash": req.content_hash}


@router.get("/baseline/{agent_id}")
async def get_baseline(
    agent_id: str,
    auth: AuthenticatedAgent = Depends(require_api_key),
    db: AsyncSession = Depends(get_db),
):
    """Get baseline metadata (not the vector itself)."""
    # Look up agent
    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")

    # Get latest baseline
    bl_result = await db.execute(
        select(Baseline)
        .where(Baseline.agent_pubkey == agent.pubkey)
        .order_by(Baseline.created_at.desc())
        .limit(1)
    )
    bl = bl_result.scalar_one_or_none()
    if not bl:
        raise HTTPException(status_code=404, detail="No baseline found")

    return {
        "content_hash": bl.content_hash,
        "agent_id": agent_id,
        "prompt_count": bl.prompt_count,
        "model_name": bl.model_name,
        "created_at": bl.created_at.isoformat(),
        "has_agent_signature": bool(bl.agent_signature),
        "has_service_signature": bool(bl.service_signature),
    }
