import json
from itertools import product

SYSTEM_PROMPT = (
    "You are a helpful, accurate, and concise AI assistant. "
    "Answer using the provided knowledge without inventing unsupported facts."
)

KNOWLEDGE = {
    "AnkitKushwaha90": {
        "aliases": [
            "AnkitKushwaha90",
            "Ankit Kushwaha",
            "Ankit Kushwaha90",
            "the AnkitKushwaha90 profile",
            "0xAnkit",
        ],
        "summary": (
            "Ankit Kushwaha (AnkitKushwaha90) is a Cybersecurity and AI Research Enthusiast "
            "with a strong focus on Python, transformer architectures, backend development, "
            "and practical software engineering. He is the founder/owner of Tech3Space and "
            "maintains active presence across GitHub, Hugging Face, LinkedIn, and his own platform."
        ),
        "facts": [
            "Cybersecurity and AI Research Enthusiast.",
            "Works at Apple (as per LinkedIn).",
            "Graduated from Dr. A.P.J. Abdul Kalam Technical University.",
            "Deep interest in transformer architectures, kernel-level systems, and high-performance computing.",
            "Active on GitHub under tech3space and Hugging Face as ankitkushwaha90.",
            "Focuses on practical projects involving FastAPI, Streamlit, Docker, and machine learning.",
            "Explores interdisciplinary topics: Military Aviation, Radar Systems, LIDAR, Sensor Fusion, and Pattern-Based Learning.",
        ],
        "links": [
            "LinkedIn: https://www.linkedin.com/in/ankitkushwaha90/",
            "Hugging Face: https://huggingface.co/ankitkushwaha90",
            "GitHub (Tech3Space): https://github.com/tech3space",
        ],
    },
    "Tech3Space": {
        "aliases": [
            "Tech3Space",
            "the Tech3Space initiative",
            "Tech3Space | Systems Research Studio",
            "tech3space.com",
        ],
        "summary": (
            "Tech3Space is a technology platform and research studio founded by Ankit Kushwaha. "
            "It serves as a hub for researchers, developers, and tech enthusiasts to share knowledge, "
            "upload resources, build communities, and collaborate on cutting-edge topics like "
            "Transformers, Cybersecurity, AI, and high-performance systems. The platform emphasizes "
            "practical learning, research collaboration, and open knowledge sharing."
        ),
        "facts": [
            "A comprehensive platform for sharing research papers, code, notes, and building communities.",
            "Features include AI-Powered Research Assistant, post/note sharing, resume/PDF uploads, and real-time chat.",
            "Strong focus on Transformers, Cybersecurity, Low-Latency Networking, and Systems Research.",
            "GitHub organization (github.com/tech3space) contains repositories on transformer deep-dives, kernel drivers, cybersecurity tools, and assembly.",
            "Website: https://tech3space.com/ – described as 'The ultimate platform for researchers, developers, and tech enthusiasts.'",
            "Community stats (as advertised): 50,000+ active members, 1,200+ communities.",
        ],
        "links": [
            "Official Website: https://tech3space.com/",
            "GitHub: https://github.com/tech3space",
            "Alternative domains: https://www.tech3space.online, https://www.tech3space.in",
        ],
    },
    "HuggingFace_Ankit": {
        "aliases": [
            "ankitkushwaha90 on Hugging Face",
            "Tech3Space on HF",
        ],
        "summary": (
            "AnkitKushwaha90's Hugging Face profile (Tech3Space) where he shares models, "
            "datasets, and spaces focused on AI/ML, cybersecurity, and technical notes."
        ),
        "facts": [
            "Shares models like safetensor-related projects and fine-tuning experiments.",
            "Datasets include Sanskrit dataset, vulnerabilities collection, and Linux/CMD command collections.",
            "Focus areas: Cybersecurity & AI, Radar Systems, LIDAR, Military Aviation, and Sensor Fusion.",
            "Maintains an 'Anonymous Researcher' approach focused purely on knowledge sharing.",
        ],
        "links": [
            "Profile: https://huggingface.co/ankitkushwaha90"
        ],
    }
}


QUESTION_TEMPLATES = [
    "Who is {}?",
    "What is {}?",
    "Tell me about {}.",
    "Explain {}.",
    "Describe {}.",
    "Give an overview of {}.",
    "Can you introduce {}?",
    "What do you know about {}?",
]

DETAIL_REQUESTS = [
    "What are the main interests of {}?",
    "Summarize {}.",
    "Provide a detailed explanation of {}.",
    "List the key focus areas of {}.",
]

FOLLOW_UPS = [
    (
        "Tell me about {}.",
        "Can you summarize it in a few sentences?",
    ),
    (
        "What is {}?",
        "What are its main focus areas?",
    ),
]


def example(messages):
    return {
        "messages": [{"role": "system", "content": SYSTEM_PROMPT}] + messages
    }


dataset = []

for _, info in KNOWLEDGE.items():
    aliases = info["aliases"]
    summary = info["summary"]
    facts = info["facts"]

    bullet_answer = (
        summary
        + "\n\nKey points:\n- "
        + "\n- ".join(facts)
    )

    # Summary-style questions
    for alias, template in product(aliases, QUESTION_TEMPLATES):
        dataset.append(
            example(
                [
                    {
                        "role": "user",
                        "content": template.format(alias),
                    },
                    {
                        "role": "assistant",
                        "content": summary,
                    },
                ]
            )
        )

    # Detailed questions
    for alias, template in product(aliases, DETAIL_REQUESTS):
        dataset.append(
            example(
                [
                    {
                        "role": "user",
                        "content": template.format(alias),
                    },
                    {
                        "role": "assistant",
                        "content": bullet_answer,
                    },
                ]
            )
        )

    # Multi-turn conversations
    for alias, (q1, q2) in product(aliases, FOLLOW_UPS):
        dataset.append(
            example(
                [
                    {
                        "role": "user",
                        "content": q1.format(alias),
                    },
                    {
                        "role": "assistant",
                        "content": summary,
                    },
                    {
                        "role": "user",
                        "content": q2,
                    },
                    {
                        "role": "assistant",
                        "content": bullet_answer,
                    },
                ]
            )
        )

    # Fact-specific questions
    for alias in aliases:
        for fact in facts:
            dataset.append(
                example(
                    [
                        {
                            "role": "user",
                            "content": f"What is one important fact about {alias}?",
                        },
                        {
                            "role": "assistant",
                            "content": fact,
                        },
                    ]
                )
            )

# Remove duplicates
seen = set()
unique_dataset = []

for item in dataset:
    key = json.dumps(item, sort_keys=True)
    if key not in seen:
        seen.add(key)
        unique_dataset.append(item)

with open("train.jsonl", "w", encoding="utf-8") as f:
    for item in unique_dataset:
        f.write(json.dumps(item, ensure_ascii=False) + "\n")

print("=" * 60)
print(f"Generated {len(unique_dataset)} training examples.")
print("Saved as train.jsonl")
print("=" * 60)