"""alice batch — Batch prompt processing from a file."""

from __future__ import annotations

import json
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import click
from botocore.exceptions import ClientError
from rich.progress import Progress

from alice_cli import personality
from alice_cli.aws.bedrock_context import create_bedrock_context
from alice_cli.bedrock_utils import build_converse_kwargs, parse_converse_response
from alice_cli.console import (
    err_console,
    print_banner,
    print_rule,
    print_status,
    print_success,
)
from alice_cli.constants import get_memory_bucket
from alice_cli.errors import AliceCLIError
from alice_cli.locker import CloudLocker
from alice_cli.models import (
    DEFAULT_MODEL_ALIAS,
    TokenUsage,
    model_help_text,
    resolve_model_id,
)
from alice_cli.session_record import build_session_record

# ── Prompt file parsing ───────────────────────────────────────────────


def parse_prompt_file(path: str) -> list[str]:
    """Parse a prompt file into a list of prompt strings.

    Supported formats:
    - .txt: one prompt per line (blank lines skipped)
    - .json: JSON array of strings
    - .jsonl: one JSON object per line with a "prompt" field
    """
    p = Path(path)
    suffix = p.suffix.lower()
    text = p.read_text(encoding="utf-8")

    if suffix == ".json":
        return _parse_json(text, path)
    elif suffix == ".jsonl":
        return _parse_jsonl(text, path)
    else:
        # Default: plain text, one prompt per line
        return _parse_txt(text)


def _parse_txt(text: str) -> list[str]:
    """Parse plain text: one prompt per non-blank line."""
    return [line for line in text.splitlines() if line.strip()]


def _parse_json(text: str, path: str) -> list[str]:
    """Parse a JSON array of prompt strings."""
    try:
        data = json.loads(text)
    except json.JSONDecodeError as exc:
        raise AliceCLIError(f"Invalid JSON in {path}: {exc}") from exc
    if not isinstance(data, list):
        raise AliceCLIError(f"Expected a JSON array in {path}, got {type(data).__name__}")
    prompts: list[str] = []
    for i, item in enumerate(data):
        if not isinstance(item, str):
            raise AliceCLIError(f"Item {i} in {path} is not a string: {type(item).__name__}")
        if item.strip():
            prompts.append(item)
    return prompts


def _parse_jsonl(text: str, path: str) -> list[str]:
    """Parse JSONL: one JSON object per line with a 'prompt' field."""
    prompts: list[str] = []
    for lineno, line in enumerate(text.splitlines(), start=1):
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except json.JSONDecodeError as exc:
            raise AliceCLIError(f"Invalid JSON on line {lineno} of {path}: {exc}") from exc
        if not isinstance(obj, dict) or "prompt" not in obj:
            raise AliceCLIError(
                f"Line {lineno} of {path} must be a JSON object with a 'prompt' field"
            )
        prompt = obj["prompt"]
        if isinstance(prompt, str) and prompt.strip():
            prompts.append(prompt)
    return prompts


# ── Serialization helpers (for round-trip property) ───────────────────


def serialize_prompts_txt(prompts: list[str]) -> str:
    """Serialize prompts to plain text format (one per line)."""
    return "\n".join(prompts)


def serialize_prompts_json(prompts: list[str]) -> str:
    """Serialize prompts to a JSON array."""
    return json.dumps(prompts, ensure_ascii=False)


def serialize_prompts_jsonl(prompts: list[str]) -> str:
    """Serialize prompts to JSONL format."""
    return "\n".join(json.dumps({"prompt": p}, ensure_ascii=False) for p in prompts)


# ── Single prompt invocation ──────────────────────────────────────────


@dataclass
class _PromptResult:
    """Result from a single prompt invocation."""

    prompt: str
    response: str | None
    model: str
    input_tokens: int
    output_tokens: int
    timestamp: str
    error: str | None = None


def _invoke_prompt(
    client: Any,
    resolved_id: str,
    model_alias: str,
    prompt_text: str,
    system_prompt: str | None,
) -> _PromptResult:
    """Invoke the Converse API for a single prompt. Returns a _PromptResult."""
    ts = datetime.now(UTC).isoformat()
    try:
        response = client.converse(
            **build_converse_kwargs(
                resolved_id,
                [{"role": "user", "content": [{"text": prompt_text}]}],
                system_prompt=system_prompt,
            )
        )
        text, tokens = parse_converse_response(response)
        return _PromptResult(
            prompt=prompt_text,
            response=text,
            model=model_alias,
            input_tokens=tokens.input,
            output_tokens=tokens.output,
            timestamp=ts,
        )
    except ClientError as exc:
        return _PromptResult(
            prompt=prompt_text,
            response=None,
            model=model_alias,
            input_tokens=0,
            output_tokens=0,
            timestamp=ts,
            error=str(exc),
        )


# ── Click command ─────────────────────────────────────────────────────


