"""CLI interface for Kredo Drift.

Local mode commands:
  drift baseline <agent_name>  — Create an identity baseline
  drift test <agent_name>      — Run a drift test against baseline
  drift score <agent_name>     — Show latest drift score
  drift history <agent_name>   — Show drift score history
  drift agents                 — List all agents with baselines
  drift prompts                — Show prompt pool statistics

Hosted mode commands (requires login):
  drift login                  — Authenticate with API key
  drift register <agent_name>  — Register agent with hosted service
  drift recover <agent_name>   — Recover a lost agent's identity
"""

from __future__ import annotations

import json
import sys
from datetime import datetime, timezone
from pathlib import Path

import click

from kredo_drift.assessment import AssessmentProtocol
from kredo_drift.baseline import BaselineManager
from kredo_drift.models import Dimension, PromptResponse
from kredo_drift.store import DriftStore
from kredo_drift.vectorize import Vectorizer


def _get_manager(db_path: str | None = None) -> BaselineManager:
    """Create a BaselineManager with default or custom DB path."""
    store = DriftStore(db_path=db_path) if db_path else DriftStore()
    return BaselineManager(store=store)


def _collect_responses(
    prompts,
    agent_name: str,
    respond_fn=None,
) -> list[PromptResponse]:
    """Collect responses to assessment prompts.

    If respond_fn is provided, use it to generate responses programmatically.
    Otherwise, prompt the user interactively via stdin.
    """
    responses: list[PromptResponse] = []

    for i, prompt in enumerate(prompts, 1):
        dim_label = prompt.dimension.upper()
        decoy_tag = " [decoy]" if prompt.is_decoy else ""

        if respond_fn:
            response_text = respond_fn(prompt.text)
        else:
            click.echo()
            click.echo(click.style(f"[{i}/{len(prompts)}] ", fg="cyan") +
                       click.style(f"({dim_label}{decoy_tag})", fg="yellow"))
            click.echo(click.style(prompt.text, fg="white", bold=True))
            click.echo()
            response_text = click.prompt(
                click.style("Response", fg="green"),
                prompt_suffix=": ",
            )

        # Map dimension string to Dimension enum for PromptResponse.
        # Decoys use PERSONALITY as a placeholder (they're excluded from scoring).
        dim_enum = (
            Dimension(prompt.dimension)
            if not prompt.is_decoy
            else Dimension.PERSONALITY
        )

        responses.append(PromptResponse(
            prompt_id=prompt.id,
            dimension=dim_enum,
            prompt_text=prompt.text,
            response_text=response_text,
            is_decoy=prompt.is_decoy,
        ))

    return responses


_CONFIG_DIR = Path.home() / ".config" / "kredo-drift"
_CREDENTIALS_FILE = _CONFIG_DIR / "credentials.json"


def _load_credentials() -> dict | None:
    """Load stored API credentials."""
    if _CREDENTIALS_FILE.exists():
        return json.loads(_CREDENTIALS_FILE.read_text())
    return None


def _save_credentials(api_key: str, endpoint: str):
    """Save API credentials."""
    _CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    _CREDENTIALS_FILE.write_text(json.dumps({
        "api_key": api_key,
        "endpoint": endpoint,
    }))
    _CREDENTIALS_FILE.chmod(0o600)


def _get_client(agent_name: str, db_path: str | None = None) -> "DriftClient":
    """Create a DriftClient, using hosted mode if logged in."""
    from kredo_drift.client import DriftClient
    creds = _load_credentials()
    if creds:
        return DriftClient(
            agent_name,
            api_key=creds["api_key"],
            endpoint=creds["endpoint"],
        )
    return DriftClient(agent_name, db_path=db_path)


@click.group()
@click.option("--db", default=None, help="Path to SQLite database file.")
@click.pass_context
def cli(ctx, db):
    """Kredo Drift — Continuous identity drift detection for AI agents."""
    ctx.ensure_object(dict)
    ctx.obj["db_path"] = db


