"""Assignment management tools for reviewing and approving work."""

import json
from mcp.types import TextContent

from ..client import get_mturk_client, format_error


def register_tools(register):
    """Register assignment-related tools."""

    async def list_assignments(args: dict) -> list[TextContent]:
        """List all assignments for a specific HIT."""
        try:
            client = get_mturk_client()
            params = {
                "HITId": args["hit_id"],
                "MaxResults": min(args.get("max_results", 100), 100),
            }
            if args.get("status"):
                params["AssignmentStatuses"] = [args["status"]]

            response = client.list_assignments_for_hit(**params)
            assignments = response.get("Assignments", [])
            result = {
                "count": len(assignments),
                "assignments": [
                    {
                        "assignment_id": a["AssignmentId"],
                        "worker_id": a["WorkerId"],
                        "status": a["AssignmentStatus"],
                        "accept_time": str(a.get("AcceptTime")),
                        "submit_time": str(a.get("SubmitTime")),
                        "answer": a.get("Answer"),
                    }
                    for a in assignments
                ],
            }
            return [TextContent(type="text", text=json.dumps(result, indent=2))]
        except Exception as e:
            return [TextContent(type="text", text=format_error(e))]

    register(
        "list_assignments",
        """Retrieves all worker submissions (assignments) for a specific HIT.

Use this tool to see who has worked on your HIT and what they submitted. This is the main way to get worker responses. Each assignment represents one worker's completion of your task.

The response includes the worker's answer in XML format (MTurk's standard format). Parse the <Answer> field to extract form field values. For example, if your form had <input name="response">, look for <QuestionIdentifier>response</QuestionIdentifier> and its <FreeText> value.

Assignment statuses:
- 'Submitted': Worker finished, awaiting your review - YOU NEED TO APPROVE OR REJECT
- 'Approved': You approved it, worker has been paid
- 'Rejected': You rejected it, worker was not paid

Filter by status to see only assignments needing review (status='Submitted') or to audit past decisions.

Returns assignment_id (needed for approve/reject), worker_id (identifies the worker), and the actual answer data.

IMPORTANT: Workers are not paid until you approve their submission. Assignments left unreviewed will auto-approve after the HIT's auto_approval_delay (default 3 days).""",
        {
            "type": "object",
            "properties": {
                "hit_id": {
                    "type": "string",
                    "description": "The unique HIT ID to list assignments for. Get this from create_hit, list_hits, or list_reviewable_hits."
                },
                "status": {
                    "type": "string",
                    "enum": ["Submitted", "Approved", "Rejected"],
                    "description": "Optional filter by assignment status. 'Submitted' shows work awaiting review (most common use). 'Approved'/'Rejected' show past reviewed work. Omit to see all assignments regardless of status."
                },
                "max_results": {
                    "type": "integer",
                    "description": "Maximum number of assignments to return, between 1 and 100. Defaults to 100.",
                    "default": 100,
                    "minimum": 1,
                    "maximum": 100
                },
            },
            "required": ["hit_id"],
        },
        list_assignments,
    )

    async def get_assignment(args: dict) -> list[TextContent]:
        """Get details of a specific assignment including the worker's answer."""
        try:
            client = get_mturk_client()
            response = client.get_assignment(AssignmentId=args["assignment_id"])
            a = response["Assignment"]
            hit = response.get("HIT", {})
            result = {
                "assignment_id": a["AssignmentId"],
                "hit_id": a["HITId"],
                "worker_id": a["WorkerId"],
                "status": a["AssignmentStatus"],
                "accept_time": str(a.get("AcceptTime")),
                "submit_time": str(a.get("SubmitTime")),
                "approval_time": str(a.get("ApprovalTime")) if a.get("ApprovalTime") else None,
                "rejection_time": str(a.get("RejectionTime")) if a.get("RejectionTime") else None,
                "answer": a.get("Answer"),
                "requester_feedback": a.get("RequesterFeedback"),
                "hit_title": hit.get("Title"),
            }
            return [TextContent(type="text", text=json.dumps(result, indent=2))]
        except Exception as e:
            return [TextContent(type="text", text=format_error(e))]

    register(
        "get_assignment",
        """Retrieves complete details for a single assignment, including the full worker response.

Use this tool when you need to examine one specific submission in detail, rather than listing all assignments. Useful for reviewing a particular worker's response or checking the status of a known assignment.

Returns all assignment data including:
- The worker's complete answer (in MTurk XML format)
- Timestamps for when they accepted, submitted, and when you approved/rejected
- Any feedback you provided when approving/rejecting
- The associated HIT's title for context

The answer field contains XML with the worker's form responses. Look for <QuestionIdentifier> tags matching your form input names, with values in <FreeText> or <SelectionIdentifier> tags.

This is helpful for:
- Detailed review of a specific submission before approving/rejecting
- Checking what feedback was given on a past assignment
- Debugging issues with a particular worker's submission""",
        {
            "type": "object",
            "properties": {
                "assignment_id": {
                    "type": "string",
                    "description": "The unique assignment ID to retrieve. Get this from list_assignments. Format is alphanumeric."
                },
            },
            "required": ["assignment_id"],
        },
        get_assignment,
    )

    async def approve_assignment(args: dict) -> list[TextContent]:
        """Approve an assignment and pay the worker."""
        try:
            client = get_mturk_client()
            params = {"AssignmentId": args["assignment_id"]}
            if args.get("feedback"):
                params["RequesterFeedback"] = args["feedback"]
            if args.get("override_rejection"):
                params["OverrideRejection"] = True

            client.approve_assignment(**params)
            return [TextContent(type="text", text=json.dumps({
                "success": True,
                "assignment_id": args["assignment_id"],
                "action": "approved",
                "override_rejection": args.get("override_rejection", False),
            }, indent=2))]
        except Exception as e:
            return [TextContent(type="text", text=format_error(e))]

    register(
        "approve_assignment",
        """Approves a worker's submission and triggers payment to the worker.

Use this tool when a worker has satisfactorily completed your task. Approval is the standard action for acceptable work - it pays the worker the HIT's reward amount and positively affects their reputation.

IMPORTANT: Approval triggers an immediate, irreversible payment from your account. The worker receives the full reward amount. MTurk fees are charged to you.

Best practices:
- Approve good-faith efforts even if not perfect (rejection hurts worker reputation)
- Reserve rejection for clear spam, gibberish, or obvious bad-faith submissions
- Provide feedback to help workers improve (optional but appreciated)

The override_rejection option lets you reverse a previous rejection - useful if you rejected by mistake or the worker appealed successfully. This will pay them and fix their reputation.

Workers can see your feedback, so keep it constructive. Example: "Thank you for your detailed response!" or "Good work, approved."

Assignments auto-approve after the HIT's auto_approval_delay if you don't manually review them.""",
        {
            "type": "object",
            "properties": {
                "assignment_id": {
                    "type": "string",
                    "description": "The unique assignment ID to approve. Get this from list_assignments. Must be in 'Submitted' status (or 'Rejected' if using override_rejection)."
                },
                "feedback": {
                    "type": "string",
                    "description": "Optional message to the worker explaining the approval. Workers can see this. Keep it brief and positive. Example: 'Great work, thank you!'"
                },
                "override_rejection": {
                    "type": "boolean",
                    "description": "Set to true to reverse a previous rejection and approve instead. Use this to fix mistaken rejections. The worker will be paid and their stats corrected.",
                    "default": False
                },
            },
            "required": ["assignment_id"],
        },
        approve_assignment,
    )

    async def reject_assignment(args: dict) -> list[TextContent]:
        """Reject an assignment. The worker will not be paid."""
        try:
            client = get_mturk_client()
            client.reject_assignment(
                AssignmentId=args["assignment_id"],
                RequesterFeedback=args["feedback"],
            )
            return [TextContent(type="text", text=json.dumps({
                "success": True,
                "assignment_id": args["assignment_id"],
                "action": "rejected",
                "feedback": args["feedback"],
            }, indent=2))]
        except Exception as e:
            return [TextContent(type="text", text=format_error(e))]

    register(
        "reject_assignment",
        """Rejects a worker's submission. The worker will NOT be paid and their reputation is negatively affected.

USE WITH CAUTION: Rejection is a serious action that hurts the worker. Only reject for clear violations:
- Spam or gibberish responses
- Obvious bot/automated submissions
- Completely off-topic or zero-effort responses
- Clear evidence of cheating or bad faith

DO NOT reject for:
- Minor mistakes or imperfect work (approve and provide feedback instead)
- Responses you disagree with (subjective tasks should be approved)
- Your own unclear instructions (that's on you, approve it)

You MUST provide feedback explaining why you're rejecting. This is required by MTurk and helps:
- Workers understand what went wrong
- Protect you if the worker disputes the rejection
- MTurk review the rejection if appealed

Workers can appeal rejections to MTurk, and frequent unjustified rejections can get your requester account flagged.

If you reject by mistake, use approve_assignment with override_rejection=true to reverse it.""",
        {
            "type": "object",
            "properties": {
                "assignment_id": {
                    "type": "string",
                    "description": "The unique assignment ID to reject. Get this from list_assignments. Must be in 'Submitted' status."
                },
                "feedback": {
                    "type": "string",
                    "description": "REQUIRED explanation for the rejection. Be specific about what was wrong. Example: 'Response was empty' or 'Answer does not address the question asked'. Workers see this message."
                },
            },
            "required": ["assignment_id", "feedback"],
        },
        reject_assignment,
    )

    async def approve_all_submitted_assignments(args: dict) -> list[TextContent]:
        """Approve all submitted (pending) assignments for a HIT."""
        try:
            client = get_mturk_client()
            hit_id = args["hit_id"]

            # Get all submitted assignments
            response = client.list_assignments_for_hit(
                HITId=hit_id,
                AssignmentStatuses=["Submitted"],
                MaxResults=100,
            )
            assignments = response.get("Assignments", [])

            approved = []
            failed = []

            for a in assignments:
                try:
                    params = {"AssignmentId": a["AssignmentId"]}
                    if args.get("feedback"):
                        params["RequesterFeedback"] = args["feedback"]
                    client.approve_assignment(**params)
                    approved.append(a["AssignmentId"])
                except Exception as e:
                    failed.append({"assignment_id": a["AssignmentId"], "error": str(e)})

            return [TextContent(type="text", text=json.dumps({
                "hit_id": hit_id,
                "approved_count": len(approved),
                "failed_count": len(failed),
                "approved_assignments": approved,
                "failed_assignments": failed,
            }, indent=2))]
        except Exception as e:
            return [TextContent(type="text", text=format_error(e))]

    register(
        "approve_all_submitted_assignments",
        """Bulk approves all pending (submitted) assignments for a HIT in one operation.

Use this tool when you want to quickly approve all work on a HIT without reviewing each submission individually. This is convenient for surveys or tasks where you trust all responses, or when you just want to pay everyone and move on.

IMPORTANT: This triggers payments for ALL pending submissions on the HIT. Make sure you:
1. Actually want to approve everything (no spam or bad submissions)
2. Have sufficient account balance (use get_account_balance to check)

The tool finds all assignments with status='Submitted' for the specified HIT and approves each one. It reports:
- How many were successfully approved
- Any that failed (with error details)

Failed approvals are usually due to the assignment already being approved/rejected, or network issues. Successfully approved assignments are paid immediately.

For HITs with many assignments, this is much faster than approving one at a time. The optional feedback message is sent to all workers.

Note: Only processes up to 100 assignments. For HITs with more pending work, run this tool multiple times.""",
        {
            "type": "object",
            "properties": {
                "hit_id": {
                    "type": "string",
                    "description": "The unique HIT ID to approve all submissions for. All pending (Submitted) assignments for this HIT will be approved."
                },
                "feedback": {
                    "type": "string",
                    "description": "Optional feedback message sent to ALL workers being approved. Keep it generic since it goes to everyone. Example: 'Thank you for your participation!'"
                },
            },
            "required": ["hit_id"],
        },
        approve_all_submitted_assignments,
    )