@click.command("batch")
@click.argument("file", type=click.Path(exists=True))
@click.option("--model-id", default=DEFAULT_MODEL_ALIAS, help=model_help_text())
@click.option(
    "--concurrency",
    default=1,
    type=int,
    help="Number of prompts to process in parallel (default: 1)",
)
@click.option(
    "-o",
    "--output",
    "output_path",
    default=None,
    type=click.Path(),
    help="Write results as JSON array to a file",
)
@click.option("--system", "system_prompt", default=None, help="System prompt for all invocations")
@click.pass_context
def batch_cmd(
    ctx: click.Context,
    file: str,
    model_id: str,
    concurrency: int,
    output_path: str | None,
    system_prompt: str | None,
) -> None:
    """Process multiple prompts from a file.

    \b
    Supported file formats:
        .txt   — one prompt per line
        .json  — JSON array of prompt strings
        .jsonl — one JSON object per line with a "prompt" field

    \b
    Examples:
        alice batch prompts.txt
        alice batch prompts.json --model-id haiku --concurrency 4
        alice batch prompts.jsonl --output results.json
        alice batch prompts.txt --system "Be concise"
    """
    config = ctx.obj["config"]

    print_banner(personality.BANNER, config.quiet)

    # ── Parse prompt file ─────────────────────────────────────────────
    prompts = parse_prompt_file(file)
    if not prompts:
        raise AliceCLIError("No prompts found in the file.")

    print_status(f"Loaded {len(prompts)} prompt(s) from {file}", config.quiet)

    # ── Authenticate and build Bedrock client ─────────────────────────
    context = create_bedrock_context(config)
    username = context.username
    resolved_id = resolve_model_id(model_id, context.inference_profile_arns)
    client = context.client

    # ── Process prompts concurrently ──────────────────────────────────
    results: list[_PromptResult] = []
    total = len(prompts)

    with Progress(console=err_console) as progress:
        task = progress.add_task("Processing prompts…", total=total)

        with ThreadPoolExecutor(max_workers=max(1, concurrency)) as executor:
            futures = {
                executor.submit(
                    _invoke_prompt, client, resolved_id, model_id, prompt, system_prompt
                ): idx
                for idx, prompt in enumerate(prompts)
            }
            for future in as_completed(futures):
                result = future.result()
                results.append(result)
                progress.advance(task)

                # Stream results to stderr when no --output file
                if not output_path:
                    if result.error:
                        err_console.print(f"[red]✗ Error:[/red] {result.error}")
                    else:
                        err_console.print(f"[green]✓[/green] {result.prompt[:60]}…")
                        print_rule(config.quiet)
                        err_console.print(result.response or "(empty)")
                        print_rule(config.quiet)

    # ── Tally results ─────────────────────────────────────────────────
    succeeded = sum(1 for r in results if r.error is None)
    failed = total - succeeded
    total_input = sum(r.input_tokens for r in results)
    total_output = sum(r.output_tokens for r in results)

    # ── Summary display ───────────────────────────────────────────────
    err_console.print()
    err_console.print(f"  [bold]Batch complete:[/bold] {total} prompts processed")
    err_console.print(f"  [green]✓[/green] Succeeded: {succeeded}")
    if failed:
        err_console.print(f"  [red]✗[/red] Failed: {failed}")
    err_console.print(
        f"  Tokens: {total_input:,} in / {total_output:,} out / {total_input + total_output:,} total"  # noqa: E501
    )

    # ── Write JSON output ─────────────────────────────────────────────
    if output_path:
        output_data = [
            {
                "prompt": r.prompt,
                "response": r.response,
                "model": r.model,
                "tokens": {
                    "input": r.input_tokens,
                    "output": r.output_tokens,
                    "total": r.input_tokens + r.output_tokens,
                },
                "timestamp": r.timestamp,
                "error": r.error,
            }
            for r in results
        ]
        try:
            with open(output_path, "w", encoding="utf-8") as f:
                json.dump(output_data, f, indent=2, ensure_ascii=False)
            print_success(f"Results written to {output_path}", config.quiet)
        except Exception as exc:
            raise AliceCLIError(f"Could not write output file: {exc}") from exc

    # ── Persist to Cloud Locker ───────────────────────────────────────
    metadata: dict[str, Any] = {
        "command": "batch",
        "source_file": file,
        "concurrency": concurrency,
        "total": total,
        "succeeded": succeeded,
        "failed": failed,
        "results": [
            {
                "prompt": r.prompt,
                "response": r.response,
                "tokens": {
                    "input": r.input_tokens,
                    "output": r.output_tokens,
                    "total": r.input_tokens + r.output_tokens,
                },
                "error": r.error,
            }
            for r in results
        ],
    }
    if system_prompt:
        metadata["system"] = system_prompt

    locker = CloudLocker(
        config=config,
        username=username,
        bucket=get_memory_bucket(config.namespace, config.environment),
    )
    record = build_session_record(
        record_type="batch",
        model=resolved_id,
        prompt=f"Batch: {total} prompts from {file}",
        response=None,
        tokens=TokenUsage(
            input=total_input,
            output=total_output,
            total=total_input + total_output,
        ),
        metadata=metadata,
    )
    try:
        locker.persist(record, "batch")
        print_success("Session saved to Cloud Locker", config.quiet)
    except Exception as exc:
        err_console.print(f"[yellow]⚠ Could not save session: {exc}[/yellow]")
