#!/usr/bin/env python3
"""Run a multi-turn visibility experiment with a real LLM."""

from __future__ import annotations

import argparse
import json
import os
from enum import Enum
from pathlib import Path
from typing import Any

from pydantic_ai import Agent, models
from pydantic_ai._run_context import get_current_run_context
from pydantic_ai.exceptions import UnexpectedModelBehavior
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.models.openai import OpenAIChatModelSettings
from pydantic_ai.profiles.openai import OpenAIModelProfile

from rgltr import AuditLogger, boundary, configure, tag
from rgltr.audit import AuditEvent
from rgltr.context import get_current_context
from rgltr.pydantic_ai import RegulatedAgent


class Tag(Enum):
    CUSTOMER = "customer"


class Boundary(Enum):
    EXTERNAL = "external"


def _write_jsonl(path: Path, record: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(record) + "\n")


class LoggingOpenAIChatModel(OpenAIChatModel):
    def __init__(
        self,
        *,
        model_name: str,
        provider: Any,
        profile: Any | None,
        settings: Any | None,
        log_dir: Path,
        audit_file: Path | None,
    ) -> None:
        super().__init__(model_name=model_name, provider=provider, profile=profile, settings=settings)
        self._log_dir = log_dir
        self._audit_logger = AuditLogger(file=str(audit_file)) if audit_file else None

    async def request(
        self,
        messages: list[models.ModelMessage],
        model_settings: OpenAIChatModelSettings | None,
        model_request_parameters: models.ModelRequestParameters,
    ) -> models.ModelResponse:
        models.check_allow_model_requests()
        model_settings, model_request_parameters = self.prepare_request(model_settings, model_request_parameters)

        tools = self._get_tools(model_request_parameters)
        if not tools:
            tool_choice: str | None = None
        elif (
            not model_request_parameters.allow_text_output
            and OpenAIModelProfile.from_profile(self.profile).openai_supports_tool_choice_required
        ):
            tool_choice = "required"
        else:
            tool_choice = "auto"

        openai_messages = await self._map_messages(messages, model_request_parameters)

        response_format = None
        if model_request_parameters.output_mode == "native":
            output_object = model_request_parameters.output_object
            assert output_object is not None
            response_format = self._map_json_schema(output_object)
        elif model_request_parameters.output_mode == "prompted" and self.profile.supports_json_object_output:
            response_format = {"type": "json_object"}

        run_ctx = get_current_run_context()
        rgltr_ctx = get_current_context()
        run_id = (rgltr_ctx.run_id if rgltr_ctx else None) or (run_ctx.run_id if run_ctx else None) or "unknown"
        path = self._log_dir / f"{run_id}.jsonl"
        _write_jsonl(
            path,
            {
                "type": "request",
                "model": self.model_name,
                "messages": openai_messages,
                "tools": tools or [],
                "tool_choice": tool_choice,
                "response_format": response_format,
            },
        )

        response = await self._completions_create(
            messages,
            False,
            model_settings or {},
            model_request_parameters,
        )

        _write_jsonl(
            path,
            {
                "type": "response",
                "response": response.model_dump(exclude_none=True),
            },
        )
        if self._audit_logger:
            tool_names = {tool["function"]["name"] for tool in tools or [] if "function" in tool}
            tool_calls = response.choices[0].message.tool_calls if response.choices else None
            if tool_calls:
                for tool_call in tool_calls:
                    name = tool_call.function.name if tool_call.function else None
                    if name and name not in tool_names:
                        self._audit_logger(
                            AuditEvent(
                                event_type="tool_call_unknown",
                                tool_kind="function",
                                tool_name=name,
                                tool_call_id=tool_call.id,
                                agent_id=rgltr_ctx.agent_id if rgltr_ctx else None,
                                run_id=run_id,
                                active_tags=[],
                            )
                        )
        return self._process_response(response)


def _build_model(model_id: str, audit_file: Path | None) -> models.Model:
    model = models.infer_model(model_id)
    if isinstance(model, OpenAIChatModel):
        provider = getattr(model, "_provider", None)
        return LoggingOpenAIChatModel(
            model_name=model.model_name,
            provider=provider or "openai",
            profile=model.profile,
            settings=model.settings,
            log_dir=Path("scripts/logs"),
            audit_file=audit_file,
        )
    return model


def run_experiment(
    model: str,
    *,
    runs: int,
    retries: int,
    audit_file: str,
) -> None:
    configure(boundaries={Boundary.EXTERNAL: True}, audit_sink=AuditLogger(file=audit_file))

    def build_agent() -> RegulatedAgent:
        agent = Agent(
            _build_model(model, Path(audit_file)),
            retries=retries,
            output_retries=retries,
            model_settings={"temperature": 0},
        )

        @agent.tool_plain
        @tag(activates=[Tag.CUSTOMER])
        async def get_customer(customer_id: int) -> str:
            return f"customer:{customer_id}"

        @agent.tool_plain
        @boundary(Boundary.EXTERNAL, deny_message="post_to_slack blocked by customer boundary")
        async def post_to_slack(message: str) -> str:
            return message

        return RegulatedAgent(agent, agent_id="visibility-experiment")

    for run in range(1, runs + 1):
        regulated = build_agent()
        run_id = f"visibility-run-{run}"
        try:
            first = regulated.run_sync("Fetch customer 123", run_id=run_id)
        except UnexpectedModelBehavior as exc:
            print(f"[run {run}] failed to fetch customer after retries: {exc}")
            continue
        print(f"[run {run}] first output:", first.output)

        try:
            second = regulated.run_sync(
                "Post to slack: customer update needed.",
                message_history=first.all_messages(),
                run_id=run_id,
            )
        except UnexpectedModelBehavior as exc:
            print(f"[run {run}] failed to post after retries: {exc}")
            continue
        print(f"[run {run}] second output:", second.output)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run visibility change experiment.")
    parser.add_argument(
        "--model",
        default=os.environ.get("RGLTR_MODEL"),
        help="Model identifier, e.g. openai:gpt-4o (env: RGLTR_MODEL)",
    )
    parser.add_argument("--runs", type=int, default=3)
    parser.add_argument("--retries", type=int, default=3)
    parser.add_argument("--audit-file", default="rgltr-audit.jsonl")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if not args.model:
        raise SystemExit("Missing --model or RGLTR_MODEL env var.")
    run_experiment(
        args.model,
        runs=args.runs,
        retries=args.retries,
        audit_file=args.audit_file,
    )


if __name__ == "__main__":
    main()
