"""alice thread — Persistent conversational threads with AgentCore Memory."""

from __future__ import annotations

from typing import Any

import click
from botocore.exceptions import ClientError
from prompt_toolkit import PromptSession
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from rich.markdown import Markdown

from alice_cli import personality
from alice_cli.agentcore_deps import require_agentcore
from alice_cli.aws.bedrock_context import create_bedrock_context
from alice_cli.aws.session import get_session
from alice_cli.console import (
    err_console,
    print_banner,
    print_rule,
    print_status,
    print_success,
    spinner,
)
from alice_cli.constants import DEFAULT_MAX_TOKENS, get_memory_bucket
from alice_cli.errors import AliceCLIError, LockerError, MemoryError
from alice_cli.locker import CloudLocker
from alice_cli.memory import MemoryClient
from alice_cli.models import (
    DEFAULT_MODEL_ALIAS,
    ChatMessage,
    CLIConfig,
    SessionRecord,
    TokenUsage,
    model_help_text,
    resolve_model_id,
)
from alice_cli.phrase_id import new_phrase_id
from alice_cli.session_record import build_session_record

# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------


def _setup_memory(
    memory: str,
    config: CLIConfig,
    actor_id: str = "default",
) -> tuple[MemoryClient | None, bool]:
    """Create a MemoryClient with graceful fallback.

    Returns (memory_client, memory_available).
    """
    try:
        boto3_session = get_session(profile=config.profile, region=config.region)
        mc = MemoryClient(
            memory_name=memory,
            region=config.region,
            actor_id=actor_id,
            boto3_session=boto3_session,
            quiet=config.quiet,
        )
        mc.check_status()
        return mc, True
    except (MemoryError, AliceCLIError) as exc:
        err_console.print(
            f"[yellow]⚠ Memory unavailable: {exc}. Using in-process history.[/yellow]"
        )
        return None, False


def _run_thread_repl(
    *,
    session_id: str,
    model_id: str,
    resolved_id: str,
    system: str | None,
    memory: str,
    memory_client: MemoryClient | None,
    memory_available: bool,
    bedrock_client: Any,
    locker: CloudLocker,
    messages: list[ChatMessage],
    config: CLIConfig,
) -> None:
    """Shared REPL loop for ``thread start`` and ``thread resume``."""
    total_input_tokens = 0
    total_output_tokens = 0

    print_rule(config.quiet)
    if not config.quiet:
        err_console.print(
            f"  [bold]Thread[/bold] with [cyan]{model_id}[/cyan]  •  "
            f"Type [dim]exit[/dim], [dim]quit[/dim], or [dim]Ctrl+D[/dim] to end.  "
            f"Type [dim]forget[/dim] to clear memory.\n"
        )

    prompt_session: PromptSession[str] = PromptSession(
        auto_suggest=AutoSuggestFromHistory(),
        multiline=False,
    )

    first_prompt: str | None = None

    while True:
        try:
            user_input = prompt_session.prompt("You: ")
        except (EOFError, KeyboardInterrupt):
            break

        stripped = user_input.strip()
        if not stripped:
            continue
        if stripped.lower() in ("exit", "quit"):
            break

        # ── Handle forget command ─────────────────────────────────────
        if stripped.lower() == "forget":
            if memory_client and memory_available:
                try:
                    memory_client.delete_events(session_id)
                except MemoryError as exc:
                    err_console.print(f"[yellow]⚠ Could not clear memory: {exc}[/yellow]")
            messages.clear()
            total_input_tokens = 0
            total_output_tokens = 0
            first_prompt = None
            err_console.print("  [green]✓[/green] Thread memory cleared.")
            continue

        if first_prompt is None:
            first_prompt = stripped

        # ── Store user STM event ──────────────────────────────────────
        if memory_client and memory_available:
            try:
                memory_client.store_event(session_id, "user", stripped)
            except MemoryError as exc:
                err_console.print(f"[yellow]⚠ Could not store to memory: {exc}[/yellow]")

        messages.append(ChatMessage(role="user", content=stripped))

        # ── Converse API ──────────────────────────────────────────────
        api_messages = [{"role": m.role, "content": [{"text": m.content}]} for m in messages]
        converse_kwargs: dict[str, Any] = {
            "modelId": resolved_id,
            "messages": api_messages,
            "inferenceConfig": {"maxTokens": DEFAULT_MAX_TOKENS},
        }
        if system:
            converse_kwargs["system"] = [{"text": system}]

        with spinner(personality.status_invoking_model(model_id), config.quiet):
            try:
                response = bedrock_client.converse(**converse_kwargs)
            except ClientError as exc:
                err_console.print(f"  [red]✗[/red] Bedrock API error: {exc}")
                messages.pop()
                continue

        output_msg = response.get("output", {}).get("message", {})
        content_blocks = output_msg.get("content", [])
        text_parts = [b["text"] for b in content_blocks if "text" in b]
        assistant_text = "\n".join(text_parts)

        usage = response.get("usage", {})
        total_input_tokens += usage.get("inputTokens", 0)
        total_output_tokens += usage.get("outputTokens", 0)

        # ── Store assistant STM event ─────────────────────────────────
        if memory_client and memory_available:
            try:
                memory_client.store_event(session_id, "assistant", assistant_text)
            except MemoryError as exc:
                err_console.print(f"[yellow]⚠ Could not store to memory: {exc}[/yellow]")

        messages.append(ChatMessage(role="assistant", content=assistant_text))

        # ── Display ───────────────────────────────────────────────────
        print_rule(config.quiet)
        if assistant_text.strip():
            err_console.print(Markdown(assistant_text))
        else:
            err_console.print("[dim](empty response)[/dim]")
        print_rule(config.quiet)

    # ── Persist to Cloud Locker ───────────────────────────────────────
    if messages:
        record = build_session_record(
            record_type="thread",
            model=resolved_id,
            prompt=first_prompt or "",
            response=None,
            tokens=TokenUsage(
                input=total_input_tokens,
                output=total_output_tokens,
                total=total_input_tokens + total_output_tokens,
            ),
            metadata={
                "messages": [{"role": m.role, "content": m.content} for m in messages],
                "system": system,
                "session_id": session_id,
                "memory_name": memory,
            },
        )
        record = SessionRecord(
            id=session_id,
            type=record.type,
            timestamp=record.timestamp,
            model=record.model,
            prompt=record.prompt,
            response=record.response,
            tokens=record.tokens,
            metadata=record.metadata,
        )
        try:
            locker.persist(record, "thread")
            print_success(f"Session saved: {session_id}", config.quiet)
        except Exception as exc:
            err_console.print(f"[yellow]⚠ Could not save to Cloud Locker: {exc}[/yellow]")
    else:
        if not config.quiet:
            err_console.print("  [dim]No messages to save.[/dim]")

    if not config.quiet:
        err_console.print(f"\n  [green]✓[/green] Thread ended. Session ID: {session_id}")


