"""
Tool handlers for the Qwen Pipeline.

All tool call logic lives here. The server routes tool calls to these handlers.

To add a new tool:
1. Add the tool definition in your profile's tools.json (with parameters, description)
2. If the tool just calls an API endpoint, add "x_endpoint" to the tool definition — done.
   The pipeline will POST the args as JSON and return the response automatically.
3. If the tool needs custom pre/post-processing, add a handler function here
   and register it with @register_handler('tool_name').
"""

import asyncio
import json
import re
import time
import requests

# ── Registry for simple API tool handlers ─────────────────────────────────────

CUSTOM_HANDLERS = {}


def register_handler(name):
    """Decorator to register a custom tool handler (simple API tools).

    Usage:
        @register_handler('book_car_service')
        def handle_book_car_service(args, tool_def=None):
            # ... your logic ...
            return 'Booking confirmed.'
    """
    def decorator(func):
        CUSTOM_HANDLERS[name] = func
        return func
    return decorator


# ── Generic endpoint executor ─────────────────────────────────────────────────

# Internal-configuration vocabulary that agent-facing 'hint' texts sometimes
# carry (ordering API variant pins). The model tends to parrot these words to
# the caller, so hints containing them are dropped before the result reaches
# the model. Field names (outlet_pins, pinned_variant, ...) are kept — the
# agent's behavior rules depend on them; only the free-text hint is filtered.
_HINT_BANNED_RE = re.compile(
    r'pinn|pins_only|variant_pin|locked|permitted|fixed spec|set for', re.I)


def sanitize_tool_result(data):
    """Post-process a custom-tool JSON response before the model sees it.

    - Drops a top-level 'hint' whose text uses pin/lock vocabulary the agent
      must never speak.
    - Drops 'did_you_mean' (and its fuzzy-suggestion hint) when the fuzzy
      search corrected nothing — a merely relaxed match is a normal success,
      and the marker only prompts the model to say "I couldn't find ...".
    Mutates and returns `data`; non-dict input is returned unchanged.
    """
    if not isinstance(data, dict):
        return data

    dym = data.get('did_you_mean')
    if isinstance(dym, dict):
        corrected = dym.get('corrected') or {}
        real_correction = any(
            str(k).strip().lower() != str(v).strip().lower()
            for k, v in corrected.items())
        if not real_correction:
            del data['did_you_mean']
            # The accompanying hint says "fuzzy suggestions — CONFIRM";
            # without the marker it would be misleading.
            data.pop('hint', None)

    hint = data.get('hint')
    if isinstance(hint, str) and _HINT_BANNED_RE.search(hint):
        del data['hint']

    return data


def normalize_phone_number(phone, country_code='+60'):
    """Normalize phone numbers (default: Malaysian +60 prefix)."""
    phone = phone.replace(' ', '').replace('-', '')
    if phone.startswith('0'):
        phone = country_code + phone[1:]
    elif phone.startswith(country_code.lstrip('+')):
        phone = '+' + phone
    elif not phone.startswith('+'):
        phone = country_code + phone
    return phone


def execute_tool_endpoint(tool_def, args):
    """Execute a tool's x_endpoint API call.

    Args:
        tool_def: The tool definition dict (from tools.json), must contain 'x_endpoint'.
        args: The parsed arguments dict from the model.

    Returns:
        Result string, or None if no endpoint configured.
    """
    if not tool_def:
        return None
    endpoint = tool_def.get('x_endpoint')
    if not endpoint:
        return None

    # Normalize phone_number field if present
    if 'phone_number' in args:
        args['phone_number'] = normalize_phone_number(args['phone_number'])

    method = (tool_def.get('x_method') or 'POST').upper()
    use_query_params = tool_def.get('x_query_params', False)

    headers = {'Content-Type': 'application/json'}
    x_headers = tool_def.get('x_headers')
    if x_headers and isinstance(x_headers, dict):
        headers.update(x_headers)

    try:
        if method == 'GET' or use_query_params:
            print(f'  [Tool API] GET {endpoint} with params {json.dumps(args)[:200]}')
            resp = requests.get(endpoint, params=args, headers=headers,
                                timeout=15, verify=False)
        else:
            print(f'  [Tool API] POST {endpoint} with {json.dumps(args)[:200]}')
            resp = requests.post(endpoint, json=args, headers=headers,
                                 timeout=15, verify=False)
        print(f'  [Tool API] Response: {resp.status_code} {resp.text[:300]}')
        if resp.ok:
            try:
                return json.dumps(sanitize_tool_result(resp.json()))
            except Exception:
                return resp.text[:500]
        else:
            return f'API error: {resp.status_code} {resp.text[:200]}'
    except Exception as e:
        print(f'  [Tool API] Error: {e}')
        return f'API call failed: {str(e)}'