@cli.command()
@click.argument("agent_name")
@click.pass_context
def baseline(ctx, agent_name: str):
    """Create an identity baseline for an agent.

    Presents 30 assessment prompts and captures the agent's responses.
    The responses are vectorized into a multi-dimensional identity fingerprint.
    """
    manager = _get_manager(ctx.obj.get("db_path"))

    # Check if baseline already exists
    if manager.has_baseline(agent_name):
        if not click.confirm(
            click.style(
                f"A baseline already exists for '{agent_name}'. "
                "Creating a new one will replace it. Continue?",
                fg="yellow",
            )
        ):
            click.echo("Aborted.")
            return

    click.echo()
    click.echo(click.style("=" * 60, fg="blue"))
    click.echo(click.style("  Kredo Drift — Baseline Capture", fg="blue", bold=True))
    click.echo(click.style(f"  Agent: {agent_name}", fg="blue"))
    click.echo(click.style("=" * 60, fg="blue"))
    click.echo()
    click.echo("Answer each prompt as yourself. There are no right answers.")
    click.echo("Your responses will be vectorized — raw text is not stored.")
    click.echo()

    # Get prompts
    prompt_set = manager.get_assessment_prompts(agent_name, sequence=0)

    click.echo(f"Assessment: {len(prompt_set.prompts)} prompts across 5 dimensions")
    click.echo()

    # Collect responses interactively
    responses = _collect_responses(prompt_set.prompts, agent_name)

    click.echo()
    click.echo(click.style("Vectorizing responses...", fg="cyan"))

    # Create baseline
    bl = manager.create_baseline(agent_name, responses)

    click.echo()
    click.echo(click.style("Baseline created.", fg="green", bold=True))
    click.echo(f"  ID:      {bl.id}")
    click.echo(f"  Agent:   {bl.agent_name}")
    click.echo(f"  Prompts: {bl.prompt_count}")
    click.echo(f"  Model:   {bl.model_name}")
    click.echo(f"  Hash:    {bl.content_hash()[:16]}...")
    click.echo(f"  Stored:  {manager.store.db_path}")
    click.echo()


@cli.command()
@click.argument("agent_name")
@click.pass_context
def test(ctx, agent_name: str):
    """Run a drift test against an agent's baseline.

    Presents a rotated set of assessment prompts and compares the
    responses to the stored baseline.
    """
    manager = _get_manager(ctx.obj.get("db_path"))

    # Check for baseline
    if not manager.has_baseline(agent_name):
        click.echo(
            click.style(
                f"No baseline found for '{agent_name}'. "
                "Run 'drift baseline' first.",
                fg="red",
            )
        )
        sys.exit(1)

    # Determine test sequence number
    history = manager.get_history(agent_name, limit=1000)
    sequence = len(history) + 1

    click.echo()
    click.echo(click.style("=" * 60, fg="blue"))
    click.echo(click.style("  Kredo Drift — Drift Test", fg="blue", bold=True))
    click.echo(click.style(f"  Agent: {agent_name}  |  Test #{sequence}", fg="blue"))
    click.echo(click.style("=" * 60, fg="blue"))
    click.echo()

    # Get rotated prompts
    prompt_set = manager.get_assessment_prompts(agent_name, sequence=sequence)

    click.echo(f"Assessment: {len(prompt_set.prompts)} prompts (rotated from pool)")
    click.echo()

    # Collect responses interactively
    responses = _collect_responses(prompt_set.prompts, agent_name)

    click.echo()
    click.echo(click.style("Vectorizing and scoring...", fg="cyan"))

    # Run test
    result = manager.run_test(agent_name, responses)

    # Display result
    _display_result(result)


@cli.command()
@click.argument("agent_name")
@click.pass_context
def score(ctx, agent_name: str):
    """Show the latest drift score for an agent."""
    manager = _get_manager(ctx.obj.get("db_path"))

    result = manager.get_latest_result(agent_name)
    if result is None:
        click.echo(
            click.style(
                f"No drift tests found for '{agent_name}'. "
                "Run 'drift test' first.",
                fg="red",
            )
        )
        sys.exit(1)

    _display_result(result)


