#!/usr/bin/env python3
import argparse
import http.cookiejar
import json
import os
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import Any


DEFAULT_WORKSPACE_ID = "ws_default"
DEFAULT_WORKSPACE_SLUG = "default"
DEFAULT_INTERNAL_UPSTREAM = "http://admin-api:3001"
MCP_PROTOCOL_VERSION = "2025-11-25"


class SmokeError(RuntimeError):
    pass


@dataclass
class JsonResponse:
    status: int
    headers: Any
    body: Any


def build_operation_payload(name: str, upstream_base_url: str) -> dict[str, Any]:
    return {
        "name": name,
        "display_name": "Internal Health Smoke",
        "category": "smoke",
        "protocol": "rest",
        "security_level": "standard",
        "target": {
            "kind": "rest",
            "base_url": upstream_base_url.rstrip("/"),
            "method": "GET",
            "path_template": "/health",
            "static_headers": {},
        },
        "input_schema": {
            "type": "object",
            "required": True,
            "fields": {
                "probe": {
                    "type": "string",
                    "required": True,
                    "description": "Smoke probe value.",
                },
            },
        },
        "output_schema": {
            "type": "object",
            "required": True,
            "fields": {
                "status": {
                    "type": "string",
                    "required": True,
                    "description": "Health status returned by admin-api.",
                },
            },
        },
        "input_mapping": {
            "rules": [
                {
                    "source": "$.mcp.probe",
                    "target": "$.request.query.probe",
                    "required": True,
                },
            ],
        },
        "output_mapping": {
            "rules": [
                {
                    "source": "$.response.body.status",
                    "target": "$.output.status",
                    "required": True,
                },
            ],
        },
        "execution_config": {
            "timeout_ms": 5000,
            "headers": {},
        },
        "tool_description": {
            "title": "Check internal service health",
            "description": "Checks the internal Crank admin-api health endpoint. Use this only for deployment smoke verification.",
            "tags": ["smoke", "health"],
            "examples": [{"input": {"probe": "ok"}}],
        },
    }


def agent_mcp_url(base_url: str, workspace_slug: str, agent_slug: str) -> str:
    return f"{base_url.rstrip('/')}/mcp/v1/{workspace_slug}/{agent_slug}"


def tools_call_payload(tool_name: str, arguments: dict[str, Any]) -> dict[str, Any]:
    return {
        "jsonrpc": "2.0",
        "id": 3,
        "method": "tools/call",
        "params": {
            "name": tool_name,
            "arguments": arguments,
        },
    }


class Client:
    def __init__(self, base_url: str, timeout_seconds: int) -> None:
        self.base_url = base_url.rstrip("/")
        self.timeout_seconds = timeout_seconds
        cookie_jar = http.cookiejar.CookieJar()
        self.opener = urllib.request.build_opener(
            urllib.request.HTTPCookieProcessor(cookie_jar)
        )

    def request_json(
        self,
        method: str,
        path_or_url: str,
        payload: dict[str, Any] | list[Any] | None = None,
        headers: dict[str, str] | None = None,
        expected: tuple[int, ...] = (200,),
    ) -> JsonResponse:
        url = (
            path_or_url
            if path_or_url.startswith("http://") or path_or_url.startswith("https://")
            else f"{self.base_url}{path_or_url}"
        )
        data = None
        request_headers = {"Accept": "application/json"}
        if headers:
            request_headers.update(headers)
        if payload is not None:
            data = json.dumps(payload).encode("utf-8")
            request_headers["Content-Type"] = "application/json"

        request = urllib.request.Request(
            url,
            data=data,
            headers=request_headers,
            method=method,
        )

        try:
            response = self.opener.open(request, timeout=self.timeout_seconds)
            status = response.status
            raw = response.read().decode("utf-8")
            body = json.loads(raw) if raw else None
            headers_obj = response.headers
        except urllib.error.HTTPError as error:
            status = error.code
            raw = error.read().decode("utf-8")
            body = json.loads(raw) if raw else None
            headers_obj = error.headers

        if status not in expected:
            raise SmokeError(f"{method} {url} returned {status}: {body}")

        return JsonResponse(status=status, headers=headers_obj, body=body)


def admin_path(workspace_id: str, suffix: str) -> str:
    return f"/api/admin/workspaces/{workspace_id}{suffix}"


def login(client: Client, email: str, password: str) -> None:
    client.request_json(
        "POST",
        "/api/auth/login",
        {"email": email, "password": password},
        expected=(200,),
    )


