from __future__ import annotations as _annotations

import os
from collections.abc import Sequence
from pathlib import Path
from typing import Final

from pydantic_acp import (
    AcpSessionContext,
    AdapterConfig,
    FileSessionStore,
    FileSystemProjectionMap,
    HookProjectionMap,
    RuntimeAgent,
    create_acp_agent,
    run_acp,
)
from pydantic_acp.models import ModelOverride
from pydantic_acp.types import (
    AgentPromptBlock,
    AudioContentBlock,
    BlobResourceContents,
    EmbeddedResourceContentBlock,
    ImageContentBlock,
    ResourceContentBlock,
)
from pydantic_ai import Agent
from pydantic_ai.capabilities import Hooks, ValidatedToolArgs
from pydantic_ai.messages import ModelResponse, ToolCallPart
from pydantic_ai.models import ModelRequestContext
from pydantic_ai.models.test import TestModel
from pydantic_ai.tools import DeferredToolRequests, RunContext, ToolDefinition

__all__ = ("TravelPromptModelProvider", "acp_agent", "agent", "config", "main")

_DEMO_ROOT: Final[Path] = Path("agent_demos")
_TRAVEL_ROOT: Final[Path] = Path.cwd() / _DEMO_ROOT / "travel-agent"
_READ_TOOL: Final[str] = "read_trip_file"
_WRITE_TOOL: Final[str] = "write_trip_file"
_SESSION_STORE_ROOT: Final[Path] = (
    Path(os.getenv("ACP_EXAMPLE_SESSION_DIR", str(_DEMO_ROOT / "acp-sessions")))
    .expanduser()
    .resolve()
    / "pydantic-travel"
)
_MEDIA_MODEL_ENV_NAMES: Final[tuple[str, ...]] = (
    "ACP_TRAVEL_MEDIA_MODEL",
    "TRAVEL_MEDIA_MODEL",
)
_DEFAULT_FILES: Final[dict[str, str]] = {
    "itinerary.md": (
        "# Travel Brief\n\n"
        "- city: Lisbon\n"
        "- trip goal: design-heavy long weekend\n"
        "- constraints: walkable neighborhoods, low-friction transit\n"
    ),
    "ideas.txt": (
        "Try these prompts:\n"
        "- list trip files\n"
        "- read trip file itinerary.md\n"
        "- write trip file scratch.txt: book the riverside hotel\n"
    ),
}


def _default_model_name() -> str | TestModel:
    configured_model = os.getenv("ACP_TRAVEL_MODEL", "").strip()
    if configured_model:
        return configured_model
    return TestModel()


def _configured_media_model_name() -> str | None:
    for env_name in _MEDIA_MODEL_ENV_NAMES:
        configured_model = os.getenv(env_name, "").strip()
        if configured_model:
            return configured_model
    return None


def _ensure_travel_workspace() -> None:
    _TRAVEL_ROOT.mkdir(parents=True, exist_ok=True)
    for relative_path, content in _DEFAULT_FILES.items():
        file_path = _resolve_trip_path(relative_path)
        if not file_path.exists():
            file_path.write_text(content, encoding="utf-8")


def _resolve_trip_path(path: str) -> Path:
    candidate = (_TRAVEL_ROOT / path).resolve()
    try:
        candidate.relative_to(_TRAVEL_ROOT)
    except ValueError as exc:
        raise ValueError("Path must stay inside the travel demo workspace.") from exc
    return candidate


def _truncate_text(text: str, *, limit: int) -> str:
    if len(text) <= limit:
        return text
    return f"{text[:limit]}\n\n...[truncated]"


def _prompt_has_binary_media(prompt: Sequence[AgentPromptBlock]) -> bool:
    for block in prompt:
        if isinstance(block, ImageContentBlock | AudioContentBlock):
            return True
        if isinstance(block, ResourceContentBlock):
            mime_type = block.mime_type
            if mime_type is not None and not mime_type.startswith("text/"):
                return True
            continue
        if isinstance(block, EmbeddedResourceContentBlock) and isinstance(
            block.resource,
            BlobResourceContents,
        ):
            return True
    return False


def _prompt_has_image_media(prompt: Sequence[AgentPromptBlock]) -> bool:
    for block in prompt:
        if isinstance(block, ImageContentBlock):
            return True
        if isinstance(block, ResourceContentBlock):
            mime_type = block.mime_type
            if mime_type is not None and mime_type.startswith("image/"):
                return True
            continue
        if isinstance(block, EmbeddedResourceContentBlock):
            resource = block.resource
            if isinstance(resource, BlobResourceContents):
                mime_type = resource.mime_type
                if mime_type is not None and mime_type.startswith("image/"):
                    return True
    return False


