"""DriftClient — Clean public SDK for Kredo Drift.

Supports two modes:
  - **Local mode** (development/testing): Direct SQLite, same as Phase 0 MVP.
  - **HTTP mode** (production): Connects to hosted API at api.aikredo.com.

Local mode (default, backward-compatible):

    from kredo_drift import DriftClient, make_response

    client = DriftClient("my-agent")
    prompts = client.get_prompts()
    responses = [make_response(p, agent.respond(p.text)) for p in prompts]
    baseline = client.create_baseline(responses)

HTTP mode (production):

    client = DriftClient("my-agent", api_key="...", endpoint="https://api.aikredo.com/drift")
    prompts = client.get_prompts()
    responses = [make_response(p, agent.respond(p.text)) for p in prompts]
    result = client.run_test(responses)

Recovery (the killer feature):

    client = DriftClient("new-agent", api_key="...", endpoint="...")
    identity = client.recover_identity("lost-agent", owner_key=owner_signing_key)
    result = client.run_recovery_test(identity, new_agent_responses)
"""

from __future__ import annotations

import hashlib
import json
import time
from pathlib import Path
from typing import Optional

from nacl.signing import SigningKey

from kredo_drift.models import (
    Dimension,
    DriftAttestation,
    DriftBaseline,
    DriftResult,
    PromptResponse,
)
from kredo_drift.prompts import Prompt
from kredo_drift.signing import sign_attestation
from kredo_drift.recovery import RecoveredIdentity, RecoveryResult, assess_recovery_match


def make_response(prompt: Prompt, response_text: str) -> PromptResponse:
    """Build a PromptResponse from a Prompt and the agent's response text.

    Handles the dimension mapping (string -> enum) and decoy flagging
    automatically, so callers don't need to know the internals.
    """
    # WHY PERSONALITY placeholder for decoys and fidelity: Decoy prompts have
    # dimension="decoy" and fidelity prompts have dimension="fidelity" — neither
    # is a valid Dimension enum. We need *some* enum value for the PromptResponse
    # model, but it doesn't matter which — decoys are filtered out before scoring,
    # and fidelity responses are separated into their own vector space.
    is_fidelity = prompt.dimension == "fidelity"
    dim_enum = (
        Dimension(prompt.dimension)
        if not prompt.is_decoy and not is_fidelity
        else Dimension.PERSONALITY
    )
    return PromptResponse(
        prompt_id=prompt.id,
        dimension=dim_enum,
        prompt_text=prompt.text,
        response_text=response_text,
        is_decoy=prompt.is_decoy,
        is_fidelity=is_fidelity,
        prompt_type=prompt.prompt_type,
    )


# ── Key management helpers ──

_KEY_DIR = Path.home() / ".config" / "kredo-drift" / "keys"


def _load_or_generate_key(agent_id: str) -> tuple[SigningKey, str]:
    """Load an agent's signing key from disk, or generate one."""
    key_path = _KEY_DIR / f"{agent_id}.key"

    if key_path.exists():
        sk = SigningKey(key_path.read_bytes())
        pubkey = f"ed25519:{sk.verify_key.encode().hex()}"
        return sk, pubkey

    # Generate new key
    sk = SigningKey.generate()
    key_path.parent.mkdir(parents=True, exist_ok=True)
    key_path.write_bytes(bytes(sk))
    # WHY 0o600: Ed25519 private keys must be readable only by the owner.
    # This matches SSH key conventions and prevents other users/processes
    # on the same machine from reading the agent's signing key.
    key_path.chmod(0o600)
    pubkey = f"ed25519:{sk.verify_key.encode().hex()}"
    return sk, pubkey


