"""
Synthetic data generator for TradeParse-LoRA / OrderIntent.

Generates clean instruction -> JSON training pairs for fine-tuning a model to
parse natural-language trade instructions into one consistent schema.

The model may output either:
- a successful trade parse
- a structured error when the instruction is missing required information

idempotency_key is intentionally excluded from model output. Generate it in
Python after validation if needed.
"""

import argparse
import json
import random
from pathlib import Path

RNG = random.Random(42)

SYSTEM_PROMPT = (
    "Convert the user trade instruction into JSON only. "
    "If the instruction has enough information, use exactly this success schema: "
    '{"symbol": str, "action_conditions": [{"trigger": "price_below"|"price_above", '
    '"value": number, "action": "buy"|"sell"}], "strategy_type": "intraday"|null}. '
    "If required information is missing or ambiguous, use exactly this error schema: "
    '{"error": "missing_symbol"|"ambiguous_symbol"|"missing_price_condition"|"missing_action"|"unsupported_strategy", '
    '"symbol": str|null, "message": str}. '
    "Rules: if strategy type is not mentioned, set strategy_type to null. "
    "Only intraday is supported as a strategy_type; futures, delivery, and swing should not be copied into strategy_type. "
    "If no numeric price condition is present, return missing_price_condition. "
    "If no symbol/company is present, return missing_symbol. "
    "If no buy/sell style action is present, return missing_action. "
    "Do not include markdown, explanations, or extra keys."
)

SYMBOL_ALIASES = {
    "INFY": ["INFY", "Infy", "Infosys", "INFY.NS"],
    "RELIANCE": ["RELIANCE", "Reliance"],
    "TCS": ["TCS", "Tcs", "Tata Consultancy Services"],
    "HDFC": ["HDFC", "Hdfc", "HDFC Bank"],
    "ICICIBANK": ["ICICIBANK", "Icicibank", "ICICI Bank"],
    "SBIN": ["SBIN", "Sbin", "SBI", "State Bank of India"],
    "WIPRO": ["WIPRO", "Wipro", "Wipro Ltd"],
    "TATASTEEL": ["TATASTEEL", "Tatasteel", "Tata Steel"],
    "BAJFINANCE": ["BAJFINANCE", "Bajfinance", "Bajaj Finance"],
    "ADANIENT": ["ADANIENT", "Adanient", "Adani Enterprises"],
    "MARUTI": ["MARUTI", "Maruti"],
    "SUNPHARMA": ["SUNPHARMA", "Sunpharma", "Sun Pharma"],
    "HCLTECH": ["HCLTECH", "Hcltech", "HCL Tech"],
    "AXISBANK": ["AXISBANK", "Axisbank", "Axis Bank"],
    "ONGC": ["ONGC", "Ongc"],
}

DOUBLE_TEMPLATES = [
    "Buy {name} shares if it drops below {below} and sell if it goes above {above}{strategy_suffix}",
    "Buy {name} when price is below {below}, sell above {above}{strategy_suffix}",
    "Buy {name} if price falls under {below}, sell above {above}{strategy_suffix}",
    "If {name} goes below {below} buy it, and if it crosses {above} sell it{strategy_suffix}",
    "{name} - buy under {below}, book profit above {above}{strategy_suffix}",
    "Get me into {name} below {below}, exit above {above}{strategy_sentence}",
    "Place buy order for {name} below {below} and sell order above {above}{strategy_suffix}",
]

BUY_ONLY_TEMPLATES = [
    "Buy {name} if it drops below {below}{strategy_suffix}",
    "Buy {name} when price is below {below}{strategy_suffix}",
    "Get {name} below {below}{strategy_suffix}",
    "buy {name} under {below}{strategy_suffix}",
    "Accumulate {name} if it falls below {below}{strategy_suffix}",
]

SELL_ONLY_TEMPLATES = [
    "Sell {name} if it goes above {above}{strategy_suffix}",
    "Sell {name} above {above}{strategy_suffix}",
    "Exit {name} once it crosses {above}{strategy_suffix}",
    "Book profit in {name} above {above}{strategy_suffix}",
]