class TravelPromptModelProvider:
    def get_prompt_model_override(
        self,
        session: AcpSessionContext,
        agent: RuntimeAgent,
        prompt: Sequence[AgentPromptBlock],
        model_override: ModelOverride | None,
    ) -> ModelOverride | None:
        del session, agent
        if not _prompt_has_binary_media(prompt):
            return model_override
        return _configured_media_model_name() or model_override


hooks = Hooks[None]()


@hooks.on.before_model_request
async def observe_before_model_request(
    ctx: RunContext[None],
    request_context: ModelRequestContext,
) -> ModelRequestContext:
    del ctx
    return request_context


@hooks.on.after_model_request
async def observe_after_model_request(
    ctx: RunContext[None],
    *,
    request_context: ModelRequestContext,
    response: ModelResponse,
) -> ModelResponse:
    del ctx, request_context
    return response


@hooks.on.before_tool_execute(tools=[_READ_TOOL])
async def observe_read_tool(
    ctx: RunContext[None],
    *,
    call: ToolCallPart,
    tool_def: ToolDefinition,
    args: ValidatedToolArgs,
) -> ValidatedToolArgs:
    del ctx, call, tool_def
    return args


@hooks.on.before_tool_execute(tools=[_WRITE_TOOL])
async def observe_write_tool(
    ctx: RunContext[None],
    *,
    call: ToolCallPart,
    tool_def: ToolDefinition,
    args: ValidatedToolArgs,
) -> ValidatedToolArgs:
    del ctx, call, tool_def
    return args


@hooks.on.after_tool_execute(tools=[_WRITE_TOOL])
async def observe_write_result(
    ctx: RunContext[None],
    *,
    call: ToolCallPart,
    tool_def: ToolDefinition,
    args: ValidatedToolArgs,
    result: object,
) -> object:
    del ctx, call, tool_def, args
    return result


agent = Agent(
    _default_model_name(),
    name="travel-agent",
    deps_type=type(None),
    capabilities=[hooks],
    output_type=[str, DeferredToolRequests],
    system_prompt=(
        "You are the ACP Kit travel example. "
        "Use the trip-file tools for grounded answers. "
        "When image or audio prompt blocks appear, the host may swap the active model override."
    ),
)


@agent.tool_plain
def describe_travel_surface() -> str:
    """Summarize the ACP-facing surfaces available in this travel example."""
    return "\n".join(
        (
            "Travel example features:",
            "- existing Hooks capability introspection rendered through HookProjectionMap",
            "- approval-gated file writes with ACP diffs",
            "- prompt-model override provider for image and audio prompts",
        ),
    )


@agent.tool_plain
def list_trip_files() -> str:
    """List the demo travel files available in the local workspace."""
    _ensure_travel_workspace()
    return "\n".join(sorted(path.name for path in _TRAVEL_ROOT.iterdir() if path.is_file()))


@agent.tool_plain(name=_READ_TOOL)
def read_trip_file(path: str, max_chars: int = 4000) -> str:
    """Read a travel file and return a bounded preview."""
    if max_chars <= 0:
        raise ValueError("max_chars must be positive.")
    _ensure_travel_workspace()
    file_path = _resolve_trip_path(Path(path).name if Path(path).is_absolute() else path)
    if not file_path.exists():
        raise ValueError(f"File not found: {file_path.name}")
    return _truncate_text(file_path.read_text(encoding="utf-8"), limit=max_chars)


@agent.tool_plain(name=_WRITE_TOOL, requires_approval=True)
def write_trip_file(path: str, content: str) -> str:
    """Write a travel file inside the local demo workspace."""
    _ensure_travel_workspace()
    file_path = _resolve_trip_path(Path(path).name if Path(path).is_absolute() else path)
    file_path.write_text(content, encoding="utf-8")
    return f"Wrote `{file_path.name}`."


config = AdapterConfig(
    session_store=FileSessionStore(_SESSION_STORE_ROOT),
    prompt_model_override_provider=TravelPromptModelProvider(),
    hook_projection_map=HookProjectionMap(
        hidden_event_ids=frozenset({"after_model_request"}),
        event_labels={
            "before_model_request": "Before Model",
            "before_tool_execute": "Before Execute",
            "after_tool_execute": "After Execute",
        },
    ),
    projection_maps=[
        FileSystemProjectionMap(
            default_read_tool=_READ_TOOL,
            default_write_tool=_WRITE_TOOL,
        ),
    ],
)
acp_agent = create_acp_agent(agent=agent, config=config)


def main() -> None:
    _ensure_travel_workspace()
    run_acp(agent=agent, config=config)


if __name__ == "__main__":  # pragma: no cover
    main()