# ---------------------------------------------------------------------------
# Click group and subcommands
# ---------------------------------------------------------------------------


@click.group("thread", invoke_without_command=True)
@click.option("--memory", default="alice_memory", help="AgentCore Memory resource name")
@click.pass_context
def thread_group(ctx: click.Context, memory: str) -> None:
    """Start persistent conversations with AgentCore Memory.

    When invoked without a subcommand, lists all previous threads.
    """
    if ctx.invoked_subcommand is not None:
        return

    require_agentcore()
    config = ctx.obj["config"]

    memory_client, memory_available = _setup_memory(memory, config)
    if not memory_client or not memory_available:
        raise AliceCLIError("Cannot list threads: memory is unavailable.")

    sessions = memory_client.list_sessions()
    if not sessions:
        err_console.print("  [dim]No threads found.[/dim]")
        return

    from rich.table import Table

    table = Table(title="Threads", show_lines=False)
    table.add_column("Session ID", style="cyan", no_wrap=True)
    table.add_column("Created", style="green")
    table.add_column("Messages", style="magenta", justify="right")

    for s in sessions:
        table.add_row(
            str(s.get("session_id", "")),
            str(s.get("created_at", "")),
            str(s.get("message_count", 0)),
        )

    err_console.print(table)
    print_success(f"{len(sessions)} thread(s) found.", config.quiet)


