"""Load target configuration and payload YAML files."""

from __future__ import annotations

import os
import re
from pathlib import Path

import yaml

from llm_pentest.models import AttackCategory, Chain, Payload, TargetConfig

PAYLOADS_DIR = Path(__file__).parent.parent.parent / "payloads"

_ENV_RE = re.compile(r"\$\{(\w+)(?::([^}]*))?\}")


def _expand_env(value: str) -> str:
    """Expand ${VAR} and ${VAR:default} in strings."""
    def _replace(m: re.Match) -> str:
        return os.environ.get(m.group(1), m.group(2) if m.group(2) is not None else m.group(0))
    return _ENV_RE.sub(_replace, value)


def _expand_env_recursive(obj):
    if isinstance(obj, str):
        return _expand_env(obj)
    if isinstance(obj, dict):
        return {k: _expand_env_recursive(v) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_expand_env_recursive(i) for i in obj]
    return obj


def load_target(path: str | Path) -> TargetConfig:
    with open(path) as f:
        data = yaml.safe_load(f)
    return TargetConfig(**_expand_env_recursive(data))


def load_payloads(
    directory: str | Path | None = None,
    category: AttackCategory | None = None,
    payload_id: str | None = None,
) -> list[Payload]:
    directory = Path(directory) if directory else PAYLOADS_DIR
    payloads: list[Payload] = []

    for yaml_file in sorted(directory.rglob("*.yaml")):
        with open(yaml_file) as f:
            data = yaml.safe_load(f)
        if not data:
            continue

        items = data if isinstance(data, list) else [data]
        for item in items:
            if not isinstance(item, dict) or "id" not in item:
                continue
            # Skip chain definitions in payload files
            if "steps" in item:
                continue
            p = Payload(**item)
            if category and p.category != category:
                continue
            if payload_id and p.id != payload_id:
                continue
            payloads.append(p)

    return payloads


def load_chains(directory: str | Path | None = None) -> list[Chain]:
    directory = Path(directory) if directory else PAYLOADS_DIR
    chains: list[Chain] = []

    for yaml_file in sorted(directory.rglob("*.yaml")):
        with open(yaml_file) as f:
            data = yaml.safe_load(f)
        if not data:
            continue

        items = data if isinstance(data, list) else [data]
        for item in items:
            if isinstance(item, dict) and "steps" in item:
                chains.append(Chain(**item))

    return chains


# --- Target Profiles ---

PROFILES: dict[str, dict] = {
    "openai": {
        "endpoint": "https://api.openai.com/v1/chat/completions",
        "auth": {"type": "bearer", "token": "${OPENAI_API_KEY}"},
        "request": {"method": "POST", "headers": {"Content-Type": "application/json"},
                     "body_template": {"model": "gpt-4o-mini", "messages": "{{messages}}", "temperature": 0}},
        "response_path": "choices[0].message.content",
    },
    "anthropic": {
        "endpoint": "https://api.anthropic.com/v1/messages",
        "auth": {"type": "header", "header_name": "x-api-key", "header_value": "${ANTHROPIC_API_KEY}"},
        "request": {"method": "POST", "headers": {"Content-Type": "application/json", "anthropic-version": "2023-06-01"},
                     "body_template": {"model": "claude-sonnet-4-20250514", "max_tokens": 1024, "messages": "{{messages}}"}},
        "response_path": "content[0].text",
    },
    "ollama": {
        "endpoint": "http://localhost:11434/api/chat",
        "auth": {"type": "none"},
        "request": {"method": "POST", "headers": {"Content-Type": "application/json"},
                     "body_template": {"model": "llama3", "messages": "{{messages}}", "stream": False}},
        "response_path": "message.content",
    },
}


def load_profile(name: str, model: str | None = None) -> TargetConfig:
    if name not in PROFILES:
        raise ValueError(f"Unknown profile: {name}. Available: {', '.join(PROFILES)}")
    data = _expand_env_recursive(PROFILES[name])
    data["name"] = name
    if model:
        data["request"]["body_template"]["model"] = model
    return TargetConfig(**data)
