import json
import hashlib
import os

# -------------------------
# CONFIG
# -------------------------
TARGET = 2500
OUTPUT_DIR = "finance_kb_shards"

# -------------------------
# DOMAIN + REGULATORY MAPPING
# -------------------------
DOMAIN_MAP = {
    "stocks": "stocks",
    "equity": "stocks",
    "crypto": "crypto",
    "insurance": "insurance",
    "iul": "insurance",
    "portfolio": "portfolio",
    "trading": "trading",
    "risk": "risk",
    "tax": "tax",
    "macro": "macro"
}

REGULATORY_SCOPE = {
    "stocks": ["SEC", "FINRA"],
    "equity": ["SEC"],
    "crypto": ["SEC", "CFTC"],
    "insurance": ["NAIC", "State Insurance Regulators"],
    "iul": ["NAIC"],
    "portfolio": ["SEC"],
    "trading": ["SEC", "FINRA"],
    "risk": ["SEC"],
    "tax": ["IRS"],
    "macro": ["Federal Reserve"]
}

SOURCE_MAP = {
    "stocks": "SEC Investor.gov",
    "equity": "SEC",
    "crypto": "CFTC / SEC Guidance",
    "insurance": "NAIC",
    "iul": "Insurance Carrier Disclosures",
    "portfolio": "FINRA",
    "trading": "FINRA",
    "risk": "Investopedia / SEC",
    "tax": "IRS Publications",
    "macro": "Federal Reserve"
}

# -------------------------
# TOPICS
# -------------------------
TOPICS = {
    "stocks": ["dividends", "P/E ratio", "market cap", "earnings reports"],
    "equity": ["growth stocks", "value stocks", "blue-chip stocks"],
    "crypto": ["bitcoin", "ethereum", "staking", "DeFi"],
    "insurance": ["life insurance", "term life", "whole life"],
    "iul": ["indexed universal life", "cap rate", "floor rate"],
    "portfolio": ["diversification", "asset allocation", "rebalancing"],
    "trading": ["RSI", "MACD", "moving averages"],
    "risk": ["volatility", "drawdowns", "liquidity risk"],
    "tax": ["capital gains tax", "dividend tax", "wash sale rule"],
    "macro": ["inflation", "interest rates", "recession"]
}

PERSONAS = [
    "beginner", "retail investor", "day trader",
    "long-term investor", "high net worth investor"
]

MARKET_CONTEXT = [
    "bull market", "bear market",
    "high inflation", "rising interest rates",
    "recession environment"
]

QUESTION_PATTERNS = [
    "What is {topic} for a {persona} during a {context}?",
    "How does {topic} impact a {persona} in a {context}?",
    "What are risks of {topic} for a {persona} in a {context}?",
    "How should a {persona} use {topic} in a {context}?"
]

# -------------------------
# HELPERS
# -------------------------
def make_id(question):
    return hashlib.sha1(question.encode()).hexdigest()[:12]

def generate_answer(topic, category, persona, context):
    return f"""
{topic.title()} is a key concept in {category}.

Explanation:
- Relevant for {persona}s
- Behavior varies in {context}
- Influences financial decisions

Example:
A {persona} evaluating {topic} during a {context} adjusts allocation and risk.

Compliance Note:
Always follow regulatory guidelines applicable to {category}.

Risk:
Improper use may lead to losses.

Tip:
Align with long-term financial goals.
""".strip()

# -------------------------
# GENERATION
# -------------------------
data = []
seen = set()

categories = list(TOPICS.keys())

i = 0
while len(data) < TARGET:
    category = categories[i % len(categories)]
    topic = TOPICS[category][i % len(TOPICS[category])]
    persona = PERSONAS[(i // 10) % len(PERSONAS)]
    context = MARKET_CONTEXT[(i // 50) % len(MARKET_CONTEXT)]
    pattern = QUESTION_PATTERNS[(i // 100) % len(QUESTION_PATTERNS)]

    question = pattern.format(topic=topic, persona=persona, context=context)

    if question not in seen:
        seen.add(question)

        entry = {
            "id": make_id(question),
            "domain": DOMAIN_MAP[category],
            "category": category,
            "topic": topic,
            "persona": persona,
            "market_context": context,
            "question": question,
            "answer": generate_answer(topic, category, persona, context),

            # NEW FIELDS
            "source": SOURCE_MAP[category],
            "regulatory_scope": REGULATORY_SCOPE[category],
            "tax_year": 2025 if category == "tax" else None,

            "keywords": [topic, category, persona],
            "difficulty": ["beginner", "intermediate", "advanced"][i % 3]
        }

        data.append(entry)

    i += 1

# -------------------------
# SHARD BY DOMAIN
# -------------------------
os.makedirs(OUTPUT_DIR, exist_ok=True)

shards = {}
for d in data:
    domain = d["domain"]
    shards.setdefault(domain, []).append(d)

for domain, items in shards.items():
    with open(f"{OUTPUT_DIR}/{domain}_qa.json", "w") as f:
        json.dump(items, f, indent=2)

# Master file
with open(f"{OUTPUT_DIR}/finance_qa_master.json", "w") as f:
    json.dump(data, f, indent=2)

print(f"Generated {len(data)} entries across {len(shards)} domains!")