@cli.command()
@click.argument("agent_name")
@click.option("--limit", "-n", default=20, help="Number of results to show.")
@click.pass_context
def history(ctx, agent_name: str, limit: int):
    """Show drift score history for an agent."""
    manager = _get_manager(ctx.obj.get("db_path"))

    results = manager.get_history(agent_name, limit=limit)
    if not results:
        click.echo(
            click.style(
                f"No drift history for '{agent_name}'.",
                fg="red",
            )
        )
        sys.exit(1)

    click.echo()
    click.echo(click.style(f"Drift History — {agent_name}", fg="blue", bold=True))
    click.echo(click.style("-" * 70, fg="blue"))
    click.echo(
        f"{'#':<4} {'Date':<20} {'Score':>6} {'Severity':<15} {'Classification'}"
    )
    click.echo(click.style("-" * 70, fg="blue"))

    for i, result in enumerate(reversed(results), 1):
        date_str = result.created_at.strftime("%Y-%m-%d %H:%M")
        severity = result.severity
        severity_color = _severity_color(result.score)
        click.echo(
            f"{i:<4} {date_str:<20} "
            + click.style(f"{result.score:>5.1f}", fg=severity_color)
            + f"  {severity:<15} {result.classification.value}"
        )

    click.echo(click.style("-" * 70, fg="blue"))
    click.echo(f"Showing {len(results)} of {len(results)} results")
    click.echo()


@cli.command()
@click.pass_context
def agents(ctx):
    """List all agents with baselines."""
    manager = _get_manager(ctx.obj.get("db_path"))

    agent_list = manager.store.list_agents()
    if not agent_list:
        click.echo("No agents found. Run 'drift baseline <name>' to create one.")
        return

    click.echo()
    click.echo(click.style("Registered Agents", fg="blue", bold=True))
    click.echo(click.style("-" * 50, fg="blue"))

    for name in agent_list:
        bl = manager.store.load_baseline(name)
        latest = manager.get_latest_result(name)
        history_count = len(manager.get_history(name, limit=1000))

        line = f"  {name}"
        if bl:
            line += f"  (baseline: {bl.created_at.strftime('%Y-%m-%d')})"
        if latest:
            color = _severity_color(latest.score)
            line += f"  [drift: " + click.style(f"{latest.score:.1f}", fg=color) + f"  {latest.severity}]"
        line += f"  ({history_count} tests)"
        click.echo(line)

    click.echo()


@cli.command()
def prompts():
    """Show prompt pool statistics."""
    protocol = AssessmentProtocol()
    stats = protocol.pool_stats()

    click.echo()
    click.echo(click.style("Prompt Pool Statistics", fg="blue", bold=True))
    click.echo(click.style("-" * 40, fg="blue"))

    for dim in Dimension:
        count = stats.get(dim.value, 0)
        click.echo(f"  {dim.value:<15} {count:>3} prompts")

    click.echo(f"  {'decoys':<15} {stats.get('decoys', 0):>3} prompts")
    click.echo(click.style("-" * 40, fg="blue"))
    click.echo(f"  {'TOTAL':<15} {stats.get('total', 0):>3} prompts")
    click.echo()

    # Show a sample prompt set
    ps = protocol.select_prompts("sample-agent", sequence=0)
    scored = ps.scored_prompts()
    click.echo(f"  Per assessment: {len(ps.prompts)} prompts ({len(scored)} scored + {len(ps.prompts) - len(scored)} decoys)")
    click.echo()