def resolve_workspace(
    client: Client,
    fallback_workspace_id: str,
    fallback_workspace_slug: str,
) -> tuple[str, str]:
    session = client.request_json("GET", "/api/auth/session").body
    memberships = session.get("memberships") or []
    current_workspace_id = session.get("current_workspace_id") or fallback_workspace_id
    membership = next(
        (
            item
            for item in memberships
            if item.get("workspace", {}).get("id") == current_workspace_id
        ),
        None,
    )
    if not membership and memberships:
        membership = memberships[0]

    workspace = membership.get("workspace", {}) if membership else {}
    workspace_id = workspace.get("id") or current_workspace_id
    workspace_slug = workspace.get("slug") or fallback_workspace_slug
    if not workspace_id:
        raise SmokeError("authenticated session does not include a workspace id")
    if not workspace_slug:
        raise SmokeError("authenticated session does not include a workspace slug")
    return workspace_id, workspace_slug


def create_operation(
    client: Client,
    workspace_id: str,
    operation_name: str,
    internal_upstream: str,
) -> str:
    created = client.request_json(
        "POST",
        admin_path(workspace_id, "/operations"),
        build_operation_payload(operation_name, internal_upstream),
    ).body
    return created["operation_id"]


def publish_operation(client: Client, workspace_id: str, operation_id: str) -> None:
    client.request_json(
        "POST",
        admin_path(workspace_id, f"/operations/{operation_id}/publish"),
        {"version": 1},
    )


def create_agent(client: Client, workspace_id: str, agent_slug: str) -> str:
    created = client.request_json(
        "POST",
        admin_path(workspace_id, "/agents"),
        {
            "slug": agent_slug,
            "display_name": "Internal Health Smoke",
            "description": "Deployment smoke agent.",
            "instructions": {},
            "tool_selection_policy": {},
        },
    ).body
    return created["agent_id"]


def bind_and_publish_agent(
    client: Client,
    workspace_id: str,
    agent_id: str,
    operation_id: str,
    tool_name: str,
) -> None:
    client.request_json(
        "POST",
        admin_path(workspace_id, f"/agents/{agent_id}/bindings"),
        [
            {
                "operation_id": operation_id,
                "operation_version": 1,
                "tool_name": tool_name,
                "tool_title": "Check internal service health",
                "tool_description_override": None,
                "enabled": True,
            }
        ],
    )
    client.request_json(
        "POST",
        admin_path(workspace_id, f"/agents/{agent_id}/publish"),
        {"version": 1},
    )


def create_agent_key(client: Client, workspace_id: str, agent_id: str) -> tuple[str, str]:
    created = client.request_json(
        "POST",
        admin_path(workspace_id, f"/agents/{agent_id}/platform-api-keys"),
        {"name": f"smoke-key-{int(time.time())}", "scopes": ["read", "write"]},
    ).body
    return created["secret"], created["api_key"]["api_key"]["id"]


def cleanup_smoke_assets(
    client: Client,
    workspace_id: str,
    operation_id: str | None,
    agent_id: str | None,
    key_id: str | None,
) -> None:
    if os.environ.get("CRANK_PRODUCT_SMOKE_KEEP_ASSETS") == "1":
        print("cleanup skipped: CRANK_PRODUCT_SMOKE_KEEP_ASSETS=1")
        return

    if agent_id and key_id:
        try:
            client.request_json(
                "DELETE",
                admin_path(workspace_id, f"/agents/{agent_id}/platform-api-keys/{key_id}"),
                expected=(200, 204, 404),
            )
        except SmokeError as error:
            print(f"cleanup warning: key delete failed: {error}")

    if agent_id:
        try:
            client.request_json(
                "DELETE",
                admin_path(workspace_id, f"/agents/{agent_id}"),
                expected=(200, 404),
            )
        except SmokeError as error:
            print(f"cleanup warning: agent delete failed: {error}")

    if operation_id:
        try:
            client.request_json(
                "DELETE",
                admin_path(workspace_id, f"/operations/{operation_id}"),
                expected=(200, 404),
            )
        except SmokeError:
            try:
                client.request_json(
                    "POST",
                    admin_path(workspace_id, f"/operations/{operation_id}/archive"),
                    expected=(200, 404),
                )
            except SmokeError as error:
                print(f"cleanup warning: operation cleanup failed: {error}")