MISSING_PRICE_TEMPLATES = [
    "Buy {name} when it is cheap, sell when it is expensive",
    "Pick up {name} at a good level and exit higher",
    "Trade {name} based on support and resistance",
    "Buy {name} when it looks attractive",
]

MISSING_SYMBOL_TEMPLATES = [
    "buy this stock below {below}, sell above {above}{strategy_suffix}",
    "enter below {below} and exit above {above}",
    "buy it under {below}, sell it over {above}",
]

MISSING_ACTION_TEMPLATES = [
    "{name} at {price}{strategy_suffix}",
    "Watch {name} near {price}",
    "Alert me about {name} around {price}",
]

AMBIGUOUS_SYMBOL_TEMPLATES = [
    "Buy Tata below {below}",
    "Sell Tata above {above}",
    "Tata below {below} buy, above {above} sell",
]


def choose_symbol():
    symbol = RNG.choice(list(SYMBOL_ALIASES.keys()))
    name = RNG.choice(SYMBOL_ALIASES[symbol])
    return symbol, name


def choose_strategy_suffix(include_strategy: bool):
    if not include_strategy:
        return "", ""
    return ", intraday", ". intraday trade."


def output(symbol, conditions, strategy_type):
    return {
        "symbol": symbol,
        "action_conditions": conditions,
        "strategy_type": strategy_type,
    }


def error(error_type, symbol, message):
    return {
        "error": error_type,
        "symbol": symbol,
        "message": message,
    }


def gen_double_condition_example():
    symbol, name = choose_symbol()
    below = RNG.randint(100, 4000)
    above = below + RNG.randint(50, 500)
    include_strategy = RNG.random() < 0.6
    strategy_suffix, strategy_sentence = choose_strategy_suffix(include_strategy)
    template = RNG.choice(DOUBLE_TEMPLATES)

    instruction = template.format(
        name=name,
        below=below,
        above=above,
        strategy_suffix=strategy_suffix,
        strategy_sentence=strategy_sentence,
    ).strip()

    return instruction, output(
        symbol,
        [
            {"trigger": "price_below", "value": below, "action": "buy"},
            {"trigger": "price_above", "value": above, "action": "sell"},
        ],
        "intraday" if include_strategy else None,
    )


def gen_buy_only_example():
    symbol, name = choose_symbol()
    below = RNG.randint(100, 4000)
    include_strategy = RNG.random() < 0.5
    strategy_suffix, _ = choose_strategy_suffix(include_strategy)
    template = RNG.choice(BUY_ONLY_TEMPLATES)
    instruction = template.format(name=name, below=below, strategy_suffix=strategy_suffix).strip()

    return instruction, output(
        symbol,
        [{"trigger": "price_below", "value": below, "action": "buy"}],
        "intraday" if include_strategy else None,
    )


def gen_sell_only_example():
    symbol, name = choose_symbol()
    above = RNG.randint(100, 4000)
    include_strategy = RNG.random() < 0.5
    strategy_suffix, _ = choose_strategy_suffix(include_strategy)
    template = RNG.choice(SELL_ONLY_TEMPLATES)
    instruction = template.format(name=name, above=above, strategy_suffix=strategy_suffix).strip()

    return instruction, output(
        symbol,
        [{"trigger": "price_above", "value": above, "action": "sell"}],
        "intraday" if include_strategy else None,
    )