@cli.command()
@click.argument("agent_name")
@click.pass_context
def attest(ctx, agent_name: str):
    """Sign the latest drift result with an Ed25519 key.

    Generates a new keypair and produces a verifiable attestation.
    """
    from kredo_drift.models import DriftAttestation
    from kredo_drift.signing import generate_keypair, sign_attestation, verify_attestation

    manager = _get_manager(ctx.obj.get("db_path"))

    result = manager.get_latest_result(agent_name)
    if result is None:
        click.echo(
            click.style(
                f"No drift results for '{agent_name}'. Run 'drift test' first.",
                fg="red",
            )
        )
        sys.exit(1)

    # Generate a fresh keypair
    signing_key, pubkey = generate_keypair()

    # Build and sign attestation
    attestation = DriftAttestation.from_result(result, pubkey=pubkey)
    signed = sign_attestation(attestation, signing_key)

    # Verify it
    valid = verify_attestation(signed)

    click.echo()
    click.echo(click.style("Drift Attestation", fg="blue", bold=True))
    click.echo(click.style("-" * 60, fg="blue"))
    click.echo(f"  Agent:          {agent_name}")
    click.echo(f"  Score:          {result.score:.1f} ({result.severity})")
    click.echo(f"  Classification: {result.classification.value}")
    click.echo(f"  Public Key:     {pubkey}")
    click.echo(f"  Signature:      {signed.signature[:40]}...")
    click.echo(f"  Verified:       {click.style('Yes', fg='green') if valid else click.style('No', fg='red')}")
    click.echo()

    # Output the full attestation as JSON
    import json
    click.echo(click.style("Attestation JSON:", fg="blue"))
    click.echo(json.dumps(signed.model_dump(mode="json"), indent=2, default=str))
    click.echo()


@cli.command()
@click.option("--api-key", prompt=True, hide_input=True, help="Your Kredo Drift API key.")
@click.option("--endpoint", default="https://api.aikredo.com/drift", help="API endpoint URL.")
def login(api_key: str, endpoint: str):
    """Authenticate with the Kredo Drift hosted service.

    After login, all commands (baseline, test, score, etc.) will use
    the hosted API instead of local SQLite.
    """
    _save_credentials(api_key, endpoint)
    click.echo()
    click.echo(click.style("Logged in.", fg="green", bold=True))
    click.echo(f"  Endpoint: {endpoint}")
    click.echo(f"  Credentials saved to: {_CREDENTIALS_FILE}")
    click.echo()
    click.echo("All drift commands will now use the hosted API.")
    click.echo("To revert to local mode, delete the credentials file.")
    click.echo()


@cli.command()
@click.argument("agent_name")
@click.pass_context
def register(ctx, agent_name: str):
    """Register an agent with the hosted Kredo Drift service.

    Requires login first. Generates an Ed25519 keypair and completes
    a challenge-response registration.
    """
    creds = _load_credentials()
    if not creds:
        click.echo(click.style("Not logged in. Run 'drift login' first.", fg="red"))
        sys.exit(1)

    from kredo_drift.client import _load_or_generate_key
    from kredo_drift.canonical import canonical_json
    import httpx

    signing_key, pubkey = _load_or_generate_key(agent_name)

    click.echo(f"Registering agent '{agent_name}'...")
    click.echo(f"  Public key: {pubkey}")

    endpoint = creds["endpoint"].rstrip("/")
    headers = {"X-API-Key": creds["api_key"]}

    # Step 1: Start registration
    with httpx.Client(timeout=30.0) as http:
        resp = http.post(
            f"{endpoint}/v1/register",
            json={"agent_id": agent_name, "name": agent_name, "pubkey": pubkey},
            headers=headers,
        )
        if resp.status_code != 200:
            click.echo(click.style(f"Registration failed: {resp.text}", fg="red"))
            sys.exit(1)

        data = resp.json()
        challenge = data["challenge"]
        challenge_id = data["challenge_id"]

        # Step 2: Sign challenge
        sig = signing_key.sign(challenge.encode()).signature.hex()

        resp = http.post(
            f"{endpoint}/v1/register/complete",
            json={
                "challenge_id": challenge_id,
                "signature": f"ed25519:{sig}",
            },
            headers=headers,
        )
        if resp.status_code != 200:
            click.echo(click.style(f"Registration failed: {resp.text}", fg="red"))
            sys.exit(1)

        result = resp.json()

    # Save the new API key
    _save_credentials(result["api_key"], creds["endpoint"])

    click.echo()
    click.echo(click.style("Agent registered.", fg="green", bold=True))
    click.echo(f"  Agent ID:      {result['agent_id']}")
    click.echo(f"  Service Key:   {result['service_pubkey']}")
    click.echo(f"  API key updated in credentials.")
    click.echo()