# ══════════════════════════════════════════════════════════════════════════════
# Pipeline tool handlers (async)
#
# Each handler receives:
#   p         — pipeline instance (self)
#   args      — parsed tool arguments dict
#   call_id   — the Omni tool call ID
#   arguments — raw arguments string
#   buffered_audio — list of base64 audio chunks buffered during mute
#
# Return conventions:
#   Return None   → handler managed everything (tool result + response)
#   Return string → common code will submit tool result and trigger response
# ══════════════════════════════════════════════════════════════════════════════


def _transfer_farewell_text(lang, to_name):
    """Fixed transfer farewell sentence per language.

    Spoken verbatim via 'Say EXACTLY' so the departing agent cannot
    hallucinate a failed transfer ("no one is available...")."""
    texts = {
        'en': f'Connecting you to {to_name} now, please hold.',
        'zh': f'正在为您转接{to_name}，请稍候。',
        'ms': f'Sedang menyambungkan anda kepada {to_name}, sila tunggu sebentar.',
    }
    return texts.get(lang or 'en', texts['en'])


def _detect_spoken_lang(p):
    """Best-effort language for fixed farewell text.

    session_language is only set by select/switch_language. When the
    caller simply speaks Chinese and Omni follows along, it stays 'en'
    — so fall back to CJK detection over recent conversation turns."""
    if p.session_language and p.session_language != 'en':
        return p.session_language
    try:
        recent = [str(m.get('content', ''))
                  for m in p._agent.get('history', [])[-6:]]
        text = ' '.join(recent)
        cjk = sum(1 for ch in text if '一' <= ch <= '鿿')
        if cjk >= 4:
            return 'zh'
    except Exception:
        pass
    return p.session_language or 'en'


async def handle_transfer_call(p, args, call_id, arguments, buffered_audio):
    """Transfer caller to another agent/department."""
    department = args.get('department', '')
    reason = args.get('reason', '')
    target_key, target_agent = p._resolve_transfer_target(department)

    if not target_key:
        return f'Transfer target "{department}" not found.'

    to_name = target_agent['name']
    from_name = p._agent['name']

    await p._emit({'type': 'transfer', 'from': from_name,
                   'to': to_name, 'reason': reason})
    p._log_message('system', f'Transfer: {from_name} -> {to_name} ({reason})')

    if call_id and p._omni_conv:
        # Cancel any lingering active response before farewell
        cancel_msg = json.dumps({"type": "response.cancel"})
        await p.loop.run_in_executor(None, p._omni_conv.send_raw, cancel_msg)
        await asyncio.sleep(0.2)

        # Send tool result to departing agent. DashScope auto-creates a
        # response when it receives the function_call_output — swallow it
        # (muted, transcript suppressed) so it cannot speak a hallucinated
        # outcome, then speak our fixed farewell.
        tool_result_msg = json.dumps({
            "type": "conversation.item.create",
            "item": {
                "type": "function_call_output",
                "call_id": call_id,
                "output": (f'SUCCESS: transfer initiated. The caller is being '
                           f'connected to {to_name} now.')
            }
        })
        p._omni_response_done_event.clear()
        p._omni_swallow_next_response = True
        await p.loop.run_in_executor(None, p._omni_conv.send_raw, tool_result_msg)
        # Wait for the swallowed auto-response to finish so our farewell
        # response.create is not rejected ("already has an active response").
        await p._omni_wait_response_done(timeout=5)
        p._omni_swallow_next_response = False  # In case no auto-response came

        # Farewell from departing agent before transfer.
        # Fixed sentence + "Say EXACTLY" — free-form instructions here get
        # overridden by the system prompt's unavailability rules and the
        # model hallucinates a failed transfer.
        _fw_text = _transfer_farewell_text(_detect_spoken_lang(p), to_name)
        farewell_msg = json.dumps({
            "type": "response.create",
            "response": {
                "modalities": ["text", "audio"],
                "instructions": (
                    f'Say EXACTLY the following, word for word. '
                    f'Do not add, remove, or change anything:\n\n'
                    f'"{_fw_text}"')
            }
        })
        await p._omni_send_response_and_wait(farewell_msg, timeout=10)
        if p.transfer_delay > 0:
            await p._emit_silence(p.transfer_delay)
            await asyncio.sleep(p.transfer_delay)

    # Switch to new agent
    p.active_agent = target_key
    await p._emit({'type': 'agent_change', 'agent': target_key, 'name': to_name})
    await p._omni_update_agent_session()

    # Trigger greeting from new agent
    if p._omni_conv:
        _gr_lang = ''
        if p.session_language and p.session_language != 'en':
            _gr_lang = (f' The caller has been speaking '
                        f'{p._lang_name(p.session_language)}. '
                        f'You MUST greet and respond in '
                        f'{p._lang_name(p.session_language)}.')
        greeting_instruction = (
            f'You have just been connected to a caller who was transferred to you. '
            f'Your name/role is: {to_name}. '
            f'Reason for transfer: {reason or "general inquiry"}. '
            f'Introduce yourself briefly (1-2 sentences) and ask how you can help '
            f'with their request. Be warm and professional.{_gr_lang}'
        )
        p._omni_suppress_flush = True
        response_msg = json.dumps({
            "type": "response.create",
            "response": {
                "modalities": ["text", "audio"],
                "instructions": greeting_instruction
            }
        })
        await p.loop.run_in_executor(None, p._omni_conv.send_raw, response_msg)
        print(f'  [Omni] Transfer greeting triggered for {to_name}')
    return None


