"""Drift test endpoints.

POST /v1/test/init     — Request rotated test prompts
POST /v1/test/submit   — Submit encrypted test responses
POST /v1/test/ack      — Agent countersigns result
GET  /v1/test/{test_id} — Test result
"""

from __future__ import annotations

import hashlib
import json
from datetime import datetime, timezone

import numpy as np
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
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, Baseline, DriftTestRecord, AuditEvent
from kredo_drift.server.db.session import get_db
from kredo_drift.server.deps import get_vectorizer, get_service_signing_key
from kredo_drift.assessment import AssessmentProtocol
from kredo_drift.canonical import canonical_json
from kredo_drift.models import Dimension, PromptResponse, DIMENSION_WEIGHTS
from kredo_drift.scoring import (
    cosine_similarity,
    similarity_to_drift_score,
    classify_drift,
)

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


class TestInitRequest(BaseModel):
    baseline_hash: str | None = None  # optional, uses latest if not provided


class TestInitResponse(BaseModel):
    prompts: list[dict]
    prompt_selection_hash: str
    baseline_hash: str
    test_sequence: int


class TestSubmitRequest(BaseModel):
    baseline_hash: str
    responses: list[dict]
    agent_signature: str


class TestSubmitResponse(BaseModel):
    content_hash: str
    score: float
    classification: str
    dimension_scores: list[dict]
    service_signature: str


class TestAckRequest(BaseModel):
    content_hash: str
    agent_signature: str


