"""Compose engine — shared logic for generating Strands agent code.

Used by both the CLI wizard (``alice compose``) and the TUI compose screen.
"""

from __future__ import annotations

from dataclasses import dataclass, field


@dataclass(frozen=True)
class ToolInfo:
    """Metadata for a Strands community tool."""

    name: str
    import_name: str
    description: str


# Curated list of commonly used strands-agents-tools.
AVAILABLE_TOOLS: list[ToolInfo] = [
    ToolInfo("calculator", "calculator", "Perform math calculations"),
    ToolInfo("python_repl", "python_repl", "Run Python code snippets"),
    ToolInfo("http_request", "http_request", "Make HTTP/API requests"),
    ToolInfo("file_read", "file_read", "Read files from disk"),
    ToolInfo("file_write", "file_write", "Write files to disk"),
    ToolInfo("shell", "shell", "Execute shell commands"),
    ToolInfo("editor", "editor", "Advanced file editing"),
    ToolInfo("memory", "memory", "Persistent agent memory"),
    ToolInfo("retrieve", "retrieve", "RAG retrieval from knowledge bases"),
    ToolInfo("journal", "journal", "Structured logging / journaling"),
]

TOOL_LOOKUP: dict[str, ToolInfo] = {t.name: t for t in AVAILABLE_TOOLS}


@dataclass
class AgentSpec:
    """User choices collected by the compose wizard."""

    name: str
    system_prompt: str
    model_provider: str = "bedrock"
    tools: list[str] = field(default_factory=list)


def _build_imports(spec: AgentSpec) -> list[str]:
    """Build the import lines for the generated agent file."""
    lines = ["from strands import Agent"]

    if spec.tools:
        tool_imports = ", ".join(TOOL_LOOKUP[t].import_name for t in spec.tools if t in TOOL_LOOKUP)
        if tool_imports:
            lines.append(f"from strands_tools import {tool_imports}")

    if spec.model_provider == "ollama":
        lines.append("from strands.models.ollama import OllamaModel")
    elif spec.model_provider == "litellm":
        lines.append("from strands.models.litellm import LiteLLMModel")

    return lines


def _build_model_block(spec: AgentSpec) -> str | None:
    """Build the model instantiation block, or None for default Bedrock."""
    if spec.model_provider == "ollama":
        return 'model = OllamaModel(host="http://localhost:11434", model_id="llama3.2")'
    if spec.model_provider == "litellm":
        return 'model = LiteLLMModel(model_id="gpt-4o")'
    if spec.model_provider == "custom":
        return "# model = YourCustomModel(...)  # TODO: configure your model provider"
    return None  # Bedrock default — no explicit model needed


def render_agent_code(spec: AgentSpec) -> str:
    """Render a complete, runnable Strands agent Python file."""
    parts: list[str] = []

    # Docstring
    parts.append(f'"""Strands Agent: {spec.name}"""')
    parts.append("")

    # Imports
    parts.extend(_build_imports(spec))
    parts.append("")

    # Model provider (if non-default)
    model_block = _build_model_block(spec)
    if model_block:
        parts.append("")
        parts.append(model_block)
        parts.append("")

    # Tool list
    if spec.tools:
        tool_names = ", ".join(TOOL_LOOKUP[t].import_name for t in spec.tools if t in TOOL_LOOKUP)
        parts.append(f"tools = [{tool_names}]")
    else:
        parts.append("tools = []")
    parts.append("")

    # Agent construction
    agent_kwargs: list[str] = [
        f'    system_prompt="{spec.system_prompt}"',
        "    tools=tools,",
    ]
    if model_block and spec.model_provider != "custom":
        agent_kwargs.append("    model=model,")

    parts.append("agent = Agent(")
    parts.extend(agent_kwargs)
    parts.append(")")
    parts.append("")

    # Main block
    parts.append('if __name__ == "__main__":')
    parts.append('    response = agent("Hello! What can you do?")')
    parts.append("    print(response)")
    parts.append("")

    return "\n".join(parts)
