"""MCP tool definitions for GA4GH specifications.

Each tool wraps a GA4GH client operation, providing MCP-compatible
input schemas and result formatting.
"""

from __future__ import annotations

import json

from mcp.types import TextContent, Tool

TOOLS: list[Tool] = [
    # Tier 1: DRS
    Tool(
        name="drs_get_object",
        description="Get a DRS object by ID from a GA4GH DRS server.",
        inputSchema={
            "type": "object",
            "properties": {
                "base_url": {"type": "string", "description": "DRS server base URL"},
                "object_id": {"type": "string", "description": "DRS object ID"},
            },
            "required": ["base_url", "object_id"],
        },
    ),
    # Tier 1: Beacon
    Tool(
        name="beacon_query",
        description="Query a GA4GH Beacon v2 server.",
        inputSchema={
            "type": "object",
            "properties": {
                "base_url": {"type": "string", "description": "Beacon server base URL"},
                "query": {"type": "string", "description": "Query string"},
                "dataset_ids": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Dataset IDs to query",
                },
            },
            "required": ["base_url", "query"],
        },
    ),
    # Tier 1: htsget
    Tool(
        name="htsget_reads",
        description="Get read data from an htsget server.",
        inputSchema={
            "type": "object",
            "properties": {
                "base_url": {"type": "string", "description": "htsget server base URL"},
                "id": {"type": "string", "description": "Read ID"},
                "reference_name": {"type": "string", "description": "Reference name (chromosome)"},
                "start": {"type": "integer", "description": "Start position"},
                "end": {"type": "integer", "description": "End position"},
            },
            "required": ["base_url", "id"],
        },
    ),
    # Tier 1: refget
    Tool(
        name="refget_sequence",
        description="Get a reference sequence by checksum from a refget server.",
        inputSchema={
            "type": "object",
            "properties": {
                "base_url": {"type": "string", "description": "refget server base URL"},
                "checksum": {"type": "string", "description": "Sequence checksum (MD5 or SHA-512)"},
                "start": {"type": "integer", "description": "Start position"},
                "end": {"type": "integer", "description": "End position"},
            },
            "required": ["base_url", "checksum"],
        },
    ),
    # Tier 1: VRS
    Tool(
        name="vrs_compute_id",
        description="Compute a GA4GH VRS identifier for a variation.",
        inputSchema={
            "type": "object",
            "properties": {
                "type": {"type": "string", "description": "VRS type (Allele, Haplotype, etc.)"},
                "definition": {"type": "string", "description": "JSON serialized VRS object"},
            },
            "required": ["type", "definition"],
        },
    ),
    # Tier 1: Phenopackets
    Tool(
        name="phenopacket_validate",
        description="Validate a Phenopacket document against the GA4GH schema.",
        inputSchema={
            "type": "object",
            "properties": {
                "phenopacket": {"type": "string", "description": "JSON Phenopacket document"},
            },
            "required": ["phenopacket"],
        },
    ),
    # Tier 2: WES
    Tool(
        name="wes_submit_run",
        description="Submit a workflow run to a GA4GH WES server.",
        inputSchema={
            "type": "object",
            "properties": {
                "base_url": {"type": "string", "description": "WES server base URL"},
                "workflow_url": {"type": "string", "description": "Workflow document URL"},
                "workflow_type": {
                    "type": "string",
                    "description": "Workflow type (CWL, WDL, Nextflow)",
                },
            },
            "required": ["base_url", "workflow_url"],
        },
    ),
    # Tier 2: TES
    Tool(
        name="tes_create_task",
        description="Create a task on a GA4GH TES server.",
        inputSchema={
            "type": "object",
            "properties": {
                "base_url": {"type": "string", "description": "TES server base URL"},
                "task": {"type": "string", "description": "JSON task definition"},
            },
            "required": ["base_url", "task"],
        },
    ),
    # Tier 2: TRS
    Tool(
        name="trs_search_tools",
        description="Search for tools/workflows in a GA4GH TRS registry.",
        inputSchema={
            "type": "object",
            "properties": {
                "base_url": {"type": "string", "description": "TRS server base URL"},
                "name": {"type": "string", "description": "Tool name filter"},
                "organization": {"type": "string", "description": "Organization filter"},
            },
            "required": ["base_url"],
        },
    ),
    # Tier 2: Data Connect
    Tool(
        name="data_connect_search",
        description="Search tabular data using SQL expressions via GA4GH Data Connect.",
        inputSchema={
            "type": "object",
            "properties": {
                "base_url": {"type": "string", "description": "Data Connect server base URL"},
                "query": {"type": "string", "description": "SQL-like query expression"},
            },
            "required": ["base_url", "query"],
        },
    ),
    # Tier 3: Passports
    Tool(
        name="passport_decode",
        description="Decode and inspect a GA4GH Passport JWT.",
        inputSchema={
            "type": "object",
            "properties": {
                "token": {"type": "string", "description": "Passport JWT string"},
            },
            "required": ["token"],
        },
    ),
    # Tier 3: DUO
    Tool(
        name="duo_match",
        description="Check if a research purpose is compatible with DUO data use conditions.",
        inputSchema={
            "type": "object",
            "properties": {
                "conditions": {
                    "type": "array",
                    "items": {"type": "object"},
                    "description": "Data use conditions (DUO codes)",
                },
                "purpose": {"type": "object", "description": "Research purpose"},
            },
            "required": ["conditions", "purpose"],
        },
    ),
    # Tier 3: Pedigree
    Tool(
        name="pedigree_validate",
        description="Validate a GA4GH Pedigree structure.",
        inputSchema={
            "type": "object",
            "properties": {
                "pedigree": {"type": "string", "description": "JSON pedigree document"},
            },
            "required": ["pedigree"],
        },
    ),
]