@router.post("/test/init")
async def test_init(
    req: TestInitRequest,
    auth: AuthenticatedAgent = Depends(require_api_key),
    db: AsyncSession = Depends(get_db),
) -> TestInitResponse:
    """Request rotated prompts for a drift test."""
    # Find agent
    agent_result = await db.execute(
        select(Agent).where(Agent.pubkey == auth.pubkey)
    )
    agent = agent_result.scalar_one_or_none()
    if not agent:
        raise HTTPException(status_code=404, detail="Agent not found")

    # Find baseline
    if req.baseline_hash:
        bl_result = await db.execute(
            select(Baseline).where(
                Baseline.content_hash == req.baseline_hash,
                Baseline.agent_pubkey == auth.pubkey,
            )
        )
    else:
        bl_result = await db.execute(
            select(Baseline)
            .where(Baseline.agent_pubkey == auth.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")

    # Determine test sequence
    count_result = await db.execute(
        select(func.count()).select_from(DriftTestRecord).where(
            DriftTestRecord.agent_pubkey == auth.pubkey
        )
    )
    sequence = count_result.scalar() + 1

    # Select rotated prompts
    protocol = AssessmentProtocol()
    prompt_set = protocol.select_prompts(auth.agent_id, sequence=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 TestInitResponse(
        prompts=prompts,
        prompt_selection_hash=selection_hash,
        baseline_hash=bl.content_hash,
        test_sequence=sequence,
    )


@router.post("/test/submit")
async def test_submit(
    req: TestSubmitRequest,
    auth: AuthenticatedAgent = Depends(require_api_key),
    db: AsyncSession = Depends(get_db),
) -> TestSubmitResponse:
    """Submit test responses. Server vectorizes, scores against baseline."""
    from kredo_drift.server.auth import verify_request_signature

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

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

    vectorizer = get_vectorizer()
    baseline_vector = np.array(json.loads(bl.encrypted_vector))
    baseline_dim_vectors = bl.dimension_vectors or {}

    # Build prompt responses
    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),
        ))

    scored = [pr for pr in prompt_responses if not pr.is_decoy]
    texts = [pr.response_text for pr in scored]
    vectors = vectorizer.encode_batch(texts)

    # Per-dimension scoring
    dimension_scores = []
    test_dim_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]
            test_avg = vectorizer.average_vectors(dim_vecs)
            test_dim_vectors[dim.value] = test_avg.tolist()

            if dim.value in baseline_dim_vectors:
                baseline_dim_vec = np.array(baseline_dim_vectors[dim.value])
                sim = cosine_similarity(baseline_dim_vec, test_avg)
                dim_score = similarity_to_drift_score(sim)
                dimension_scores.append({
                    "dimension": dim.value,
                    "score": round(dim_score, 2),
                    "similarity": round(float(sim), 4),
                    "prompt_count": len(dim_responses),
                })

    # Aggregate score
    overall_test_vector = vectorizer.average_vectors(list(vectors))
    overall_sim = cosine_similarity(baseline_vector, overall_test_vector)
    overall_score = similarity_to_drift_score(overall_sim)
    classification = classify_drift(overall_score)

    # Weighted aggregate from dimension scores
    if dimension_scores:
        weighted_sum = 0.0
        weight_sum = 0.0
        for ds in dimension_scores:
            dim = Dimension(ds["dimension"])
            w = DIMENSION_WEIGHTS.get(dim, 0.0)
            weighted_sum += ds["score"] * w
            weight_sum += w
        if weight_sum > 0:
            overall_score = weighted_sum / weight_sum
            classification = classify_drift(overall_score)

    # Content hash
    response_hashes = [pr.response_hash() for pr in prompt_responses]
    prompt_ids = [pr.prompt_id for pr in prompt_responses]

    content_data = canonical_json({
        "agent_pubkey": auth.pubkey,
        "baseline_hash": req.baseline_hash,
        "model_name": vectorizer.model_name,
        "vector": overall_test_vector.tolist(),
        "dimension_vectors": test_dim_vectors,
        "response_hashes": sorted(response_hashes),
        "prompt_ids": sorted(prompt_ids),
        "score": round(overall_score, 2),
        "classification": classification.value,
    })
    content_hash = hashlib.sha256(content_data).hexdigest()

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

    # Store test record
    test_record = DriftTestRecord(
        content_hash=content_hash,
        agent_pubkey=auth.pubkey,
        baseline_hash=req.baseline_hash,
        encrypted_vector=json.dumps(overall_test_vector.tolist()),
        dimension_vectors=test_dim_vectors,
        response_hashes=response_hashes,
        prompt_ids=prompt_ids,
        prompt_count=len(prompt_responses),
        score=round(overall_score, 2),
        dimension_scores=dimension_scores,
        classification=classification.value,
        model_name=vectorizer.model_name,
        agent_signature=req.agent_signature,
        service_signature=f"ed25519:{service_sig}",
    )
    db.add(test_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, "test", content_hash, db)
    test_record.chain_hash = chain_entry.chain_hash

    db.add(AuditEvent(
        event_type="test_submit",
        agent_pubkey=auth.pubkey,
        detail={"content_hash": content_hash, "score": round(overall_score, 2)},
    ))

    return TestSubmitResponse(
        content_hash=content_hash,
        score=round(overall_score, 2),
        classification=classification.value,
        dimension_scores=dimension_scores,
        service_signature=f"ed25519:{service_sig}",
    )


@router.post("/test/ack")
async def test_ack(
    req: TestAckRequest,
    auth: AuthenticatedAgent = Depends(require_api_key),
    db: AsyncSession = Depends(get_db),
):
    """Agent countersigns the test result."""
    from kredo_drift.server.auth import verify_request_signature

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

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

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


@router.get("/test/{test_id}")
async def get_test(
    test_id: str,
    auth: AuthenticatedAgent = Depends(require_api_key),
    db: AsyncSession = Depends(get_db),
):
    """Get a specific test result by content hash."""
    result = await db.execute(
        select(DriftTestRecord).where(
            DriftTestRecord.content_hash == test_id,
            DriftTestRecord.agent_pubkey == auth.pubkey,
        )
    )
    test = result.scalar_one_or_none()
    if not test:
        raise HTTPException(status_code=404, detail="Test not found")

    return {
        "content_hash": test.content_hash,
        "baseline_hash": test.baseline_hash,
        "score": test.score,
        "classification": test.classification,
        "dimension_scores": test.dimension_scores,
        "prompt_count": test.prompt_count,
        "created_at": test.created_at.isoformat(),
        "has_agent_signature": bool(test.agent_signature),
        "has_service_signature": bool(test.service_signature),
    }
