import logging
import re
from textwrap import dedent
from typing import Any, Optional, Tuple
from uuid import uuid4

import boto3
from aws_durable_execution_sdk_python import (
    DurableContext,
    StepContext,
    durable_execution,
    durable_step,
)
from aws_durable_execution_sdk_python.config import CallbackConfig, Duration
from bedrock_agentcore.memory import MemoryClient
from bedrock_agentcore.memory.integrations.strands.config import AgentCoreMemoryConfig
from bedrock_agentcore.memory.integrations.strands.session_manager import (
    AgentCoreMemorySessionManager,
)
from mcp.client.streamable_http import streamable_http_client
from pydantic import BaseModel, Field
from strands import Agent, tool
from strands.models import BedrockModel
from strands.tools.mcp import MCPClient

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

BUCKET = "your-bucket"
PREFIX = "agent-reports/"
MODEL_ID = "global.anthropic.claude-haiku-4-5-20251001-v1:0"
MEM_ID = "your-memory-id"
ACTOR_ID = "sample-htl-agent"

s3_client = boto3.client("s3")
memory_client = MemoryClient()

mcp_client = MCPClient(
    lambda: streamable_http_client("https://knowledge-mcp.global.api.aws")
)

ACTOR_ID_REGEX = re.compile(
    r"^[A-Za-z0-9][A-Za-z0-9\-_\/]*(?::[A-Za-z0-9\-_\/]+)*[A-Za-z0-9\-_\/]*$"
)
SESSION_ID_REGEX = re.compile(r"^[A-Za-z0-9][A-Za-z0-9\-_]*$")


def sanitize_identifier(
    raw_value: str,
    *,
    allowed_specials: str,
    pattern: re.Pattern[str],
    fallback_prefix: str,
) -> str:
    """IDのサニタイズ."""

    allow_colon = ":" in allowed_specials
    allowed_chars = set(allowed_specials)
    if allow_colon:
        allowed_chars.add(":")

    cleaned = "".join(ch for ch in raw_value if ch.isalnum() or ch in allowed_chars)

    if allow_colon and ":" in cleaned:
        cleaned = ":".join(segment for segment in cleaned.split(":") if segment)
    elif not allow_colon and ":" in cleaned:
        cleaned = cleaned.replace(":", "")

    cleaned = cleaned.lstrip(allowed_specials.replace(":", ""))
    if not cleaned or not cleaned[0].isalnum():
        cleaned = ""

    if cleaned and pattern.fullmatch(cleaned):
        return cleaned

    return f"{fallback_prefix}-{uuid4().hex[:10]}"


def ensure_safe_ids(actor_id: str, session_id: str) -> Tuple[str, str]:
    """actor_idとsession_idのサニタイズ."""

    safe_actor_id = sanitize_identifier(
        actor_id,
        allowed_specials="-_/:",
        pattern=ACTOR_ID_REGEX,
        fallback_prefix="actor",
    )

    safe_session_id = sanitize_identifier(
        session_id,
        allowed_specials="-_",
        pattern=SESSION_ID_REGEX,
        fallback_prefix="session",
    )

    return safe_actor_id, safe_session_id


class AgentOutput(BaseModel):
    """レポート作成結果のスキーマ定義"""

    report_content: str = Field(
        description="レポートの内容。マークダウン形式で作成すること"
    )
    file_name: str = Field(description="ファイル名。拡張子.mdをつけること")
    denied_message: Optional[str] = Field(
        default=None,
        description="S3への保存が承認されなかった場合のメッセージ。保存しようとしていた内容の概要と、承認されなかったことを簡潔に説明すること。",
    )


@tool(context=True)
def upload_report(doc: str, file_name) -> bool:
    """S3に調査レポートをアップロードします

    Args:
        doc (str): レポート内容
        file_name (_type_): アップロードするファイル名

    Returns:
        bool: アップロード成功有無
    """
    logger.info("Uploading report to %s/%s", BUCKET, PREFIX + file_name)
    try:
        s3_client.put_object(
            Bucket=BUCKET,
            Key=f"{PREFIX}{file_name}",
            Body=doc.encode("utf-8"),
            ContentType="text/markdown",
        )
    except Exception:
        logger.exception("Failed to upload report to %s/%s", BUCKET, PREFIX + file_name)
        return False
    return True


def create_session(session_id: str) -> AgentCoreMemorySessionManager:
    # AgentCore Memoryの設定値
    safe_actor_id, safe_session_id = ensure_safe_ids(ACTOR_ID, session_id)
    agentcore_memory_config = AgentCoreMemoryConfig(
        memory_id=MEM_ID, session_id=safe_session_id, actor_id=safe_actor_id
    )

    # セッションマネージャーを作成
    session_manager = AgentCoreMemorySessionManager(
        agentcore_memory_config=agentcore_memory_config, region_name="us-east-1"
    )

    return session_manager