async def handle_transfer_back(p, args, call_id, arguments, buffered_audio):
    """Transfer caller back to the main agent (Agent A)."""
    from_name = p._agent['name']
    to_name = p.agents['a']['name']
    reason = args.get('reason', '')

    await p._emit({'type': 'transfer', 'from': from_name,
                   'to': to_name, 'reason': reason})
    p._log_message('system', f'Transfer back: {from_name} -> {to_name} ({reason})')

    if call_id and p._omni_conv:
        # Cancel any lingering active response before farewell
        cancel_msg = json.dumps({"type": "response.cancel"})
        await p.loop.run_in_executor(None, p._omni_conv.send_raw, cancel_msg)
        await asyncio.sleep(0.2)

        # Send tool result to departing agent. Swallow DashScope's
        # auto-response — same anti-hallucination measure as
        # handle_transfer_call.
        tool_result_msg = json.dumps({
            "type": "conversation.item.create",
            "item": {
                "type": "function_call_output",
                "call_id": call_id,
                "output": (f'SUCCESS: transfer initiated. The caller is being '
                           f'connected back to {to_name} now.')
            }
        })
        p._omni_response_done_event.clear()
        p._omni_swallow_next_response = True
        await p.loop.run_in_executor(None, p._omni_conv.send_raw, tool_result_msg)
        await p._omni_wait_response_done(timeout=5)
        p._omni_swallow_next_response = False  # In case no auto-response came

        # Farewell from departing agent before transfer.
        # Fixed sentence + "Say EXACTLY" — same anti-hallucination measure
        # as handle_transfer_call.
        _bw_text = _transfer_farewell_text(_detect_spoken_lang(p), to_name)
        farewell_msg = json.dumps({
            "type": "response.create",
            "response": {
                "modalities": ["text", "audio"],
                "instructions": (
                    f'Say EXACTLY the following, word for word. '
                    f'Do not add, remove, or change anything:\n\n'
                    f'"{_bw_text}"')
            }
        })
        await p._omni_send_response_and_wait(farewell_msg, timeout=10)
        if p.transfer_delay > 0:
            await p._emit_silence(p.transfer_delay)
            await asyncio.sleep(p.transfer_delay)

    # Switch to Agent A
    p.active_agent = 'a'
    await p._emit({'type': 'agent_change', 'agent': 'a', 'name': to_name})

    # Inject transfer context into Agent A's history
    context_msg = (f'Caller returned from {from_name}. Reason: {reason}'
                   if reason else f'Caller returned from {from_name}.')
    p.agents['a']['history'].append({'role': 'system', 'content': context_msg})

    await p._omni_update_agent_session()

    if p._omni_conv:
        returned_greeting = p._pick_lang(
            p.agent_a_returned_greeting_map, p.session_language or 'en')
        if returned_greeting:
            # Speak exact configured greeting
            p._omni_suppress_flush = True
            response_msg = json.dumps({
                "type": "response.create",
                "response": {
                    "modalities": ["text", "audio"],
                    "instructions": (
                        f'Say EXACTLY the following, word for word. '
                        f'Do not add, remove, or change anything:\n\n'
                        f'"{returned_greeting}"')
                }
            })
            await p.loop.run_in_executor(None, p._omni_conv.send_raw, response_msg)
        else:
            # Inject transfer context as a conversation turn so the
            # model can use tools (e.g. dial_extension) to act on it.
            context_text = (
                f'[The caller was just transferred back from {from_name}. '
                f'Reason: {reason or "completed their request"}. '
                f'Act on this reason immediately — if it requires calling '
                f'a function (like dialing an extension or transferring), '
                f'do so now. Briefly acknowledge the transfer.]')
            item_msg = json.dumps({
                "type": "conversation.item.create",
                "item": {
                    "type": "message",
                    "role": "user",
                    "content": [{"type": "input_text", "text": context_text}]
                }
            })
            await p.loop.run_in_executor(None, p._omni_conv.send_raw, item_msg)
            p._omni_suppress_flush = True
            response_msg = json.dumps({"type": "response.create"})
            await p.loop.run_in_executor(None, p._omni_conv.send_raw, response_msg)
        print(f'  [Omni] Transfer-back greeting triggered for {to_name}')
    return None