def gen_error_example():
    kind = RNG.choice(["missing_price", "missing_symbol", "missing_action", "ambiguous_symbol"])
    below = RNG.randint(100, 4000)
    above = below + RNG.randint(50, 500)
    strategy_suffix, _ = choose_strategy_suffix(RNG.random() < 0.3)

    if kind == "missing_price":
        symbol, name = choose_symbol()
        instruction = RNG.choice(MISSING_PRICE_TEMPLATES).format(name=name)
        return instruction, error("missing_price_condition", symbol, "No numeric price condition found.")

    if kind == "missing_symbol":
        instruction = RNG.choice(MISSING_SYMBOL_TEMPLATES).format(
            below=below,
            above=above,
            strategy_suffix=strategy_suffix,
        )
        return instruction, error("missing_symbol", None, "No symbol or company name found.")

    if kind == "missing_action":
        symbol, name = choose_symbol()
        instruction = RNG.choice(MISSING_ACTION_TEMPLATES).format(
            name=name,
            price=below,
            strategy_suffix=strategy_suffix,
        )
        return instruction, error("missing_action", symbol, "No buy or sell action found.")

    instruction = RNG.choice(AMBIGUOUS_SYMBOL_TEMPLATES).format(below=below, above=above)
    return instruction, error("ambiguous_symbol", "TATA", "Symbol is ambiguous. Specify the exact Tata company.")


def fixed_examples():
    """High-value examples that exactly match likely manual tests."""
    return [
        (
            "Buy Infy when price below 1500 and sell above 1700, intraday",
            output("INFY", [
                {"trigger": "price_below", "value": 1500, "action": "buy"},
                {"trigger": "price_above", "value": 1700, "action": "sell"},
            ], "intraday"),
        ),
        (
            "Buy INFY if it drops below 1500, sell above 1700, intraday",
            output("INFY", [
                {"trigger": "price_below", "value": 1500, "action": "buy"},
                {"trigger": "price_above", "value": 1700, "action": "sell"},
            ], "intraday"),
        ),
        (
            "Buy Infosys if it drops below 1500, sell above 1700, intraday",
            output("INFY", [
                {"trigger": "price_below", "value": 1500, "action": "buy"},
                {"trigger": "price_above", "value": 1700, "action": "sell"},
            ], "intraday"),
        ),
        (
            "Buy Infy when price below 1500",
            output("INFY", [
                {"trigger": "price_below", "value": 1500, "action": "buy"},
            ], None),
        ),
        (
            "Buy ONGC when it's cheap, sell when it's expensive",
            error("missing_price_condition", "ONGC", "No numeric price condition found."),
        ),
        (
            "buy this stock below 500, sell above 600, futures",
            error("missing_symbol", None, "No symbol or company name found."),
        ),
        (
            "Do something with Infosys at 1500",
            error("missing_action", "INFY", "No buy or sell action found."),
        ),
        (
            "Buy Tata below 500",
            error("ambiguous_symbol", "TATA", "Symbol is ambiguous. Specify the exact Tata company."),
        ),
    ]


def to_chat_format(instruction, output_json):
    return {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": instruction},
            {"role": "assistant", "content": json.dumps(output_json, ensure_ascii=False)},
        ]
    }


def generate_dataset(n_total=1600):
    examples = fixed_examples()

    while len(examples) < n_total:
        r = RNG.random()
        if r < 0.52:
            examples.append(gen_double_condition_example())
        elif r < 0.68:
            examples.append(gen_buy_only_example())
        elif r < 0.84:
            examples.append(gen_sell_only_example())
        else:
            examples.append(gen_error_example())

    RNG.shuffle(examples)
    return [to_chat_format(instruction, parsed) for instruction, parsed in examples]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--n", type=int, default=1600, help="Total examples to generate")
    parser.add_argument("--train-split", type=float, default=0.85)
    parser.add_argument("--out-dir", type=str, default="data_v2")
    args = parser.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    data = generate_dataset(args.n)
    split_idx = int(len(data) * args.train_split)
    train_data = data[:split_idx]
    valid_data = data[split_idx:]

    train_path = out_dir / "train.jsonl"
    valid_path = out_dir / "valid.jsonl"

    with train_path.open("w", encoding="utf-8") as f:
        for row in train_data:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    with valid_path.open("w", encoding="utf-8") as f:
        for row in valid_data:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    print(f"Generated {len(train_data)} training examples -> {train_path}")
    print(f"Generated {len(valid_data)} validation examples -> {valid_path}")
    print(json.dumps(train_data[0], indent=2))


if __name__ == "__main__":
    main()