def initialize_mcp_session(client: Client, mcp_url: str, api_key: str) -> str:
    initialized = client.request_json(
        "POST",
        mcp_url,
        {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {"protocolVersion": MCP_PROTOCOL_VERSION},
        },
        headers={
            "Accept": "application/json, text/event-stream",
            "Authorization": f"Bearer {api_key}",
        },
    )
    session_id = initialized.headers.get("MCP-Session-Id")
    if not session_id:
        raise SmokeError("MCP initialize response did not include MCP-Session-Id")

    client.request_json(
        "POST",
        mcp_url,
        {
            "jsonrpc": "2.0",
            "method": "notifications/initialized",
            "params": {},
        },
        headers={
            "Accept": "application/json, text/event-stream",
            "Authorization": f"Bearer {api_key}",
            "MCP-Session-Id": session_id,
            "MCP-Protocol-Version": MCP_PROTOCOL_VERSION,
        },
        expected=(202,),
    )
    return session_id


def call_mcp(
    client: Client,
    mcp_url: str,
    api_key: str,
    session_id: str,
    payload: dict[str, Any],
) -> Any:
    return client.request_json(
        "POST",
        mcp_url,
        payload,
        headers={
            "Accept": "application/json, text/event-stream",
            "Authorization": f"Bearer {api_key}",
            "MCP-Session-Id": session_id,
            "MCP-Protocol-Version": MCP_PROTOCOL_VERSION,
        },
    ).body


def run(args: argparse.Namespace) -> None:
    admin_email = os.environ.get("CRANK_STAGING_ADMIN_EMAIL")
    admin_password = os.environ.get("CRANK_STAGING_ADMIN_PASSWORD")
    if not admin_email:
        raise SmokeError("CRANK_STAGING_ADMIN_EMAIL is required")
    if not admin_password:
        raise SmokeError("CRANK_STAGING_ADMIN_PASSWORD is required")

    timestamp = int(time.time())
    operation_name = f"internal_health_smoke_{timestamp}"
    agent_slug = f"health-smoke-{timestamp}"
    client = Client(args.base_url, args.timeout_seconds)
    operation_id = None
    agent_id = None
    key_id = None

    print(f"authenticated product smoke: {args.base_url.rstrip('/')}")
    login(client, admin_email, admin_password)
    print("login: ok")
    workspace_id, workspace_slug = resolve_workspace(
        client,
        args.workspace_id,
        args.workspace_slug,
    )
    print(f"workspace: {workspace_id} / {workspace_slug}")

    operation_id = create_operation(
        client,
        workspace_id,
        operation_name,
        args.internal_upstream,
    )
    print(f"operation created: {operation_id}")
    publish_operation(client, workspace_id, operation_id)
    print("operation published: v1")

    agent_id = create_agent(client, workspace_id, agent_slug)
    bind_and_publish_agent(client, workspace_id, agent_id, operation_id, operation_name)
    print(f"agent published: {agent_id}")

    api_key, key_id = create_agent_key(client, workspace_id, agent_id)
    mcp_url = agent_mcp_url(args.base_url, workspace_slug, agent_slug)
    session_id = initialize_mcp_session(client, mcp_url, api_key)
    print("mcp initialized: ok")

    tools = call_mcp(
        client,
        mcp_url,
        api_key,
        session_id,
        {"jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {}},
    )
    tool_names = [tool["name"] for tool in tools.get("result", {}).get("tools", [])]
    if operation_name not in tool_names:
        raise SmokeError(f"published tool {operation_name} not found in tools/list: {tool_names}")
    print("tools/list: ok")

    result = call_mcp(
        client,
        mcp_url,
        api_key,
        session_id,
        tools_call_payload(operation_name, {"probe": "ok"}),
    )
    if "error" in result:
        raise SmokeError(f"tools/call returned error: {result['error']}")
    print("tools/call: ok")
    cleanup_smoke_assets(client, workspace_id, operation_id, agent_id, key_id)
    print("authenticated product smoke completed")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run deterministic authenticated product smoke against deployed Crank."
    )
    parser.add_argument("base_url", help="Public Crank base URL, for example https://crank.example.com")
    parser.add_argument("--workspace-id", default=DEFAULT_WORKSPACE_ID)
    parser.add_argument("--workspace-slug", default=DEFAULT_WORKSPACE_SLUG)
    parser.add_argument("--internal-upstream", default=DEFAULT_INTERNAL_UPSTREAM)
    parser.add_argument(
        "--timeout-seconds",
        type=int,
        default=int(os.environ.get("SMOKE_TIMEOUT_SECONDS", "20")),
    )
    return parser.parse_args()


def main() -> int:
    try:
        run(parse_args())
    except SmokeError as error:
        print(f"smoke failed: {error}", flush=True)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())