"""Autopilot agents for TiDB Cloud NextGen workflows."""

from pathlib import Path
from string import Template
from typing import Any

from pydantic_ai import Agent, ToolOutput
from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.models import Model
from pydantic_ai_harness import CodeMode
from pydantic_ai_skills import SkillsCapability

from agent_output import AgentIssueUpdate
from capabilities.fs import AgentDeps, FSCapability
from capabilities.github_issues import DiagnoseIssuesCapability, GitHubIssuesCapability
from capabilities.github_workflows import GitHubWorkflowsCapability
from capabilities.sleep import SleepCapability
from config import (
    APP_DELIVERY_DIR,
    AUTO_DEPLOY_DIR,
    BEDROCK_AWS_ACCESS_KEY_ID,
    BEDROCK_AWS_PROFILE,
    BEDROCK_AWS_REGION,
    BEDROCK_AWS_SECRET_ACCESS_KEY,
    BEDROCK_AWS_SESSION_TOKEN,
    GITHUB_LOGIN,
    MODEL,
    OBSERVABILITY_DIR,
)

github_workflows = GitHubWorkflowsCapability()
sleep = SleepCapability()

_SKILLS_DIR = Path(__file__).parent / "skills"
_INSTRUCTIONS_DIR = Path(__file__).parent / "instructions"


def _resolve_model() -> str | Model[Any]:
    """Return a BedrockModel with isolated credentials when MODEL is bedrock:*,
    otherwise return the model string as-is for openrouter/openai providers.

    This keeps Bedrock boto3 credentials separate from the default AWS credential
    chain, so destroy skill scripts can continue to access S3 with the default profile.
    """
    if MODEL.startswith("bedrock:"):
        from pydantic_ai.models.bedrock import BedrockConverseModel
        from pydantic_ai.providers.bedrock import BedrockProvider

        model_id = MODEL[len("bedrock:") :]
        provider = BedrockProvider(
            profile_name=BEDROCK_AWS_PROFILE,
            aws_access_key_id=BEDROCK_AWS_ACCESS_KEY_ID,
            aws_secret_access_key=BEDROCK_AWS_SECRET_ACCESS_KEY,
            aws_session_token=BEDROCK_AWS_SESSION_TOKEN,
            region_name=BEDROCK_AWS_REGION,
        )
        return BedrockConverseModel(model_id, provider=provider)
    return MODEL


_chat_tmpl = Template((_INSTRUCTIONS_DIR / "chat.tmpl").read_text())
CHAT_INSTRUCTIONS = _chat_tmpl.substitute(
    github_login=GITHUB_LOGIN,
    auto_deploy_dir=AUTO_DEPLOY_DIR,
    app_delivery_dir=APP_DELIVERY_DIR,
    observability_dir=OBSERVABILITY_DIR,
)

_diag_tmpl = Template((_INSTRUCTIONS_DIR / "diagnose.tmpl").read_text())
DIAGNOSE_INSTRUCTIONS = _diag_tmpl.substitute(
    github_login=GITHUB_LOGIN,
    auto_deploy_dir=AUTO_DEPLOY_DIR,
    app_delivery_dir=APP_DELIVERY_DIR,
    observability_dir=OBSERVABILITY_DIR,
)


def make_agent(
    capabilities: list[AbstractCapability[Any]],
    instructions: str,
) -> Agent[AgentDeps, AgentIssueUpdate]:
    return Agent(
        _resolve_model(),
        output_type=ToolOutput(
            AgentIssueUpdate,
            name="publish_issue_update",
            description="Return the final issue update to publish for this run",
        ),
        deps_type=AgentDeps,
        capabilities=capabilities,
        instructions=instructions,
        retries={"output": 3},
    )


diagnose_agent: Agent[AgentDeps, AgentIssueUpdate] = make_agent(
    [
        CodeMode(),
        FSCapability,
        SkillsCapability(directories=[str(_SKILLS_DIR)]),
        DiagnoseIssuesCapability(),
        github_workflows,
    ],
    instructions=DIAGNOSE_INSTRUCTIONS,
)

chat_agent: Agent[AgentDeps, str] = Agent(
    _resolve_model(),
    output_type=str,
    deps_type=AgentDeps,
    capabilities=[
        CodeMode(),
        FSCapability,
        SkillsCapability(directories=[str(_SKILLS_DIR)]),
        GitHubIssuesCapability(),
        github_workflows,
        sleep,
    ],
    instructions=CHAT_INSTRUCTIONS,
)