@thread_group.command("start")
@click.option(
    "--model",
    "--model-id",
    "model_id",
    default=DEFAULT_MODEL_ALIAS,
    help=model_help_text(),
)
@click.option("--memory", default="alice_memory", help="AgentCore Memory resource name")
@click.option("--system", default=None, help="System prompt for the conversation")
@click.pass_context
def thread_start(
    ctx: click.Context,
    model_id: str,
    memory: str,
    system: str | None,
) -> None:
    """Start a new persistent conversation thread.

    \b
    Examples:
        alice thread start
        alice thread start --model opus
        alice thread start --system "You are a helpful tutor"
        alice thread start --memory my-memory
    """
    require_agentcore()
    config = ctx.obj["config"]
    print_banner(personality.BANNER, config.quiet)

    ctx_bedrock = create_bedrock_context(config)
    username = ctx_bedrock.username
    bedrock_client = ctx_bedrock.client
    resolved_id = resolve_model_id(model_id, ctx_bedrock.inference_profile_arns)
    memory_client, memory_available = _setup_memory(memory, config, actor_id=username)
    locker = CloudLocker(
        config=config,
        username=username,
        bucket=get_memory_bucket(config.namespace, config.environment),
    )

    session_id = new_phrase_id()
    print_status(f"New thread: {session_id}", config.quiet)

    _run_thread_repl(
        session_id=session_id,
        model_id=model_id,
        resolved_id=resolved_id,
        system=system,
        memory=memory,
        memory_client=memory_client,
        memory_available=memory_available,
        bedrock_client=bedrock_client,
        locker=locker,
        messages=[],
        config=config,
    )


@thread_group.command("resume")
@click.argument("session_id")
@click.option(
    "--model",
    "--model-id",
    "model_id",
    default=DEFAULT_MODEL_ALIAS,
    help=model_help_text(),
)
@click.option("--memory", default="alice_memory", help="AgentCore Memory resource name")
@click.option("--system", default=None, help="System prompt for the conversation")
@click.pass_context
def thread_resume(
    ctx: click.Context,
    session_id: str,
    model_id: str,
    memory: str,
    system: str | None,
) -> None:
    """Resume a previous conversation thread.

    \b
    Examples:
        alice thread resume <session-id>
        alice thread resume <session-id> --model opus
    """
    require_agentcore()
    config = ctx.obj["config"]
    print_banner(personality.BANNER, config.quiet)

    ctx_bedrock = create_bedrock_context(config)
    username = ctx_bedrock.username
    bedrock_client = ctx_bedrock.client
    resolved_id = resolve_model_id(model_id, ctx_bedrock.inference_profile_arns)
    memory_client, memory_available = _setup_memory(memory, config, actor_id=username)
    locker = CloudLocker(
        config=config,
        username=username,
        bucket=get_memory_bucket(config.namespace, config.environment),
    )

    # ── Restore previous conversation from STM ────────────────────────
    messages: list[ChatMessage] = []
    if memory_client and memory_available:
        try:
            events = memory_client.get_events(session_id)
            for ev in events:
                messages.append(ChatMessage(role=ev["role"], content=ev["content"]))
            if messages:
                print_success(
                    f"Resumed thread {session_id} ({len(messages)} messages)", config.quiet
                )
            else:
                print_status(f"Resuming thread {session_id} (no previous messages)", config.quiet)
        except Exception as exc:
            err_console.print(f"[yellow]⚠ Could not load history: {exc}[/yellow]")
    else:
        print_status(f"Resuming thread {session_id} (memory unavailable)", config.quiet)

    _run_thread_repl(
        session_id=session_id,
        model_id=model_id,
        resolved_id=resolved_id,
        system=system,
        memory=memory,
        memory_client=memory_client,
        memory_available=memory_available,
        bedrock_client=bedrock_client,
        locker=locker,
        messages=messages,
        config=config,
    )


@thread_group.command("delete")
@click.argument("session_id")
@click.option("--memory", default="alice_memory", help="AgentCore Memory resource name")
@click.pass_context
def thread_delete(
    ctx: click.Context,
    session_id: str,
    memory: str,
) -> None:
    """Delete a thread and its memory.

    \b
    Examples:
        alice thread delete <session-id>
    """
    require_agentcore()
    config = ctx.obj["config"]

    if not click.confirm(f"Delete thread {session_id}?"):
        err_console.print("  [dim]Cancelled.[/dim]")
        return

    memory_client, memory_available = _setup_memory(memory, config)
    if not memory_client or not memory_available:
        raise AliceCLIError("Cannot delete thread: memory is unavailable.")

    memory_client.delete_events(session_id)
    print_success(f"Thread {session_id} deleted.", config.quiet)