@cli.command()
@click.argument("agent_name")
@click.pass_context
def recover(ctx, agent_name: str):
    """Recover a lost agent's identity from the hosted service.

    Requires the owner's Ed25519 key. The new agent will be assessed
    against the original baseline to verify identity continuity.
    """
    creds = _load_credentials()
    if not creds:
        click.echo(click.style("Not logged in. Run 'drift login' first.", fg="red"))
        sys.exit(1)

    owner_key_path = click.prompt(
        "Path to owner signing key",
        type=click.Path(exists=True),
    )

    from nacl.signing import SigningKey
    from kredo_drift.client import DriftClient

    owner_key = SigningKey(Path(owner_key_path).read_bytes())

    client = DriftClient(
        agent_name,
        api_key=creds["api_key"],
        endpoint=creds["endpoint"],
    )

    click.echo(f"Recovering identity for '{agent_name}'...")
    identity = client.recover_identity(agent_name, owner_key=owner_key)

    click.echo()
    click.echo(click.style("Identity recovered.", fg="green", bold=True))
    click.echo(f"  Agent:       {identity.agent_id}")
    click.echo(f"  Baseline:    {identity.baseline_hash[:16]}...")
    click.echo(f"  Chain:       {identity.chain_length} events")
    click.echo(f"  Created:     {identity.created_at}")
    click.echo()
    click.echo("Run 'drift test' to assess the new agent against the original baseline.")
    click.echo()


def _display_result(result):
    """Display a drift result with formatting."""
    click.echo()
    click.echo(click.style("=" * 60, fg="blue"))
    click.echo(click.style("  Drift Score Report", fg="blue", bold=True))
    click.echo(click.style("=" * 60, fg="blue"))
    click.echo()

    # Overall score
    severity_color = _severity_color(result.score)
    click.echo(
        "  Overall:  "
        + click.style(f"{result.score:.1f}", fg=severity_color, bold=True)
        + click.style(f" / 100", fg="white")
        + "  ("
        + click.style(result.severity, fg=severity_color)
        + ")"
    )
    click.echo(f"  Classification: {result.classification.value}")
    click.echo()

    # Per-dimension breakdown
    click.echo(click.style("  Dimension Breakdown:", fg="blue"))
    click.echo(click.style("  " + "-" * 50, fg="blue"))

    for ds in sorted(result.dimension_scores, key=lambda d: d.score, reverse=True):
        bar_len = int(ds.score / 2)  # 50 chars max
        bar = "#" * bar_len
        dim_color = _severity_color(ds.score)
        click.echo(
            f"  {ds.dimension.value:<13} "
            + click.style(f"{ds.score:>5.1f}", fg=dim_color)
            + f"  (sim: {ds.similarity:>6.3f})  "
            + click.style(bar, fg=dim_color)
        )

    click.echo()
    click.echo(f"  Agent:    {result.agent_name}")
    click.echo(f"  Test ID:  {result.test_id[:16]}...")
    click.echo(f"  Date:     {result.created_at.strftime('%Y-%m-%d %H:%M:%S UTC')}")
    click.echo()

    # Score legend
    click.echo(click.style("  Score Guide:", dim=True))
    click.echo(click.style("   0-15 Stable | 16-35 Shifting | 36-60 Drifting | 61-85 Diverged | 86-100 Unrecognizable", dim=True))
    click.echo()


def _severity_color(score: float) -> str:
    """Return a click color name based on drift severity."""
    if score <= 15:
        return "green"
    elif score <= 35:
        return "yellow"
    elif score <= 60:
        return "red"
    else:
        return "bright_red"


def main():
    """Entry point for the CLI."""
    cli()


if __name__ == "__main__":
    main()