def list_tools() -> list[Tool]:
    """Return all registered GA4GH tools."""
    return TOOLS


async def handle_tool_call(name: str, arguments: dict) -> list[TextContent]:
    """Handle an MCP tool call.

    Routes the call to the appropriate GA4GH client and formats the result.
    """
    handler = _TOOL_HANDLERS.get(name)
    if handler is None:
        return [TextContent(type="text", text=f"Unknown tool: {name}")]
    return await handler(arguments)


async def _drs_get_object(args: dict) -> list[TextContent]:
    from ai4gh.specs.drs.client import DRSClient

    client = DRSClient(base_url=args["base_url"])
    try:
        obj = await client.get_object(args["object_id"])
        return [TextContent(type="text", text=obj.model_dump_json(indent=2))]
    finally:
        await client.close()


async def _beacon_query(args: dict) -> list[TextContent]:
    from ai4gh.specs.beacon.client import BeaconClient

    client = BeaconClient(base_url=args["base_url"])
    try:
        result = await client.query(args["query"])
        return [TextContent(type="text", text=json.dumps(result, indent=2, default=str))]
    finally:
        await client.close()


async def _htsget_reads(args: dict) -> list[TextContent]:
    from ai4gh.specs.htsget.client import HtsgetClient

    client = HtsgetClient(base_url=args["base_url"])
    try:
        ticket = await client.get_reads(
            args["id"],
            reference_name=args.get("reference_name"),
            start=args.get("start"),
            end=args.get("end"),
        )
        return [TextContent(type="text", text=ticket.model_dump_json(indent=2))]
    finally:
        await client.close()


async def _refget_sequence(args: dict) -> list[TextContent]:
    from ai4gh.specs.refget.client import RefgetClient

    client = RefgetClient(base_url=args["base_url"])
    try:
        seq = await client.get_sequence(
            args["checksum"],
            start=args.get("start"),
            end=args.get("end"),
        )
        return [TextContent(type="text", text=seq)]
    finally:
        await client.close()


async def _vrs_compute_id(args: dict) -> list[TextContent]:
    from ai4gh.specs.vrs.models import ga4gh_identify

    vrs_obj = json.loads(args["definition"])
    identifier = ga4gh_identify(args["type"], json.dumps(vrs_obj, sort_keys=True))
    return [TextContent(type="text", text=identifier)]