def define_agent(tools: Any, memory_session: AgentCoreMemorySessionManager) -> Agent:
    agent = Agent(
        model=BedrockModel(
            model_id=MODEL_ID,
            streaming=False,
        ),
        system_prompt=dedent("""
                                AWSについて調査レポートを作成してください。
                                調査時には必ずAWS Knowledge MCPを使用してください。
                                マークダウン形式で出力し、レポート結果をS3にアップロードしてください。
                                """),
        tools=tools,
        session_manager=memory_session,
        callback_handler=None,
        structured_output_model=AgentOutput,
    )
    return agent


@durable_step
def create_report(step_context: StepContext, prompt: str, session_id: str) -> Any:
    step_context.logger.info(
        "【CREATE_REPORT】調査レポートを作成します。",
    )
    session = create_session(session_id=session_id)
    with mcp_client:
        aws_search_tools = mcp_client.list_tools_sync()
        agent: Agent = define_agent(aws_search_tools, memory_session=session)

        result: AgentOutput = agent(prompt).structured_output

    step_context.logger.info(
        "MESSAGE: %s", result.model_dump_json(indent=2, ensure_ascii=False)
    )
    return result.model_dump()


@durable_step
def send_for_approval(
    step_context: StepContext, callback_id: str, file_name: str
) -> dict:
    """承認リクエストをユーザに送信"""
    step_context.logger.info(
        "【SEND_APPROVAL】%sのS3へのアップロードの承認をお願いします。 %s",
        file_name,
        callback_id,
    )

    # 実際に利用する場合には、ここでメール送信やSNS通知などを実装する
    # send_email(to=approver_email, subject="レポート保存承認リクエスト", body=message)

    return {
        "file_name": file_name,
        "callback_id": callback_id,
        "status": "sent_for_approval",
    }


@durable_step
def save_report(
    step_context: StepContext,
    session_id,
) -> Any:
    """レポートをS3に保存する"""
    step_context.logger.info(
        "【SAVE_REPORT】S3にレポートを保存します。",
    )
    session = create_session(session_id=session_id)
    agent: Agent = define_agent([upload_report], memory_session=session)

    result: AgentOutput = agent("作成したレポートを保存してください").structured_output

    step_context.logger.info(
        "MESSAGE: %s", result.model_dump_json(indent=2, ensure_ascii=False)
    )
    return result.model_dump()


@durable_step
def denied_report(
    step_context: StepContext,
    session_id: str,
) -> Any:
    """レポートをS3に保存する"""
    step_context.logger.info(
        "【DENIED_REPORT】S3への保存が承認されませんでした。",
    )
    session = create_session(session_id=session_id)
    agent: Agent = define_agent([], memory_session=session)

    result: AgentOutput = agent(
        "レポートの保存が承認されませんでした。"
    ).structured_output

    step_context.logger.info(
        "MESSAGE: %s", result.model_dump_json(indent=2, ensure_ascii=False)
    )
    return result.model_dump()


@durable_execution
def lambda_handler(event, context: DurableContext):
    prompt = event["prompt"]
    session_id = context.state.durable_execution_arn.split(":")[-1]

    # Step1: レポートを作成
    created_report = context.step(
        create_report(prompt, session_id), name="create-report"
    )

    # Step2: コールバックを作成する
    context.logger.info("コールバックを作成完了")
    callback = context.create_callback(
        name="user_approval",
        config=CallbackConfig(
            timeout=Duration.from_hours(24), heartbeat_timeout=Duration.from_hours(1)
        ),
    )
    context.logger.info("コールバックを作成完了: %s", callback.callback_id)
    
    context.logger.info("承認リクエストを送信開始")
    approval_request = context.step(
        send_for_approval(callback.callback_id, created_report["file_name"]),
    )
    context.logger.info("承認リクエストを送信完了: %s", approval_request)

    # Step3: コールバックの結果を取得
    context.logger.info("コールバック結果を取得: %s", callback.callback_id)
    approval_result = callback.result()
    context.logger.info("承認結果を取得: %s", approval_result)

    # Step4: 承認結果に基づいて条件分岐 - Remove lambda wrappers
    if approval_result:
        context.logger.info("承認されました。レポートを保存します。")
        saved_report = context.step(save_report(session_id), name="save-report")
        return saved_report
    else:
        context.logger.info("承認されませんでした。レポートは保存しません。")
        saved_report = context.step(denied_report(session_id), name="denied-report")
        return saved_report