#!/usr/bin/env python3
"""Map a CSV to add_campaign_leads calls. Planning is safe; --execute sends MCP requests."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import os
import sys
import time
import urllib.error
import urllib.request
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import Any

MAX_LEADS_PER_BATCH = 1_000
MAX_SERIALIZED_BYTES = 1_048_576
MCP_URL = "https://api.acquisity.ai/mcp"
MCP_PROTOCOL_VERSION = "2025-06-18"
TERMINAL_STATUSES = {"succeeded", "failed", "cancelled", "needs_attention"}

# Change these explicit CSV mappings before running. Custom values are preserved verbatim.
STANDARD_COLUMNS = {
    "Email": "email", "First name": "firstName", "Last name": "lastName",
    "Company": "company", "Title": "title", "Industry": "industry", "Phone": "phone",
    "LinkedIn": "linkedin", "Twitter": "twitter", "Facebook": "facebook",
    "Instagram": "instagram", "Employees": "employees", "Street": "street",
    "City": "city", "State": "state", "Postal code": "postalCode", "Country": "country",
    "Website": "website", "Company description": "companyDescription", "Tags": "tags",
}
CUSTOM_VARIABLE_COLUMNS = {"Opening line": "icebreaker", "Reason": "reasonToReachOut"}


class ImportPlanError(ValueError):
    pass


class McpError(RuntimeError):
    pass


@dataclass(frozen=True)
class ImportBatch:
    idempotency_key: str
    envelope: dict[str, Any]


def canonical_json(value: object) -> bytes:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode("utf-8")


def map_csv_rows(rows: Iterable[dict[str, str | None]]) -> tuple[list[dict[str, Any]], list[str]]:
    """Map only selected columns; do not trim or drop custom-variable whitespace/blanks."""
    leads: list[dict[str, Any]] = []
    warnings: list[str] = []
    for row_number, row in enumerate(rows, start=2):
        lead = {
            destination: value
            for column, destination in STANDARD_COLUMNS.items()
            if (value := row.get(column)) not in (None, "")
        }
        if not str(lead.get("email", "")).strip():
            warnings.append(f"row {row_number}: skipped because Email is empty")
            continue
        custom_variables = {
            variable: row[column]
            for column, variable in CUSTOM_VARIABLE_COLUMNS.items()
            if column in row and row[column] is not None
        }
        if custom_variables:
            lead["customVariables"] = custom_variables
        leads.append(lead)
    return leads, warnings


def unmapped_columns(fieldnames: Iterable[str | None]) -> list[str]:
    mapped = set(STANDARD_COLUMNS) | set(CUSTOM_VARIABLE_COLUMNS)
    return sorted(name for name in fieldnames if name and name not in mapped)


def batch_idempotency_key(workspace_id: str, campaign_id: str, body: dict[str, Any]) -> str:
    digest = hashlib.sha256(canonical_json({"workspaceId": workspace_id, "campaignId": campaign_id, "body": body})).hexdigest()
    return f"campaign-lead-upload:{digest[:40]}"


def mcp_envelope(workspace_id: str, campaign_id: str, body: dict[str, Any]) -> dict[str, Any]:
    return {
        "workspaceId": workspace_id,
        "campaignId": campaign_id,
        "idempotencyKey": batch_idempotency_key(workspace_id, campaign_id, body),
        "body": body,
    }


def tool_call_request(envelope: dict[str, Any]) -> dict[str, Any]:
    """The complete JSON-RPC request whose UTF-8 size is limited by the public API."""
    return {
        "jsonrpc": "2.0",
        "id": f"upload-{envelope['idempotencyKey']}",
        "method": "tools/call",
        "params": {"name": "add_campaign_leads", "arguments": envelope},
    }


def serialized_request_bytes(envelope: dict[str, Any]) -> int:
    return len(canonical_json(tool_call_request(envelope)))


def build_batches(
    leads: Iterable[dict[str, Any]], workspace_id: str, campaign_id: str, *, max_rows: int = MAX_LEADS_PER_BATCH, max_bytes: int = MAX_SERIALIZED_BYTES
) -> list[ImportBatch]:
    """Split by row count and the full JSON-RPC tools/call request size."""
    if not 1 <= max_rows <= MAX_LEADS_PER_BATCH:
        raise ImportPlanError(f"max_rows must be between 1 and {MAX_LEADS_PER_BATCH}")
    batches: list[ImportBatch] = []
    current: list[dict[str, Any]] = []

    def add_current() -> None:
        body = {"leads": current, "skipExistingInOtherCampaigns": True}
        envelope = mcp_envelope(workspace_id, campaign_id, body)
        batches.append(ImportBatch(envelope["idempotencyKey"], envelope))

    empty_overhead = serialized_request_bytes(
        mcp_envelope(workspace_id, campaign_id, {"leads": [], "skipExistingInOtherCampaigns": True})
    )
    current_bytes = empty_overhead
    for lead in leads:
        lead_bytes = len(canonical_json(lead))
        candidate_bytes = current_bytes + lead_bytes + (1 if current else 0)
        if len(current) < max_rows and candidate_bytes <= max_bytes:
            current.append(lead)
            current_bytes = candidate_bytes
            continue
        if not current:
            raise ImportPlanError("one lead exceeds the full JSON-RPC request limit")
        add_current()
        current = [lead]
        current_bytes = empty_overhead + lead_bytes
        if current_bytes > max_bytes:
            raise ImportPlanError("one lead exceeds the full JSON-RPC request limit")
    if current:
        add_current()
    return batches


def tool_result(response: dict[str, Any]) -> tuple[dict[str, Any], bool]:
    if "error" in response:
        error = response["error"]
        raise McpError(f"MCP protocol error {error.get('code', 'unknown')}: {error.get('message', 'unknown error')}")
    result = response.get("result")
    if not isinstance(result, dict):
        raise McpError("MCP response did not contain a tool result")
    content = result.get("content")
    if not isinstance(content, list) or not content or not isinstance(content[0], dict):
        raise McpError("MCP tool result did not contain JSON content")
    text = content[0].get("text")
    if not isinstance(text, str):
        raise McpError("MCP tool result did not contain text")
    try:
        payload = json.loads(text)
    except json.JSONDecodeError as error:
        raise McpError("MCP tool result was not JSON") from error
    if not isinstance(payload, dict):
        raise McpError("MCP tool result JSON was not an object")
    return payload, result.get("isError") is True


def operation_from_response(response: dict[str, Any]) -> tuple[dict[str, Any], bool]:
    payload, is_error = tool_result(response)
    operation = payload.get("operation", payload)
    if not isinstance(operation, dict) or not isinstance(operation.get("id"), str):
        raise McpError("MCP tool result did not contain an operation id")
    return operation, is_error


def read_operation(response: dict[str, Any]) -> dict[str, Any]:
    operation, is_error = operation_from_response(response)
    if is_error:
        raise McpError("get_operation returned an MCP error")
    return operation


class McpClient:
    def __init__(self, url: str, api_key: str, timeout_seconds: float):
        self.url = url
        self.api_key = api_key
        self.timeout_seconds = timeout_seconds
        self.protocol_version: str | None = None
        self.next_id = 1

    def post(self, payload: dict[str, Any], *, notification: bool = False) -> dict[str, Any] | None:
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
        }
        if self.protocol_version:
            headers["MCP-Protocol-Version"] = self.protocol_version
        request = urllib.request.Request(self.url, data=canonical_json(payload), headers=headers, method="POST")
        try:
            with urllib.request.urlopen(request, timeout=self.timeout_seconds) as response:
                if notification or response.status == 202:
                    return None
                body = response.read()
        except urllib.error.HTTPError as error:
            raise McpError(f"MCP HTTP request failed with status {error.code}") from error
        except urllib.error.URLError as error:
            raise McpError("MCP network request failed") from error
        try:
            parsed = json.loads(body)
        except json.JSONDecodeError as error:
            raise McpError("MCP returned a non-JSON response") from error
        if not isinstance(parsed, dict):
            raise McpError("MCP returned an invalid JSON response")
        return parsed

    def initialize(self) -> None:
        response = self.post({
            "jsonrpc": "2.0", "id": self.next_id, "method": "initialize",
            "params": {"protocolVersion": MCP_PROTOCOL_VERSION, "capabilities": {}, "clientInfo": {"name": "acquisity-csv-import-example", "version": "1.0"}},
        })
        self.next_id += 1
        if not response or not isinstance(response.get("result"), dict):
            raise McpError("MCP initialization failed")
        version = response["result"].get("protocolVersion")
        if not isinstance(version, str):
            raise McpError("MCP did not negotiate a protocol version")
        self.protocol_version = version
        self.post({"jsonrpc": "2.0", "method": "notifications/initialized"}, notification=True)

    def call_tool(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]:
        response = self.post({"jsonrpc": "2.0", "id": self.next_id, "method": "tools/call", "params": {"name": name, "arguments": arguments}})
        self.next_id += 1
        if response is None:
            raise McpError("MCP unexpectedly accepted a tool call without a result")
        return response


def wait_for_operations(operation_ids: Iterable[str], get_operation: Callable[[str], dict[str, Any]], *, timeout_seconds: float, interval_seconds: float) -> list[dict[str, Any]]:
    deadline = time.monotonic() + timeout_seconds
    pending = list(operation_ids)
    results: list[dict[str, Any]] = []
    while pending:
        if time.monotonic() >= deadline:
            raise McpError("Timed out while waiting for operations; retain their ids and re-poll with get_operation")
        operation_id = pending.pop(0)
        operation = get_operation(operation_id)
        status = operation.get("status")
        if status in {"queued", "running"}:
            pending.append(operation_id)
            time.sleep(min(interval_seconds, max(0, deadline - time.monotonic())))
        elif status in TERMINAL_STATUSES:
            results.append(operation)
        else:
            raise McpError(f"Operation {operation_id} returned an unknown status")
    return results


def submit_batches(
    batches: Iterable[ImportBatch],
    call_add_campaign_leads: Callable[[dict[str, Any]], dict[str, Any]],
    on_accepted: Callable[[dict[str, Any]], None] | None = None,
) -> tuple[list[dict[str, Any]], McpError | None]:
    """Keep every accepted operation, including one returned beside MCP isError."""
    accepted: list[dict[str, Any]] = []
    for batch in batches:
        try:
            operation, is_error = operation_from_response(
                call_add_campaign_leads(batch.envelope)
            )
        except McpError as error:
            return accepted, error
        accepted.append(operation)
        if on_accepted:
            on_accepted(operation)
        if is_error:
            return accepted, McpError(
                "MCP reported an error after accepting an operation; no later batches were submitted"
            )
    return accepted, None


def main() -> int:
    parser = argparse.ArgumentParser(description="Plan CSV imports; --execute performs explicit MCP writes.")
    parser.add_argument("csv_path", type=Path)
    parser.add_argument("--workspace-id", required=True)
    parser.add_argument("--campaign-id", required=True)
    parser.add_argument("--allow-unmapped", action="store_true")
    parser.add_argument("--execute", action="store_true", help="Send the planned batches through MCP.")
    parser.add_argument("--api-key-env", default="ACQUISITY_API_KEY")
    parser.add_argument("--mcp-url", default=MCP_URL)
    # The API has a 10s request budget and async tools spend nearly all of it polling.
    parser.add_argument("--request-timeout", type=float, default=30.0)
    parser.add_argument("--poll-timeout", type=float, default=300.0)
    parser.add_argument("--poll-interval", type=float, default=2.0)
    args = parser.parse_args()
    if args.request_timeout <= 0 or args.poll_timeout <= 0 or args.poll_interval <= 0:
        parser.error("timeouts and polling interval must be positive")

    with args.csv_path.open("r", encoding="utf-8-sig", newline="") as source:
        reader = csv.DictReader(source)
        unknown = unmapped_columns(reader.fieldnames or [])
        if unknown and not args.allow_unmapped:
            print("Unmapped columns: " + ", ".join(unknown), file=sys.stderr)
            print("Add them to the explicit mapping, then run again.", file=sys.stderr)
            return 2
        leads, warnings = map_csv_rows(reader)
    try:
        batches = build_batches(leads, args.workspace_id, args.campaign_id)
    except ImportPlanError as error:
        print(str(error), file=sys.stderr)
        return 2
    print("Standard mapping: " + json.dumps(STANDARD_COLUMNS, sort_keys=True))
    print("Custom-variable mapping: " + json.dumps(CUSTOM_VARIABLE_COLUMNS, sort_keys=True))
    print(f"Planned {len(leads)} leads in {len(batches)} request batches.")
    for warning in warnings:
        print(f"Warning: {warning}", file=sys.stderr)
    for index, batch in enumerate(batches, start=1):
        print(json.dumps({"batch": index, "leadCount": len(batch.envelope["body"]["leads"]), "serializedBytes": serialized_request_bytes(batch.envelope), "idempotencyKey": batch.idempotency_key, "tool": "add_campaign_leads"}, sort_keys=True))
    if not args.execute:
        print("No requests were sent. Re-run with --execute only after reviewing the plan.")
        return 0

    api_key = os.getenv(args.api_key_env)
    if not api_key:
        print(f"Set {args.api_key_env} before using --execute.", file=sys.stderr)
        return 2
    accepted: list[dict[str, Any]] = []
    try:
        client = McpClient(args.mcp_url, api_key, args.request_timeout)
        client.initialize()
        accepted, submission_error = submit_batches(
            batches,
            lambda envelope: client.call_tool("add_campaign_leads", envelope),
            lambda operation: print(json.dumps({"event": "operationAccepted", "operationId": operation.get("id"), "status": operation.get("status")}, sort_keys=True)),
        )
        operation_ids = [operation["id"] for operation in accepted if operation.get("status") not in TERMINAL_STATUSES]
        settled = wait_for_operations(
            operation_ids,
            lambda operation_id: read_operation(client.call_tool("get_operation", {"operationId": operation_id})),
            timeout_seconds=args.poll_timeout,
            interval_seconds=args.poll_interval,
        )
        outcomes = [*accepted, *settled]
        failed = [operation for operation in outcomes if operation.get("status") in TERMINAL_STATUSES - {"succeeded"}]
        for operation in outcomes:
            print(json.dumps({"operationId": operation.get("id"), "status": operation.get("status")}, sort_keys=True))
        submission_incomplete = submission_error is not None and len(accepted) < len(batches)
        if failed or submission_incomplete:
            print("One or more operations did not succeed; inspect their ids with get_operation before retrying.", file=sys.stderr)
            return 1
    except McpError as error:
        for operation in accepted:
            print(json.dumps({"operationId": operation.get("id"), "status": operation.get("status")}, sort_keys=True))
        print(f"Upload stopped: {error}", file=sys.stderr)
        return 1
    return 0


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