#!/usr/bin/env python3
"""
scripts/omniroute_freight_synthesizer.py — OmniRoute 'paid-premium' Multi-Agent Freight Dataset Synthesizer.

Generates hyper-realistic multi-turn freight negotiation training dialogues using OmniRoute's
high-reasoning 'paid-premium' combo model.

Output complies with:
- OpenAI tool-calling chat completions format
- LoadETA tool schemas (fmcsa_verify_mc, calculate_rate_floor, get_truck_route_distance, book_load_offer)
- LiveKit voice conversational requirements (phonetic currencies, spoken cadence, no markdown)

Usage:
    python3 scripts/omniroute_freight_synthesizer.py --count 10 --output data/freight_negotiation_omniroute.jsonl
"""

import os
import sys
import json
import random
import time
import argparse
import urllib.request
import urllib.error
from typing import Dict, Any, List

OMNIROUTE_URL = "http://100.70.158.21:20128/v1/chat/completions"
OMNIROUTE_MODEL = "paid-premium"

SEED_SCENARIOS = [
    {
        "lane": "Atlanta, GA to Chicago, IL",
        "miles": 715,
        "equipment": "53ft Reefer",
        "temp": "-10F continuous",
        "commodity": "Frozen Poultry",
        "weight": "43,000 lbs",
        "broker_persona": "Aggressive lowball broker trying to anchor rate at $1,800 on a $2,400 lane",
        "special_condition": "Requires strict 2 hours free detention agreement and temp recorder verification",
    },
    {
        "lane": "Laredo, TX to Dallas, TX",
        "miles": 430,
        "equipment": "53ft Dry Van",
        "temp": "ambient",
        "commodity": "Cross-border Automotive Parts",
        "weight": "42,500 lbs",
        "broker_persona": "Urgent hot-load broker with auto assembly line shutdown risk if not picked up in 90 mins",
        "special_condition": "High rate elasticity, dispatcher pushes for premium rate ($1,650)",
    },
    {
        "lane": "Allentown, PA to Richmond, VA",
        "miles": 260,
        "equipment": "48ft Flatbed",
        "temp": "ambient",
        "commodity": "Structural Steel & Beams with 8ft Tarp",
        "weight": "46,000 lbs",
        "broker_persona": "Mid-tier broker trying to avoid paying $150 tarp fee",
        "special_condition": "Dispatcher stands firm on tarping surcharge and Northeast toll reimbursement",
    },
    {
        "lane": "Los Angeles, CA to Phoenix, AZ",
        "miles": 375,
        "equipment": "53ft Reefer",
        "temp": "+34F pre-cooled",
        "commodity": "Fresh Organic Berries",
        "weight": "38,000 lbs",
        "broker_persona": "Double-broker suspect with newly registered MC (42 days old)",
        "special_condition": "FMCSA tool flags new authority and lack of credit score; dispatcher politely refuses load",
    },
    {
        "lane": "Houston, TX to Savannah, GA",
        "miles": 840,
        "equipment": "53ft Dry Van",
        "temp": "ambient",
        "commodity": "Retail Consumer Goods",
        "weight": "36,000 lbs",
        "broker_persona": "Standard C.H. Robinson corporate broker negotiating multi-stop delivery",
        "special_condition": "2 pick-up locations and 2 drop-offs; dispatcher negotiates $100 extra per stop",
    },
]


def get_omniroute_api_key() -> str:
    key = os.environ.get("OMNIROUTE_API_KEY")
    if key:
        return key
    secrets_path = os.path.expanduser("~/.mcp/secrets/omniroute.env")
    if os.path.exists(secrets_path):
        with open(secrets_path, "r", encoding="utf-8") as f:
            for line in f:
                if line.startswith("OMNIROUTE_MCP_KEY="):
                    return line.strip().split("=", 1)[1]
    return ""


def call_omniroute(prompt: str, system: str = "") -> str:
    key = get_omniroute_api_key()
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": OMNIROUTE_MODEL,
        "messages": [
            {"role": "system", "content": system or "You are an expert synthetic dataset engineer for autonomous freight negotiation voice agents."},
            {"role": "user", "content": prompt},
        ],
        "temperature": 0.7,
        "max_tokens": 2500,
    }

    req = urllib.request.Request(
        OMNIROUTE_URL,
        data=json.dumps(payload).encode("utf-8"),
        headers=headers,
    )
    with urllib.request.urlopen(req, timeout=120) as resp:
        res = json.loads(resp.read().decode("utf-8"))
        return res["choices"][0]["message"]["content"]


def extract_json(raw_text: str) -> Dict[str, Any]:
    import re
    # 1. Try markdown code blocks first (often the cleanest JSON)
    blocks = re.findall(r"```(?:json)?\s*([\s\S]*?)\s*```", raw_text)
    for block in reversed(blocks):
        try:
            val = json.loads(block.strip())
            if isinstance(val, dict) and ("messages" in val or "turns" in val or "id" in val):
                return val
        except Exception:
            pass

    for block in blocks:
        try:
            val = json.loads(block.strip())
            if isinstance(val, dict):
                return val
        except Exception:
            pass

    # 2. Try whole string directly
    try:
        return json.loads(raw_text.strip())
    except Exception:
        pass

    # 3. Scan from first { to last }
    s = raw_text.find("{")
    e = raw_text.rfind("}")
    if s != -1 and e != -1 and e > s:
        try:
            return json.loads(raw_text[s:e+1])
        except Exception:
            pass

    raise ValueError(f"No valid JSON object could be extracted from: {raw_text[:120]}")


