"""Global Bedrock LLM client — shared across all pipeline steps.

Usage:
    from shared.bedrock_client import build_chat_model, get_opus_model_id, ...
"""
import os

try:
    import boto3
except ModuleNotFoundError as exc:  # pragma: no cover
    raise ModuleNotFoundError(
        "Missing dependency 'boto3'. Install project dependencies and run via the project environment.\n"
        "Recommended:\n"
        "  1) uv sync\n"
        "  2) uv run python <entry_point>.py"
    ) from exc

try:
    from dotenv import load_dotenv
except ModuleNotFoundError as exc:  # pragma: no cover
    raise ModuleNotFoundError(
        "Missing dependency 'python-dotenv'. Run `uv sync` and use `uv run ...`."
    ) from exc

try:
    from langchain_aws import ChatBedrockConverse
except ModuleNotFoundError as exc:  # pragma: no cover
    raise ModuleNotFoundError(
        "Missing dependency 'langchain-aws'. Run `uv sync` and use `uv run ...`."
    ) from exc

try:
    from botocore.config import Config
except ModuleNotFoundError:  # pragma: no cover
    Config = None  # type: ignore[assignment]

load_dotenv(override=True)

from shared.config import (
    BEDROCK_CONNECT_TIMEOUT,
    BEDROCK_MAX_ATTEMPTS,
    BEDROCK_READ_TIMEOUT,
    DEFAULT_HAIKU_MODEL,
    DEFAULT_OPUS_MODEL,
    DEFAULT_SONNET_MODEL,
    IR_MODEL_ID,
    REACT_MODEL_ID,
)

# Public aliases
DEFAULT_CLAUDE_MODEL       = DEFAULT_SONNET_MODEL
DEFAULT_CLAUDE_OPUS_MODEL  = DEFAULT_OPUS_MODEL
DEFAULT_CLAUDE_HAIKU_MODEL = DEFAULT_HAIKU_MODEL
_OPUS_45_MODEL = "global.anthropic.claude-opus-4-5-20251101-v1:0"
_OPUS_46_MODEL = DEFAULT_OPUS_MODEL

_CONNECT_TIMEOUT = BEDROCK_CONNECT_TIMEOUT
_READ_TIMEOUT    = BEDROCK_READ_TIMEOUT
_MAX_ATTEMPTS    = BEDROCK_MAX_ATTEMPTS

_MODEL_ALIASES = {
    # Sonnet 4.5
    "sonnet 4.5": DEFAULT_CLAUDE_MODEL,
    "sonnet-4.5": DEFAULT_CLAUDE_MODEL,
    "claude sonnet 4.5": DEFAULT_CLAUDE_MODEL,
    "claude-sonnet-4.5": DEFAULT_CLAUDE_MODEL,
    # Opus 4.5
    "opus 4.5": _OPUS_45_MODEL,
    "opus-4.5": _OPUS_45_MODEL,
    "claude opus 4.5": _OPUS_45_MODEL,
    "claude-opus-4.5": _OPUS_45_MODEL,
    # Opus 4.6
    "opus": _OPUS_46_MODEL,
    "opus 4.6": _OPUS_46_MODEL,
    "opus-4.6": _OPUS_46_MODEL,
    "claude opus 4.6": _OPUS_46_MODEL,
    "claude-opus-4.6": _OPUS_46_MODEL,
    # Haiku 4.5
    "haiku": DEFAULT_CLAUDE_HAIKU_MODEL,
    "haiku 4.5": DEFAULT_CLAUDE_HAIKU_MODEL,
    "haiku-4.5": DEFAULT_CLAUDE_HAIKU_MODEL,
    "claude haiku 4.5": DEFAULT_CLAUDE_HAIKU_MODEL,
    "claude-haiku-4.5": DEFAULT_CLAUDE_HAIKU_MODEL,
}


def _env_model(env_var: str, default: str) -> str:
    return (os.getenv(env_var) or "").strip() or default


def get_opus_model_id() -> str:
    return DEFAULT_OPUS_MODEL


def get_ir_model_id() -> str:
    return IR_MODEL_ID


def get_haiku_model_id() -> str:
    return DEFAULT_HAIKU_MODEL


def get_react_model_id() -> str:
    return REACT_MODEL_ID


def _first_non_empty(*keys: str) -> str | None:
    for key in keys:
        value = os.getenv(key)
        if value is not None:
            stripped = value.strip()
            if stripped:
                return stripped
    return None


def _resolve_model_name(model_name: str | None) -> str:
    configured = _first_non_empty("BEDROCK_MODEL_ID", "ANTHROPIC_BEDROCK_MODEL") or ""
    raw = (model_name or "").strip()
    target = raw or configured or DEFAULT_CLAUDE_MODEL
    return _MODEL_ALIASES.get(target.lower(), target)


def build_chat_model(
    model_name: str | None = None,
    temperature: float = 0,
    max_tokens: int | None = None,
) -> ChatBedrockConverse:
    profile = _first_non_empty("AWS_PROFILE", "aws_profile")
    access_key = _first_non_empty("AWS_ACCESS_KEY_ID", "AWS_ACCESS_KEY", "aws_access_key")
    secret_key = _first_non_empty("AWS_SECRET_ACCESS_KEY", "AWS_SECRET_KEY", "aws_secret_key")
    session_token = _first_non_empty("AWS_SESSION_TOKEN", "aws_session_token")
    region = _first_non_empty(
        "BEDROCK_AWS_REGION", "AWS_REGION", "AWS_DEFAULT_REGION", "aws_region",
    ) or "us-east-1"

    if access_key and secret_key:
        session = boto3.Session(
            aws_access_key_id=access_key,
            aws_secret_access_key=secret_key,
            aws_session_token=session_token,
            region_name=region,
        )
    elif profile:
        session = boto3.Session(profile_name=profile, region_name=region)
    else:
        session = boto3.Session(region_name=region)

    if session.get_credentials() is None:
        raise RuntimeError(
            "AWS credentials were not found. Set AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY "
            "(or aws_access_key/aws_secret_key), plus optional AWS_SESSION_TOKEN, "
            "or configure AWS_PROFILE."
        )

    kwargs: dict = {
        "model": _resolve_model_name(model_name),
        "region_name": region,
        "temperature": temperature,
    }

    if max_tokens is not None:
        kwargs["max_tokens"] = max_tokens

    if Config is not None:
        kwargs["config"] = Config(
            connect_timeout=_CONNECT_TIMEOUT,
            read_timeout=_READ_TIMEOUT,
            retries={"max_attempts": _MAX_ATTEMPTS, "mode": "standard"},
            tcp_keepalive=True,
        )

    if access_key and secret_key:
        kwargs["aws_access_key_id"] = access_key
        kwargs["aws_secret_access_key"] = secret_key
        if session_token:
            kwargs["aws_session_token"] = session_token
    elif profile:
        kwargs["credentials_profile_name"] = profile

    return ChatBedrockConverse(**kwargs)