"""
corpus_generator.py  —  Generate Tamil domain corpus using Gemini Pro

Usage
-----
    python corpus_generator.py --api_key YOUR_KEY
    python corpus_generator.py --api_key YOUR_KEY --domains politics cricket
    python corpus_generator.py --api_key YOUR_KEY --words_per_domain 1500

Output
------
  data/<domain>/gemini_<domain>.txt   (one file per domain)

After running, execute prepare_corpus.py as usual to merge into
data/tamil_corpus.txt, then retrain FastText.
"""

import os
import time
import argparse
import requests
import json

# ---------------------------------------------------------------------------
# CONFIG
# ---------------------------------------------------------------------------

GEMINI_URL = (
    "https://generativelanguage.googleapis.com/v1beta/models/"
    "gemini-2.0-flash:generateContent"
)

OUTPUT_BASE = "data"
WORDS_PER_DOMAIN = 1000
SLEEP_BETWEEN_CALLS = 2   # seconds — stay within free-tier rate limits

# ---------------------------------------------------------------------------
# DOMAINS — mirrors the cluster groups in cluster_store.py
# ---------------------------------------------------------------------------

DOMAINS = {
    "politics": (
        "தமிழ்நாட்டு அரசியல், தேர்தல் முறை, சட்டமன்றம், நாடாளுமன்றம், "
        "அரசியல் கட்சிகள், கூட்டணிகள், வாக்காளர்கள், அரசமைப்பு, "
        "இட ஒதுக்கீடு, சமூக நீதி, ஜனநாயகம், அரசு நிர்வாகம்"
    ),
    "history": (
        "தமிழ் வரலாறு, சோழர் சேரர் பாண்டியர் ஆட்சி, சங்க காலம், "
        "விஜயநகர பேரரசு, முகலாயர், காலனித்துவம், சுதந்திர போராட்டம், "
        "தமிழக வரலாற்று நிகழ்வுகள், கல்வெட்டுகள், தொல்லியல்"
    ),
    "cricket": (
        "இந்திய கிரிக்கெட், IPL, உலகக்கோப்பை, டெஸ்ட் ஒருநாள் T20 போட்டிகள், "
        "பிரபல வீரர்கள், பேட்டிங் பவுலிங் புலத்தடுப்பு, தமிழ்நாடு கிரிக்கெட், "
        "பந்து வகைகள், ஷாட் வகைகள், கிரிக்கெட் விதிகள்"
    ),
    "cinema": (
        "தமிழ் திரைப்படம், கோலிவுட், நடிகர் நடிகையர், இயக்குநர்கள், "
        "இசையமைப்பாளர்கள், விருதுகள், திரை விமர்சனம், படத்தயாரிப்பு, "
        "பாடல்கள், வசனம், திரை வரலாறு"
    ),
    "science": (
        "விஞ்ஞானம் தொழில்நுட்பம், இயற்பியல் வேதியியல் உயிரியல், "
        "விண்வெளி ஆராய்ச்சி, ISRO, அணு ஆற்றல், கணினி அறிவியல், "
        "செயற்கை நுண்ணறிவு, மருத்துவ ஆராய்ச்சி, சுற்றுச்சூழல் அறிவியல்"
    ),
    "education": (
        "தமிழ்நாடு கல்வி முறை, பள்ளி கல்லூரி பல்கலைக்கழகம், "
        "தேர்வு முறை, உதவித்தொகை, ஆசிரியர் பயிற்சி, "
        "தொலைதூர கல்வி, தொழிற்கல்வி, மதிப்பெண் முறை"
    ),
    "health": (
        "மருத்துவம் சுகாதாரம், நோய்கள் மற்றும் சிகிச்சை, "
        "தடுப்பூசி, ஆயுர்வேதம் சித்த மருத்துவம், "
        "மனநல ஆரோக்கியம், உணவு ஊட்டம், பொது சுகாதாரம், "
        "மருத்துவமனை நிர்வாகம்"
    ),
    "environment": (
        "சுற்றுச்சூழல் பாதுகாப்பு, காலநிலை மாற்றம், மரம் வளர்ப்பு, "
        "நீர் ஆதாரம் நிர்வாகம், கடல் சூழல், வனவிலங்கு பாதுகாப்பு, "
        "மாசுபாடு கட்டுப்பாடு, புதுப்பிக்கத்தக்க ஆற்றல்"
    ),
    "economics": (
        "இந்திய பொருளாதாரம், தமிழ்நாடு தொழில் வளர்ச்சி, "
        "வேளாண்மை பொருளாதாரம், ஏற்றுமதி இறக்குமதி, "
        "பட்ஜெட் திட்டமிடல், வங்கி முறை, சிறு தொழில்கள், "
        "பங்குச்சந்தை, பொருளாதார சீர்திருத்தம்"
    ),
    "literature": (
        "தமிழ் இலக்கியம், சங்க இலக்கியம், திருக்குறள், "
        "காப்பியங்கள், நவீன தமிழ் இலக்கியம், கவிஞர்கள், "
        "சிறுகதை நாவல் கவிதை, இலக்கண நூல்கள், "
        "தமிழ் செய்யுள் உரைநடை"
    ),
    "geography": (
        "தமிழ்நாடு மாவட்டங்கள் நகரங்கள், இந்திய மாநிலங்கள், "
        "ஆறுகள் மலைகள் சமவெளிகள், தமிழக வரைபடம், "
        "இந்திய புவியியல், உலக நாடுகள் தலைநகரங்கள்"
    ),
}