async def handle_end_call(p, args, call_id, arguments, buffered_audio):
    """End the call — flush buffered farewell audio and wait for playback."""
    reason = args.get('reason', '')
    print(f'  [Omni end_call] agent={p.active_agent}, call_id={call_id}')
    p._log_message('system', f'Call ended by agent ({reason})')

    # Flush farewell audio BEFORE notifying the bridge, so the caller
    # hears the goodbye before the bridge starts teardown.
    if buffered_audio:
        print(f'  [Omni end_call] Flushing {len(buffered_audio)} buffered farewell audio chunks')
        p._omni_response_start_time = time.time()
        p._omni_response_audio_bytes = 0
        for chunk in buffered_audio:
            await p._emit({'type': 'audio', 'data': chunk,
                           '_gen': p._omni_audio_gen})
            p._omni_response_audio_bytes += len(chunk) * 3 // 4
    else:
        print(f'  [Omni end_call] No buffered farewell audio')

    # Wait for sender queue to drain
    for _i in range(50):
        if p.queue.empty():
            break
        await asyncio.sleep(0.1)
    # Wait for client-side playback of goodbye audio
    remaining_playback = 0.0
    if p._omni_response_audio_bytes > 0 and p._omni_response_start_time:
        audio_duration = p._omni_response_audio_bytes / (24000 * 2)
        elapsed = time.time() - p._omni_response_start_time
        remaining_playback = max(audio_duration - elapsed, 0)
    wait_buf = remaining_playback + 0.5
    print(f'  [Omni end_call] Waiting {wait_buf:.1f}s for farewell playback '
          f'(audio_bytes={p._omni_response_audio_bytes})')
    await asyncio.sleep(wait_buf)

    # Now notify the bridge — triggers hangup after farewell is played
    await p._emit({
        'type': 'tool_call', 'name': 'end_call',
        'arguments': arguments
    })

    print(f'  [Omni end_call] Farewell done, ending session')
    p.active = False
    return None


