"""Agent tool definitions for OpenAI function calling / LangChain."""

from __future__ import annotations

import json
from typing import Any, Callable

from fhir_agent_tools import allergy_intolerance
from fhir_agent_tools import bundle
from fhir_agent_tools import care_plan
from fhir_agent_tools import care_team
from fhir_agent_tools import condition
from fhir_agent_tools import coverage
from fhir_agent_tools import diagnostic_report
from fhir_agent_tools import encounter
from fhir_agent_tools import family_member_history
from fhir_agent_tools import goal
from fhir_agent_tools import immunization
from fhir_agent_tools import location
from fhir_agent_tools import medication_request
from fhir_agent_tools import medication_statement
from fhir_agent_tools import observation
from fhir_agent_tools import organization
from fhir_agent_tools import patient
from fhir_agent_tools import practitioner
from fhir_agent_tools import procedure


def _wrap_extractor(extractor: Callable[..., dict], param_name: str) -> Callable[[str], str]:
    """Wrap extractor to accept fhir_json string and return JSON string for agent."""

    def _fn(fhir_json: str) -> str:
        result = extractor(fhir_json)
        return json.dumps(result)

    return _fn


# Registry: (name, description, extractor, param_description)
_TOOL_REGISTRY: list[tuple[str, str, Callable[..., dict], str]] = [
    ("extract_bundle_entries", "Extract all resources from a FHIR Bundle (e.g. search results). Returns list of resources.", bundle.get_bundle_entries, "FHIR Bundle as JSON string"),
    ("extract_resources_by_type", "Filter Bundle entries by resource type (e.g. Patient, Observation).", bundle.get_resources_by_type, "FHIR Bundle as JSON string"),
    ("extract_bundle_total", "Get total match count from a searchset Bundle.", bundle.get_bundle_total, "FHIR Bundle as JSON string"),
    ("extract_patient_summary", "Extract key patient demographics: name, birth date, gender, identifiers, contact info.", patient.get_patient_summary, "FHIR Patient resource as JSON string"),
    ("extract_patient_name", "Get formatted patient name.", patient.get_patient_name, "FHIR Patient resource as JSON string"),
    ("extract_patient_identifiers", "Get patient identifiers (MRN, etc.).", patient.get_patient_identifiers, "FHIR Patient resource as JSON string"),
    ("extract_practitioner_summary", "Extract practitioner info: name, identifiers.", practitioner.get_practitioner_summary, "FHIR Practitioner resource as JSON string"),
    ("extract_observation_summary", "Extract observation: code, value, effective date.", observation.get_observation_summary, "FHIR Observation resource as JSON string"),
    ("extract_observation_value", "Get observation value (numeric or text).", observation.get_observation_value, "FHIR Observation resource as JSON string"),
    ("extract_condition_summary", "Extract condition: code, status, onset.", condition.get_condition_summary, "FHIR Condition resource as JSON string"),
    ("extract_medication_request_summary", "Extract medication request: medication, dosage, status.", medication_request.get_medication_request_summary, "FHIR MedicationRequest resource as JSON string"),
    ("extract_medication_statement_summary", "Extract medication statement.", medication_statement.get_medication_statement_summary, "FHIR MedicationStatement resource as JSON string"),
    ("extract_encounter_summary", "Extract encounter: type, period, status.", encounter.get_encounter_summary, "FHIR Encounter resource as JSON string"),
    ("extract_allergy_summary", "Extract allergy: substance, reaction, criticality.", allergy_intolerance.get_allergy_summary, "FHIR AllergyIntolerance resource as JSON string"),
    ("extract_diagnostic_report_summary", "Extract diagnostic report: code, conclusion.", diagnostic_report.get_diagnostic_report_summary, "FHIR DiagnosticReport resource as JSON string"),
    ("extract_immunization_summary", "Extract immunization: vaccine, date.", immunization.get_immunization_summary, "FHIR Immunization resource as JSON string"),
    ("extract_care_plan_summary", "Extract care plan: activities, status.", care_plan.get_care_plan_summary, "FHIR CarePlan resource as JSON string"),
    ("extract_care_team_summary", "Extract care team: members, roles.", care_team.get_care_team_summary, "FHIR CareTeam resource as JSON string"),
    ("extract_goal_summary", "Extract goal: description, target.", goal.get_goal_summary, "FHIR Goal resource as JSON string"),
    ("extract_procedure_summary", "Extract procedure: code, performed date.", procedure.get_procedure_summary, "FHIR Procedure resource as JSON string"),
    ("extract_family_history_summary", "Extract family member history: conditions.", family_member_history.get_family_history_summary, "FHIR FamilyMemberHistory resource as JSON string"),
    ("extract_coverage_summary", "Extract coverage: payor, period.", coverage.get_coverage_summary, "FHIR Coverage resource as JSON string"),
    ("extract_organization_summary", "Extract organization: name, identifiers.", organization.get_organization_summary, "FHIR Organization resource as JSON string"),
    ("extract_location_summary", "Extract location: name, address.", location.get_location_summary, "FHIR Location resource as JSON string"),
]


def get_tools() -> list[dict[str, Any]]:
    """Return list of tool definitions for OpenAI function calling / LangChain.

    Each tool has: name, description, parameters (JSON Schema).
    """
    tools = []
    for name, description, _, param_desc in _TOOL_REGISTRY:
        # Special case: extract_resources_by_type needs resource_type param
        if name == "extract_resources_by_type":
            tools.append({
                "type": "function",
                "function": {
                    "name": name,
                    "description": description,
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "fhir_json": {"type": "string", "description": param_desc},
                            "resource_type": {"type": "string", "description": "Resource type to filter (e.g. Patient, Observation)"},
                        },
                        "required": ["fhir_json", "resource_type"],
                    },
                },
            })
        else:
            tools.append({
                "type": "function",
                "function": {
                    "name": name,
                    "description": description,
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "fhir_json": {"type": "string", "description": param_desc},
                        },
                        "required": ["fhir_json"],
                    },
                },
            })
    return tools


def get_extractor_functions() -> list[tuple[str, Callable[..., str]]]:
    """Return list of (name, callable) for direct use.

    Each callable accepts fhir_json (and resource_type for extract_resources_by_type)
    and returns JSON string of the result.
    """
    result = []
    for name, _, extractor, _ in _TOOL_REGISTRY:
        if name == "extract_resources_by_type":

            def _make_resources_fn(ext):
                def _fn(fhir_json: str, resource_type: str) -> str:
                    r = ext(fhir_json, resource_type)
                    return json.dumps(r)
                return _fn

            result.append((name, _make_resources_fn(extractor)))
        else:
            result.append((name, _wrap_extractor(extractor, "fhir_json")))
    return result