# ---------------------------------------------------------------------------
# PROMPT BUILDER
# ---------------------------------------------------------------------------

def build_prompt(domain: str, keywords: str, word_count: int) -> str:
    return f"""நீங்கள் ஒரு தமிழ் உள்ளடக்க எழுத்தாளர். கீழே கொடுக்கப்பட்ட தலைப்பில் சுமார் {word_count} தமிழ் வார்த்தைகள் கொண்ட ஒரு விரிவான கட்டுரை எழுதுங்கள்.

தலைப்பு: {domain.upper()}
முக்கிய தலைப்புகள்: {keywords}

வழிகாட்டுதல்கள்:
- தூய தமிழில் எழுதுங்கள் (ஆங்கில வார்த்தைகளை தவிர்க்கவும், தொழில்நுட்ப சொற்கள் மட்டும் ஆங்கிலத்தில் இருக்கலாம்)
- தெளிவான, நேரடியான வாக்கியங்கள் எழுதுங்கள்
- ஒவ்வொரு வாக்கியமும் முழுமையானதாகவும் சுயாதீனமானதாகவும் இருக்க வேண்டும்
- இது/அது/இவை/அவை போன்ற சுட்டுப்பெயர்களில் வாக்கியங்களை தொடங்காதீர்கள்
- குறிப்பிட்ட பெயர்கள் (நபர்கள், இடங்கள், நிறுவனங்கள், சட்டங்கள்) அடங்கிய வாக்கியங்கள் அதிகம் எழுதுங்கள்
- கட்டுரையை நேரடியாக தமிழிலேயே தொடங்குங்கள், எந்த ஆங்கில தலைப்பும் வேண்டாம்"""


# ---------------------------------------------------------------------------
# GEMINI API CALL
# ---------------------------------------------------------------------------

def call_gemini(prompt: str, api_key: str) -> str:
    headers = {"Content-Type": "application/json"}
    params  = {"key": api_key}
    body    = {
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": {
            "temperature": 0.7,
            "maxOutputTokens": 2048,
        },
    }

    resp = requests.post(GEMINI_URL, headers=headers, params=params, json=body, timeout=60)
    resp.raise_for_status()
    data = resp.json()

    try:
        return data["candidates"][0]["content"]["parts"][0]["text"]
    except (KeyError, IndexError) as e:
        raise ValueError(f"Unexpected Gemini response structure: {data}") from e


# ---------------------------------------------------------------------------
# MAIN
# ---------------------------------------------------------------------------

def generate_corpus(api_key: str, domains: list[str] | None = None,
                    words_per_domain: int = WORDS_PER_DOMAIN) -> None:

    selected = {k: v for k, v in DOMAINS.items()
                if domains is None or k in domains}

    print(f"[Corpus] Generating text for {len(selected)} domains "
          f"(~{words_per_domain} words each)\n")

    for domain, keywords in selected.items():
        out_dir  = os.path.join(OUTPUT_BASE, domain)
        out_file = os.path.join(out_dir, f"gemini_{domain}.txt")
        os.makedirs(out_dir, exist_ok=True)

        print(f"  [{domain}] Calling Gemini...", end=" ", flush=True)

        try:
            prompt = build_prompt(domain, keywords, words_per_domain)
            text   = call_gemini(prompt, api_key)

            with open(out_file, "w", encoding="utf-8") as f:
                f.write(text)

            word_count = len(text.split())
            print(f"✅  {word_count} words → {out_file}")

        except Exception as e:
            print(f"❌  Error: {e}")

        time.sleep(SLEEP_BETWEEN_CALLS)

    print("\n[Corpus] Done. Run prepare_corpus.py to merge into tamil_corpus.txt")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate Tamil domain corpus with Gemini")
    parser.add_argument("--api_key",          required=True,  help="Gemini API key")
    parser.add_argument("--domains",          nargs="+",      help="Specific domains (default: all)")
    parser.add_argument("--words_per_domain", type=int, default=WORDS_PER_DOMAIN,
                        help=f"Target word count per domain (default: {WORDS_PER_DOMAIN})")
    args = parser.parse_args()

    generate_corpus(
        api_key=args.api_key,
        domains=args.domains,
        words_per_domain=args.words_per_domain,
    )