async def _phenopacket_validate(args: dict) -> list[TextContent]:
    from ai4gh.specs.phenopackets.models import Phenopacket

    try:
        data = json.loads(args["phenopacket"])
        pp = Phenopacket.model_validate(data)
        return [TextContent(type="text", text=f"Valid Phenopacket: {pp.id or 'unnamed'}")]
    except Exception as e:
        return [TextContent(type="text", text=f"Validation error: {e}")]


async def _wes_submit_run(args: dict) -> list[TextContent]:
    from ai4gh.specs.wes.client import WESClient

    client = WESClient(base_url=args["base_url"])
    try:
        result = await client.submit_run(
            workflow_url=args["workflow_url"],
            workflow_type=args.get("workflow_type", "CWL"),
        )
        return [TextContent(type="text", text=result.model_dump_json(indent=2))]
    finally:
        await client.close()


async def _tes_create_task(args: dict) -> list[TextContent]:
    from ai4gh.specs.tes.client import TESClient

    client = TESClient(base_url=args["base_url"])
    try:
        task_def = json.loads(args["task"])
        result = await client.create_task(task_def)
        return [TextContent(type="text", text=result.model_dump_json(indent=2))]
    finally:
        await client.close()


async def _trs_search_tools(args: dict) -> list[TextContent]:
    from ai4gh.specs.trs.client import TRSClient

    client = TRSClient(base_url=args["base_url"])
    try:
        tools = await client.search_tools(
            name=args.get("name"),
            organization=args.get("organization"),
        )
        result = [t.model_dump() for t in tools]
        return [TextContent(type="text", text=json.dumps(result, indent=2, default=str))]
    finally:
        await client.close()


async def _data_connect_search(args: dict) -> list[TextContent]:
    from ai4gh.specs.data_connect.client import DataConnectClient

    client = DataConnectClient(base_url=args["base_url"])
    try:
        data = await client.search(args["query"])
        return [TextContent(type="text", text=data.model_dump_json(indent=2))]
    finally:
        await client.close()


async def _passport_decode(args: dict) -> list[TextContent]:
    from ai4gh.specs.passports.validator import PassportValidator

    validator = PassportValidator()
    claims = validator.decode(args["token"], verify=False)
    return [TextContent(type="text", text=claims.model_dump_json(indent=2))]


async def _duo_match(args: dict) -> list[TextContent]:
    from ai4gh.specs.duo.matcher import DUOMatcher
    from ai4gh.specs.duo.models import DataUseCondition, ResearchPurpose

    matcher = DUOMatcher()
    conditions = [DataUseCondition.model_validate(c) for c in args["conditions"]]
    purpose = ResearchPurpose.model_validate(args["purpose"])
    result = matcher.check_compatibility(conditions, purpose)
    return [TextContent(type="text", text=result.model_dump_json(indent=2))]


async def _pedigree_validate(args: dict) -> list[TextContent]:
    from ai4gh.specs.pedigree.models import Pedigree

    try:
        data = json.loads(args["pedigree"])
        ped = Pedigree.model_validate(data)
        return [
            TextContent(
                type="text",
                text=f"Valid pedigree: {len(ped.individuals)} individuals, "
                f"{len(ped.get_affected())} affected",
            )
        ]
    except Exception as e:
        return [TextContent(type="text", text=f"Validation error: {e}")]


_TOOL_HANDLERS = {
    "drs_get_object": _drs_get_object,
    "beacon_query": _beacon_query,
    "htsget_reads": _htsget_reads,
    "refget_sequence": _refget_sequence,
    "vrs_compute_id": _vrs_compute_id,
    "phenopacket_validate": _phenopacket_validate,
    "wes_submit_run": _wes_submit_run,
    "tes_create_task": _tes_create_task,
    "trs_search_tools": _trs_search_tools,
    "data_connect_search": _data_connect_search,
    "passport_decode": _passport_decode,
    "duo_match": _duo_match,
    "pedigree_validate": _pedigree_validate,
}