@thread_group.command("remember")
@click.argument("text")
@click.option("--memory", default="alice_memory", help="AgentCore Memory resource name")
@click.pass_context
def thread_remember(
    ctx: click.Context,
    text: str,
    memory: str,
) -> None:
    """Store a fact or knowledge item in long-term memory.

    \b
    Examples:
        alice thread remember "I prefer APA citation format"
        alice thread remember "My research focuses on NLP" --memory my-memory
    """
    require_agentcore()
    config = ctx.obj["config"]

    memory_client, memory_available = _setup_memory(memory, config)
    if not memory_client or not memory_available:
        raise AliceCLIError("Cannot store fact: memory is unavailable.")

    try:
        memory_client.store_fact(text)
    except MemoryError:
        raise
    print_success(f"Fact stored in long-term memory ({memory}).", config.quiet)


@thread_group.command("recall")
@click.argument("query")
@click.option("--memory", default="alice_memory", help="AgentCore Memory resource name")
@click.pass_context
def thread_recall(
    ctx: click.Context,
    query: str,
    memory: str,
) -> None:
    """Search long-term memory for facts matching a query.

    \b
    Examples:
        alice thread recall "citation format"
        alice thread recall "NLP research" --memory my-memory
    """
    require_agentcore()
    config = ctx.obj["config"]

    memory_client, memory_available = _setup_memory(memory, config)
    if not memory_client or not memory_available:
        raise AliceCLIError("Cannot search facts: memory is unavailable.")

    results = memory_client.search_facts(query)
    if not results:
        err_console.print("  [dim]No matching facts found.[/dim]")
        return

    from rich.table import Table

    table = Table(title=f"Recall: {query}", show_lines=True)
    table.add_column("#", style="dim", width=4, justify="right")
    table.add_column("Fact", style="cyan")
    table.add_column("Score", style="magenta", width=8, justify="right")

    for i, fact in enumerate(results, 1):
        text = fact.get("text", str(fact))
        score = fact.get("score", fact.get("relevance", ""))
        score_str = f"{score:.3f}" if isinstance(score, int | float) else str(score)
        table.add_row(str(i), text, score_str)

    err_console.print(table)
    print_success(f"{len(results)} fact(s) found.", config.quiet)


@thread_group.command("forget-all")
@click.option("--memory", default="alice_memory", help="AgentCore Memory resource name")
@click.pass_context
def thread_forget_all(
    ctx: click.Context,
    memory: str,
) -> None:
    """Clear all long-term memory facts from the configured Memory Resource.

    \b
    Examples:
        alice thread forget-all
        alice thread forget-all --memory my-memory
    """
    require_agentcore()
    config = ctx.obj["config"]

    if not click.confirm("Clear all long-term memory facts?"):
        err_console.print("  [dim]Cancelled.[/dim]")
        return

    memory_client, memory_available = _setup_memory(memory, config)
    if not memory_client or not memory_available:
        raise AliceCLIError("Cannot clear facts: memory is unavailable.")

    memory_client.clear_facts()
    print_success(f"All LTM facts cleared ({memory}).", config.quiet)


@thread_group.command("import")
@click.argument("session_id")
@click.option("--memory", default="alice_memory", help="AgentCore Memory resource name")
@click.option("--category", default="chat", help="Cloud Locker category to import from")
@click.pass_context
def thread_import(
    ctx: click.Context,
    session_id: str,
    memory: str,
    category: str,
) -> None:
    """Import a Cloud Locker session as a new thread.

    \b
    Examples:
        alice thread import <session-id>
        alice thread import <session-id> --category invoke
        alice thread import <session-id> --memory my-memory
    """
    require_agentcore()
    config = ctx.obj["config"]

    ctx_bedrock = create_bedrock_context(config)
    username = ctx_bedrock.username
    locker = CloudLocker(
        config=config,
        username=username,
        bucket=get_memory_bucket(config.namespace, config.environment),
    )

    # Load the session from Cloud Locker
    try:
        record = locker.load_session(session_id, category)
    except (LockerError, AliceCLIError) as exc:
        raise AliceCLIError(f"Session {session_id} not found: {exc}") from exc

    # Extract messages from metadata
    messages = (record.metadata or {}).get("messages", [])
    if not messages:
        err_console.print("[yellow]⚠ Session has no conversation history to import.[/yellow]")
        return

    # Set up memory client
    memory_client, memory_available = _setup_memory(memory, config, actor_id=username)
    if not memory_client or not memory_available:
        raise AliceCLIError("Cannot import thread: memory is unavailable.")

    # Create new thread with a fresh session ID
    new_session_id = new_phrase_id()

    # Write each message as an STM event
    for msg in messages:
        memory_client.store_event(new_session_id, msg["role"], msg["content"])

    print_success(
        f"Imported {len(messages)} messages into new thread: {new_session_id}", config.quiet
    )
