#!/usr/bin/env python3
"""Render a grounded RAG template safely into separate system and user chat messages.

Reads all client style settings exclusively from style.yaml (single source of truth),
splits the Prompty v2 system: and user: sections before rendering, and renders each
role independently inside Jinja2 SandboxedEnvironment with StrictUndefined. Returning
structured chat messages ensures retrieved chunks containing 'user:' or 'system:'
lines can never create an extra message role.
"""
from __future__ import annotations

import argparse
import json
import pathlib
import re
import sys
from typing import Any

HERE = pathlib.Path(__file__).resolve().parent
DEFAULT_STYLE_PATH = HERE / "style.yaml"
DEFAULT_TEMPLATE_PATH = HERE / "grounded_qa.prompty"

REQUIRED_STYLE_KEYS = (
    "client_name",
    "assistant_persona",
    "tone",
    "language",
    "citation_format",
    "max_answer_words",
    "refusal_message",
    "out_of_scope_message",
    "untrusted_data_directive",
)

DEFAULT_SAMPLE_CHUNKS = [
    {
        "id": "doc1-p3",
        "title": "Leave policy",
        "url": "https://intranet.example.com/hr/leave-policy#page=3",
        "content": "Full-time staff get 25 days of paid annual leave per year.",
    },
    {
        "id": "doc2-p1",
        "title": "Public holiday schedule",
        "url": "https://intranet.example.com/hr/holidays#page=1",
        "content": "There are 8 standard public bank holidays across England and Wales.",
    },
]


def load_yaml(file_path: pathlib.Path | str) -> dict[str, Any]:
    """Load YAML key-value pairs from style.yaml and fail fast if missing."""
    path = pathlib.Path(file_path)
    if not path.exists():
        sys.stderr.write(f"Error: Style file not found at {path}\n")
        sys.exit(1)

    try:
        import yaml

        with path.open("r", encoding="utf-8") as handle:
            data = yaml.safe_load(handle) or {}
    except ImportError:
        data = {}
        for raw_line in path.read_text(encoding="utf-8").splitlines():
            line = raw_line.strip()
            if line and not line.startswith("#") and ":" in line:
                key, val = line.split(":", 1)
                data[key.strip()] = val.strip().strip("\"'")

    if not isinstance(data, dict):
        sys.stderr.write(f"Error: Style file {path} must be a YAML mapping.\n")
        sys.exit(1)
    return data


def split_role_templates(template_path: pathlib.Path) -> tuple[str, str]:
    """Read a .prompty file, strip YAML front matter, and return (system_tpl, user_tpl)."""
    if not template_path.exists():
        sys.stderr.write(f"Error: Template file not found at {template_path}\n")
        sys.exit(1)

    raw_text = template_path.read_text(encoding="utf-8")
    fm_match = re.match(r"\A---\s*\n.*?\n---\s*\n(.*)\Z", raw_text, re.S)
    body = fm_match.group(1) if fm_match else raw_text

    role_match = re.search(
        r"^system:\s*\n(.*?)\nuser:\s*\n(.*)\Z",
        body.strip(),
        flags=re.S | re.M,
    )
    if not role_match:
        sys.stderr.write(
            f"Error: Template {template_path} must contain both 'system:' and 'user:' sections.\n"
        )
        sys.exit(1)

    return role_match.group(1).strip(), role_match.group(2).strip()


def build_context(
    style_config: dict[str, Any],
    sources: list[dict[str, Any]],
    question: str,
) -> dict[str, Any]:
    """Validate required style.yaml keys without hiding missing values behind defaults."""
    context: dict[str, Any] = {"question": question, "sources": sources}
    for key in REQUIRED_STYLE_KEYS:
        if key in style_config and style_config[key] not in (None, ""):
            context[key] = style_config[key]
    return context


def render_messages(
    style_config: dict[str, Any],
    sources: list[dict[str, Any]],
    question: str = "How many days of annual leave do full-time staff get?",
    template_path: pathlib.Path = DEFAULT_TEMPLATE_PATH,
) -> list[dict[str, str]]:
    """Render system and user turns separately in SandboxedEnvironment with StrictUndefined."""
    try:
        from jinja2 import StrictUndefined, TemplateError
        from jinja2.sandbox import SandboxedEnvironment
    except ImportError:
        sys.stderr.write("Error: jinja2 is required. Run: pip install jinja2\n")
        sys.exit(1)

    system_tpl, user_tpl = split_role_templates(template_path)
    env = SandboxedEnvironment(undefined=StrictUndefined, autoescape=False)
    context = build_context(style_config, sources, question)

    try:
        system_content = env.from_string(system_tpl).render(**context).strip()
        user_content = env.from_string(user_tpl).render(**context).strip()
    except TemplateError as exc:
        sys.stderr.write(f"Error rendering template: {exc}\n")
        sys.exit(1)

    return [
        {"role": "system", "content": system_content},
        {"role": "user", "content": user_content},
    ]


def parse_args() -> argparse.Namespace:
    """Parse CLI arguments."""
    parser = argparse.ArgumentParser(
        description="Render a grounded RAG template into isolated system and user chat messages."
    )
    parser.add_argument(
        "--style",
        "-s",
        type=pathlib.Path,
        default=DEFAULT_STYLE_PATH,
        help="Path to client style.yaml (default: snippets/prompts/style.yaml)",
    )
    parser.add_argument(
        "--template",
        "-t",
        type=pathlib.Path,
        default=DEFAULT_TEMPLATE_PATH,
        help="Path to .prompty template (default: snippets/prompts/grounded_qa.prompty)",
    )
    parser.add_argument(
        "--sources",
        type=pathlib.Path,
        default=None,
        help="Path to JSON file containing retrieved chunks (each with id, title, url, content)",
    )
    parser.add_argument(
        "--question",
        "-q",
        type=str,
        default="How many days of annual leave do full-time staff get?",
        help="Sample user question to render into the user message",
    )
    parser.add_argument(
        "--output",
        "-o",
        type=pathlib.Path,
        default=None,
        help="Optional file path to write the rendered messages JSON (default: stdout)",
    )
    return parser.parse_args()


def main() -> int:
    """Run CLI prompt rendering and print structured chat messages as JSON."""
    args = parse_args()
    style_data = load_yaml(args.style)

    sources = DEFAULT_SAMPLE_CHUNKS
    if args.sources is not None:
        if not args.sources.exists():
            sys.stderr.write(f"Error: Sources file not found at {args.sources}\n")
            return 1
        sources = json.loads(args.sources.read_text(encoding="utf-8"))

    messages = render_messages(
        style_config=style_data,
        sources=sources,
        question=args.question,
        template_path=args.template,
    )
    payload = json.dumps(messages, indent=2, ensure_ascii=False)

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(payload + "\n", encoding="utf-8")
        print(f"Rendered {len(messages)} chat messages to {args.output}")
    else:
        print(payload)

    return 0


if __name__ == "__main__":
    sys.exit(main())