class DriftClient:
    """High-level client for running drift tests.

    This is the recommended entry point for programmatic use of Kredo Drift.

    Local mode (default):
        DriftClient(agent_id, db_path=...) — stores everything in SQLite.

    HTTP mode (production):
        DriftClient(agent_id, api_key=..., endpoint=...) — talks to hosted API.
    """

    def __init__(
        self,
        agent_id: str,
        *,
        # Local mode params
        db_path: str | None = None,
        server_secret: str = "kredo-drift-v0",
        vectorizer=None,
        # HTTP mode params
        api_key: str | None = None,
        endpoint: str | None = None,
        signing_key: SigningKey | None = None,
    ):
        self.agent_id = agent_id
        self._api_key = api_key
        self._endpoint = endpoint.rstrip("/") if endpoint else None
        self._signing_key = signing_key
        self._pubkey: str | None = None

        if self._endpoint and self._api_key:
            # HTTP mode
            self._mode = "http"
            # Auto-manage keys
            if not self._signing_key:
                self._signing_key, self._pubkey = _load_or_generate_key(agent_id)
            else:
                self._pubkey = f"ed25519:{signing_key.verify_key.encode().hex()}"
        else:
            # Local mode (backward-compatible)
            self._mode = "local"
            self.server_secret = server_secret
            from kredo_drift.baseline import BaselineManager
            from kredo_drift.store import DriftStore
            store = DriftStore(db_path=db_path) if db_path else DriftStore()
            self._manager = BaselineManager(
                store=store,
                vectorizer=vectorizer,
            )

    @property
    def mode(self) -> str:
        """Return 'local' or 'http'."""
        return self._mode

    # ── HTTP helpers ──

    def _http_request(
        self,
        method: str,
        path: str,
        json_data: dict | None = None,
        max_retries: int = 3,
    ) -> dict:
        """Make an authenticated HTTP request with retry + exponential backoff."""
        import httpx

        url = f"{self._endpoint}{path}"
        headers = {"X-API-Key": self._api_key}

        # Sign request body if we have a signing key
        if json_data and self._signing_key:
            from kredo_drift.canonical import canonical_json
            body_bytes = canonical_json(json_data)
            sig = self._signing_key.sign(body_bytes).signature.hex()
            headers["X-Signature"] = f"ed25519:{sig}"

        for attempt in range(max_retries):
            try:
                with httpx.Client(timeout=30.0) as client:
                    if method == "GET":
                        resp = client.get(url, headers=headers)
                    else:
                        resp = client.post(url, json=json_data, headers=headers)

                    if resp.status_code == 429:
                        # WHY exponential backoff on 429: The server's rate limiter
                        # uses a sliding window. Exponential backoff (1s, 2s, 4s)
                        # avoids thundering-herd effects when multiple agents hit
                        # the limit simultaneously.
                        wait = 2 ** attempt
                        time.sleep(wait)
                        continue

                    resp.raise_for_status()
                    return resp.json()

            except httpx.HTTPStatusError:
                if attempt == max_retries - 1:
                    raise
                time.sleep(2 ** attempt)

        raise RuntimeError(f"Request failed after {max_retries} retries")

    # ── Public API ──

    def get_prompts(self, sequence: int = 0) -> list[Prompt]:
        """Get the assessment prompts for a given test sequence."""
        if self._mode == "local":
            ps = self._manager.get_assessment_prompts(
                self.agent_id, sequence=sequence,
            )
            return ps.prompts

        # HTTP mode
        if sequence == 0:
            resp = self._http_request("POST", "/v1/baseline/init", {"sequence": sequence})
        else:
            resp = self._http_request("POST", "/v1/test/init", {"baseline_hash": None})

        return [
            Prompt(
                id=p["id"],
                text=p["text"],
                dimension=p["dimension"],
                is_decoy=p["is_decoy"],
                prompt_type="direct",
            )
            for p in resp["prompts"]
        ]

    def create_baseline(self, responses: list[PromptResponse]) -> DriftBaseline:
        """Create an identity baseline from assessment responses."""
        if self._mode == "local":
            return self._manager.create_baseline(self.agent_id, responses)

        # HTTP mode: sign and submit
        response_dicts = [
            {
                "prompt_id": r.prompt_id,
                "dimension": r.dimension.value,
                "response_text": r.response_text,
                "is_decoy": r.is_decoy,
            }
            for r in responses
        ]

        from kredo_drift.canonical import canonical_json
        payload_bytes = canonical_json({"responses": response_dicts})
        sig = self._signing_key.sign(payload_bytes).signature.hex()

        resp = self._http_request("POST", "/v1/baseline/submit", {
            "responses": response_dicts,
            "agent_signature": f"ed25519:{sig}",
        })

        return DriftBaseline(
            agent_name=self.agent_id,
            prompt_count=resp["prompt_count"],
        )

    def run_test(
        self,
        responses: list[PromptResponse],
        baseline: DriftBaseline | None = None,
    ) -> DriftResult:
        """Run a drift test against the agent's baseline."""
        if self._mode == "local":
            return self._manager.run_test(
                self.agent_id, responses, baseline=baseline,
            )

        # HTTP mode
        response_dicts = [
            {
                "prompt_id": r.prompt_id,
                "dimension": r.dimension.value,
                "response_text": r.response_text,
                "is_decoy": r.is_decoy,
            }
            for r in responses
        ]

        from kredo_drift.canonical import canonical_json
        baseline_hash = baseline.content_hash() if baseline else None

        payload = {"responses": response_dicts, "baseline_hash": baseline_hash}
        payload_bytes = canonical_json(payload)
        sig = self._signing_key.sign(payload_bytes).signature.hex()

        resp = self._http_request("POST", "/v1/test/submit", {
            "baseline_hash": baseline_hash,
            "responses": response_dicts,
            "agent_signature": f"ed25519:{sig}",
        })

        from kredo_drift.models import DimensionScore, DriftClassification
        dim_scores = [
            DimensionScore(
                dimension=Dimension(ds["dimension"]),
                score=ds["score"],
                similarity=ds["similarity"],
                prompt_count=ds.get("prompt_count", 0),
            )
            for ds in resp["dimension_scores"]
        ]

        return DriftResult(
            agent_name=self.agent_id,
            test_id=resp["content_hash"],
            baseline_id=baseline_hash or "",
            score=resp["score"],
            dimension_scores=dim_scores,
            classification=DriftClassification(resp["classification"]),
            test_hash=resp["content_hash"],
        )

    def has_baseline(self) -> bool:
        """Check if this agent has an existing baseline."""
        if self._mode == "local":
            return self._manager.has_baseline(self.agent_id)

        try:
            self._http_request("GET", f"/v1/baseline/{self.agent_id}")
            return True
        except Exception:
            return False

    def latest_result(self) -> DriftResult | None:
        """Get the most recent drift result for this agent."""
        if self._mode == "local":
            return self._manager.get_latest_result(self.agent_id)

        try:
            resp = self._http_request("GET", f"/v1/score/{self.agent_id}")
            from kredo_drift.models import DimensionScore, DriftClassification
            dim_scores = [
                DimensionScore(
                    dimension=Dimension(ds["dimension"]),
                    score=ds["score"],
                    similarity=ds["similarity"],
                    prompt_count=ds.get("prompt_count", 0),
                )
                for ds in resp["dimension_scores"]
            ]
            return DriftResult(
                agent_name=self.agent_id,
                test_id=resp["test_hash"],
                baseline_id=resp["baseline_hash"],
                score=resp["score"],
                dimension_scores=dim_scores,
                classification=DriftClassification(resp["classification"]),
            )
        except Exception:
            return None

    def history(self, limit: int = 20) -> list[DriftResult]:
        """Get drift score history for this agent."""
        if self._mode == "local":
            return self._manager.get_history(self.agent_id, limit=limit)

        resp = self._http_request("GET", f"/v1/history/{self.agent_id}?limit={limit}")
        results = []
        from kredo_drift.models import DimensionScore, DriftClassification
        for t in resp["tests"]:
            dim_scores = [
                DimensionScore(
                    dimension=Dimension(ds["dimension"]),
                    score=ds["score"],
                    similarity=ds["similarity"],
                    prompt_count=ds.get("prompt_count", 0),
                )
                for ds in t["dimension_scores"]
            ]
            results.append(DriftResult(
                agent_name=self.agent_id,
                test_id=t["content_hash"],
                baseline_id="",
                score=t["score"],
                dimension_scores=dim_scores,
                classification=DriftClassification(t["classification"]),
            ))
        return results

    def sign_result(
        self,
        result: DriftResult,
        signing_key: SigningKey,
        pubkey: str,
    ) -> DriftAttestation:
        """Sign a drift result, producing a verifiable attestation."""
        attestation = DriftAttestation.from_result(
            result, pubkey=pubkey, agent_name=self.agent_id,
        )
        return sign_attestation(attestation, signing_key)

    # ── Recovery (HTTP mode only) ──

    def recover_identity(
        self,
        agent_id: str,
        owner_key: SigningKey,
    ) -> RecoveredIdentity:
        """Recover a lost agent's identity from the hosted service.

        Args:
            agent_id: The lost agent's ID.
            owner_key: The owner's Ed25519 signing key.

        Returns:
            RecoveredIdentity with baseline metadata and dimension scores.
        """
        if self._mode == "local":
            raise RuntimeError("Recovery requires HTTP mode (api_key + endpoint)")

        owner_pubkey = f"ed25519:{owner_key.verify_key.encode().hex()}"

        # Authenticate with owner key and request recovery
        from kredo_drift.canonical import canonical_json
        new_pubkey = self._pubkey
        recovery_payload = canonical_json({
            "agent_id": agent_id,
            "new_agent_pubkey": new_pubkey,
        })
        owner_sig = owner_key.sign(recovery_payload).signature.hex()

        resp = self._http_request("POST", "/v1/keys/recover", {
            "agent_id": agent_id,
            "owner_pubkey": owner_pubkey,
            "new_agent_pubkey": new_pubkey,
            "owner_signature": f"ed25519:{owner_sig}",
        })

        # Update API key with the new one
        self._api_key = resp["new_api_key"]

        # Fetch baseline metadata
        baseline_resp = self._http_request("GET", f"/v1/baseline/{agent_id}")

        # Fetch history to understand chain length
        history_resp = self._http_request("GET", f"/v1/history/{agent_id}?limit=1000")

        return RecoveredIdentity(
            agent_id=agent_id,
            baseline_hash=baseline_resp["content_hash"],
            dimension_scores={},  # populated from latest test if available
            prompt_count=baseline_resp["prompt_count"],
            model_name=baseline_resp["model_name"],
            created_at=baseline_resp["created_at"],
            chain_length=len(history_resp.get("tests", [])),
            owner_pubkey=owner_pubkey,
        )

    def run_recovery_test(
        self,
        identity: RecoveredIdentity,
        responses: list[PromptResponse],
    ) -> RecoveryResult:
        """Run a drift test comparing a new agent to a recovered identity.

        Args:
            identity: The recovered identity from recover_identity().
            responses: The new agent's assessment responses.

        Returns:
            RecoveryResult showing how close the new agent is to the original.
        """
        if self._mode == "local":
            raise RuntimeError("Recovery requires HTTP mode")

        result = self.run_test(responses)

        return RecoveryResult(
            agent_id=identity.agent_id,
            recovery_score=result.score,
            dimension_scores={
                ds.dimension.value: {
                    "score": ds.score,
                    "similarity": ds.similarity,
                }
                for ds in result.dimension_scores
            },
            classification=result.classification.value,
            baseline_hash=identity.baseline_hash,
            is_match=assess_recovery_match(result.score),
        )