def synthesize_conversation(scenario: Dict[str, Any], index: int) -> Dict[str, Any]:
    prompt = f"""
Generate an authentic, complete multi-turn freight negotiation training conversation between a US Freight Broker and LoadETA (an autonomous dispatcher voice agent).

Scenario Parameters:
- Lane: {scenario['lane']} ({scenario['miles']} miles)
- Equipment: {scenario['equipment']}
- Commodity & Weight: {scenario['commodity']} ({scenario['weight']})
- Temperature: {scenario['temp']}
- Broker Profile: {scenario['broker_persona']}
- Special Case: {scenario['special_condition']}

Strict Rules for LoadETA Spoken Output:
1. Format exclusively for human voice synthesis (LiveKit TTS): No markdown bolding, no bullet points, no emojis, no robotic phrasing.
2. Spoken numbers and currencies: use natural spoken words for money (e.g. "twenty-four hundred", "eighteen fifty all in", "seventy-five dollars an hour").
3. Tool Calling: Assistant MUST include tool calls with exact arguments:
   - 'fmcsa_verify_mc' with arguments {{"mc_number": "...", "broker_name": "..."}}
   - 'calculate_rate_floor' with arguments {{"origin": "...", "destination": "...", "equipment_type": "...", "mileage": {scenario['miles']}, "weight": 42000}}
   - 'book_load_offer' with arguments {{"broker_name": "...", "mc_number": "...", "rate_agreed": 2400, "detention_rate_per_hr": 75}}

You must output a JSON object with this exact key structure:
{{
  "id": "omniroute_freight_{index:05d}",
  "messages": [
    {{"role": "system", "content": "You are LoadETA, an expert freight dispatcher and live negotiation voice agent for US commercial trucking. Protect driver margins, verify broker authority, and negotiate rates naturally for voice without markdown."}},
    {{"role": "user", "content": "Broker opening pitch..."}},
    {{"role": "assistant", "content": null, "tool_calls": [{{"id": "call_fmcsa_{index:04d}", "type": "function", "function": {{"name": "fmcsa_verify_mc", "arguments": "{{\\"mc_number\\": \\"782914\\", \\"broker_name\\": \\"Apex Logistics\\"}}"}}}}]}},
    {{"role": "tool", "tool_call_id": "call_fmcsa_{index:04d}", "content": "{{\\"status\\": \\"CLEAN\\", \\"active\\": true, \\"credit_score\\": 94}}"}}
    {{"role": "assistant", "content": "Natural counter-offer in spoken voice..."}},
    {{"role": "user", "content": "Broker counter..."}},
    {{"role": "assistant", "content": null, "tool_calls": [{{"id": "call_book_{index:04d}", "type": "function", "function": {{"name": "book_load_offer", "arguments": "{{\\"rate_agreed\\": 2400}}"}}}}]}},
    {{"role": "tool", "tool_call_id": "call_book_{index:04d}", "content": "{{\\"status\\": \\"BOOKED\\"}}"}},
    {{"role": "assistant", "content": "Final confirmation..."}}
  ]
}}

OUTPUT JSON ONLY. Do not write explanations or notes.
"""

    raw_response = call_omniroute(prompt)
    try:
        data = extract_json(raw_response)
        return data
    except Exception as e:
        print(f"JSON decode failed for scenario {index}: {e}. Raw head: {repr(raw_response[:100])}")
        return None


def main():
    parser = argparse.ArgumentParser(description="OmniRoute Paid-Premium Freight Dataset Synthesizer")
    parser.add_argument("--count", type=int, default=5, help="Number of dialogues to synthesize")
    parser.add_argument("--output", default="freight/data/freight_negotiation_omniroute.jsonl", help="Output JSONL file")
    args = parser.parse_args()

    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    print(f"Connecting to OmniRoute ({OMNIROUTE_URL}) with model '{OMNIROUTE_MODEL}'...")

    results = []
    for i in range(1, args.count + 1):
        scenario = random.choice(SEED_SCENARIOS)
        print(f"[{i}/{args.count}] Synthesizing: {scenario['lane']} ({scenario['equipment']})...")
        conv = synthesize_conversation(scenario, i)
        if conv:
            results.append(conv)
            with open(args.output, "a" if i > 1 else "w", encoding="utf-8") as f:
                f.write(json.dumps(conv) + "\n")
            print(f"  -> Generated {len(conv.get('messages', []))} turns.")
        time.sleep(1)

    print(f"\nSuccessfully generated {len(results)} high-reasoning dialogues in {args.output}")


if __name__ == "__main__":
    main()