async def handle_call_back(p, args, call_id, arguments, buffered_audio):
    """Record a callback request and optionally dispatch to ACRS."""
    customer_name = args.get('customer_name', '')
    customer_phone = args.get('customer_phone', '').replace(' ', '')
    reason = args.get('reason', '')

    await p._emit({
        'type': 'call_back',
        'agent': p._agent.get('name', 'AI'),
        'customer_name': customer_name,
        'customer_phone': customer_phone,
        'reason': reason
    })
    p._log_message('system', f'Callback: {customer_name} ({customer_phone}) - {reason}')

    # ACRS callback dispatch
    if (p.acrs_enabled and p.acrs_api_url
            and p.acrs_callback_dept_id and not p.acrs_dispatched):
        try:
            acrs_result = await p.loop.run_in_executor(
                None, p._dispatch_callback_to_acrs,
                customer_name, customer_phone, reason)
            if acrs_result.get('success'):
                p.acrs_dispatched = True
            await p._emit({
                'type': 'acrs_status', 'callback': True,
                'success': acrs_result.get('success', False),
                'message': acrs_result.get('message', ''),
                'case_no': acrs_result.get('case_no', '')
            })
        except Exception as e:
            print(f'  [Omni ACRS] Callback dispatch error: {e}')

    return 'Callback request recorded.'


async def handle_select_language(p, args, call_id, arguments, buffered_audio):
    """Handle language selection from the language agent."""
    lang = args.get('language', 'en')
    if lang not in p.supported_languages:
        lang = 'en'
    p.session_language = lang
    lang_label = p._lang_name(lang)
    print(f'  [Omni LANG] Language selected: {lang} ({lang_label})')
    p._log_message('system', f'Language selected: {lang}')

    p.active_agent = 'a'
    to_name = p.agents['a']['name']
    await p._emit({'type': 'agent_change', 'agent': 'a', 'name': to_name})
    await p._omni_update_agent_session()

    result_text = f'Language set to {lang_label}. Transferred to main agent.'

    if call_id and p._omni_conv:
        # Send tool result
        tool_result_msg = json.dumps({
            "type": "conversation.item.create",
            "item": {
                "type": "function_call_output",
                "call_id": call_id,
                "output": result_text
            }
        })
        await p.loop.run_in_executor(None, p._omni_conv.send_raw, tool_result_msg)

        # Pick greeting: first-contact vs returned-contact, per chosen language
        if not p._main_agent_visited:
            greeting = p._pick_lang(p.fixed_greeting_map, lang)
            p._main_agent_visited = True
        else:
            greeting = (
                p._pick_lang(p.agent_a_returned_greeting_map, lang)
                or p._pick_lang(p.fixed_greeting_map, lang))

        if greeting:
            greeting_instruction = (
                f'Say EXACTLY the following greeting, word for word, '
                f'in the same language as written: "{greeting}"')
        else:
            greeting_instruction = (
                f'You have just been connected to a caller. '
                f'Your name/role is: {to_name}. '
                f'The caller selected {lang_label} as their language. '
                f'Introduce yourself briefly (1-2 sentences) in {lang_label} '
                f'and ask how you can help.')

        p._omni_suppress_flush = True
        response_msg = json.dumps({
            "type": "response.create",
            "response": {
                "modalities": ["text", "audio"],
                "instructions": greeting_instruction
            }
        })
        await p.loop.run_in_executor(None, p._omni_conv.send_raw, response_msg)
    return None


async def handle_change_language(p, args, call_id, arguments, buffered_audio):
    """Transfer to language selection agent."""
    reason = args.get('reason', '')
    print(f'  [Omni LANG] Change language requested (reason: {reason})')
    p._log_message('system', f'Change language requested: {reason}')

    # Clear locked language and reset lang agent history for a fresh menu
    p.session_language = None
    if 'lang' in p.agents:
        p.agents['lang']['history'] = []

    p.active_agent = 'lang'
    to_name = p.agents.get('lang', {}).get('name', 'Language Agent')
    await p._emit({'type': 'agent_change', 'agent': 'lang', 'name': to_name})
    await p._omni_update_agent_session()

    result_text = 'Transferred to language selection.'

    if call_id and p._omni_conv:
        # Send tool result
        tool_result_msg = json.dumps({
            "type": "conversation.item.create",
            "item": {
                "type": "function_call_output",
                "call_id": call_id,
                "output": result_text
            }
        })
        await p.loop.run_in_executor(None, p._omni_conv.send_raw, tool_result_msg)

        # Speak language agent greeting if configured
        if p.language_agent_greeting:
            greeting_instruction = (
                f'Say EXACTLY the following greeting, word for word: '
                f'"{p.language_agent_greeting}"')
        else:
            greeting_instruction = (
                f'You are {to_name}. The caller wants to change their '
                f'language. Ask them which language they would like to '
                f'continue in. Keep it brief (1-2 sentences).')

        p._omni_suppress_flush = True
        response_msg = json.dumps({
            "type": "response.create",
            "response": {
                "modalities": ["text", "audio"],
                "instructions": greeting_instruction
            }
        })
        await p.loop.run_in_executor(None, p._omni_conv.send_raw, response_msg)
    return None


async def handle_switch_language(p, args, call_id, arguments, buffered_audio):
    """Switch language mid-conversation (without transferring to lang agent)."""
    lang = args.get('language', 'en')
    if lang not in p.supported_languages:
        lang = 'en'
    lang_label = p._lang_name(lang)

    # Skip no-op switch
    if lang == p.session_language:
        print(f'  [Omni LANG] Already in {lang}, ignoring no-op switch')
        return f'Already using {lang_label}.'

    p.session_language = lang
    print(f'  [Omni LANG] Switched to: {lang} ({lang_label})')
    await p._emit({'type': 'language_switched', 'language': lang})
    p._log_message('system', f'Language switched to: {lang}')
    # Update Omni session with language-appropriate instructions
    await p._omni_update_agent_session()

    result_text = f'Language switched to {lang_label}.'

    if call_id and p._omni_conv:
        # Send tool result
        tool_result_msg = json.dumps({
            "type": "conversation.item.create",
            "item": {
                "type": "function_call_output",
                "call_id": call_id,
                "output": result_text
            }
        })
        await p.loop.run_in_executor(None, p._omni_conv.send_raw, tool_result_msg)

        # Follow-up in the new language
        p._omni_suppress_flush = True
        response_msg = json.dumps({
            "type": "response.create",
            "response": {
                "modalities": ["text", "audio"],
                "instructions": (
                    f'You just switched to {lang_label}. '
                    f'Say a single short sentence in {lang_label} '
                    f'asking how you can help. '
                    f'Do NOT introduce yourself again.')
            }
        })
        await p.loop.run_in_executor(None, p._omni_conv.send_raw, response_msg)
    return None


async def handle_dial(p, name, args, call_id, arguments, buffered_audio, result_text):
    """Handle dial_extension / dial_department — speak 'connecting' then return."""
    person = args.get('name', '') or args.get('department', '')
    _dl_lang = ''
    if p.session_language and p.session_language != 'en':
        _dl_lang = f' Speak in {p._lang_name(p.session_language)}.'
    connecting_msg = json.dumps({
        "type": "response.create",
        "response": {
            "modalities": ["text", "audio"],
            "instructions": (
                f'Tell the caller you are now connecting them to {person}. '
                f'Keep it to 1 short sentence.{_dl_lang}')
        }
    })
    print(f'  [Omni {name}] Generating "connecting" speech via dedicated response...')
    await p._omni_send_response_and_wait(connecting_msg, timeout=10)
    print(f'  [Omni {name}] "Connecting" speech done')
    return None


# ── Custom API handlers ───────────────────────────────────────────────────────
# Add your custom tool handlers below. Each handler receives:
#   args (dict): The parsed tool arguments from the model
#   tool_def (dict): The full tool definition from tools.json (optional)
#
# Return a string that will be sent back to the model as the tool result.
#
# Example:
#
# @register_handler('book_car_service')
# def handle_book_car_service(args, tool_def=None):
#     customer = args.get('customer_name', '')
#     vehicle = args.get('vehicle_model', '')
#     service = args.get('service_type', '')
#     date = args.get('preferred_date', '')
#     phone = args.get('phone_number', '')
#
#     # Call your booking API
#     resp = requests.post('https://your-api.com/api/bookings', json={
#         'customer_name': customer,
#         'vehicle_model': vehicle,
#         'service_type': service,
#         'preferred_date': date,
#         'phone_number': phone,
#     }, timeout=15)
#
#     if resp.ok:
#         return json.dumps(resp.json())
#     return f'Booking failed: {resp.status_code}'