import os
import csv
import json
import asyncio
from openai import AsyncOpenAI
from dotenv import load_dotenv

load_dotenv()

# Configure Teacher Model
AI_API_KEY = os.getenv("QWEN_API_KEY", "your-api-key-here")
AI_BASE_URL = os.getenv("QWEN_BASE_URL", "https://dashscope-intl.aliyuncs.com/compatible-mode/v1")
TEACHER_MODEL_NAME = os.getenv("TEACHER_MODEL_NAME", "qwen3-coder-480b-a35b")

client = AsyncOpenAI(
    api_key=AI_API_KEY,
    base_url=AI_BASE_URL
)

IEDID_DIR = "IEDID"
OUTPUT_FILE = "synthetic_slm_dataset.jsonl"

async def generate_synthetic_data(dialect, utterance, clarification):
    system_prompt = (
        f"You are a master linguist specializing in {dialect} and data augmentation for SLM fine-tuning.\n"
        f"I will provide a human-verified conversation pair.\n"
        f"Dialect Utterance: '{utterance}'\n"
        f"Standard English Meaning: '{clarification}'\n\n"
        f"Generate 3 highly accurate, synthetically varied conversation pairs that capture the exact same "
        f"linguistic rules, phonetic syntax, and cultural tone as the example.\n"
        f"Respond ONLY in valid JSON array format, where each object has 'prompt' (Standard English) and 'completion' ({dialect})."
    )
    
    try:
        response = await client.chat.completions.create(
            model=TEACHER_MODEL_NAME,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": "Generate 3 synthetic pairs in a JSON array."}
            ],
            temperature=0.7,
            max_tokens=1024
        )
        
        content = response.choices[0].message.content.strip()
        # Fallback parsing if wrapped in markdown
        if content.startswith("```json"):
            content = content[7:-3]
        elif content.startswith("```"):
            content = content[3:-3]
            
        data = json.loads(content)
        # Expected format: [{"prompt": "...", "completion": "..."}, ...]
        
        # If wrapped in a key
        if isinstance(data, dict):
            for key in data:
                if isinstance(data[key], list):
                    return data[key]
            return []
            
        return data if isinstance(data, list) else []
    except Exception as e:
        print(f"Error generating data for '{utterance}': {e}")
        return []

async def process_csv(filepath):
    print(f"Processing {filepath}...")
    dataset = []
    
    with open(filepath, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        rows = list(reader)[:10] # Limit to first 10 per file for demonstration/testing speed
        
        for row in rows:
            utterance = row.get('Utterance')
            clarification = row.get('Clarification')
            dialect = row.get('Dialect', 'Unknown Dialect')
            
            if not utterance or not clarification:
                continue
                
            print(f"  Distilling variations for: {utterance}")
            synthetic_pairs = await generate_synthetic_data(dialect, utterance, clarification)
            
            for pair in synthetic_pairs:
                if "prompt" in pair and "completion" in pair:
                    dataset.append({
                        "system": f"You are an expert polyglot interpreter. Translate Standard English to {dialect}.",
                        "user": pair["prompt"],
                        "assistant": pair["completion"]
                    })
                    
    return dataset

PRIORITY_DIALECTS = [
    "African American Vernacular English",
    "Nigerian Pidgin English",
    "Gyeongsang Satoori",
    "Jeolla Satoori",
    "Chungcheong Satoori",
    "Jeju Satoori",
    "Gangwon Satoori",
    "Hamgyong Satoori"
]

async def main():
    print(f"Starting Knowledge Distillation Pipeline (Teacher Model: {TEACHER_MODEL_NAME})...")
    all_synthetic_data = []
    
    # Read IEDID CSVs
    if os.path.exists(IEDID_DIR):
        for filename in os.listdir(IEDID_DIR):
            # Only process if it's a CSV, not minted, and matches a priority dialect
            is_priority = any(dialect in filename for dialect in PRIORITY_DIALECTS)
            if filename.endswith(".csv") and "minted" not in filename and is_priority:
                filepath = os.path.join(IEDID_DIR, filename)
                data = await process_csv(filepath)
                all_synthetic_data.extend(data)
                
    # Save to JSONL
    with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
        for entry in all_synthetic_data:
            f.write(json.dumps(entry) + '\n')
            
    print(f"\nSuccessfully generated {len(all_synthetic_data)} synthetic training pairs!")
    print(f"Saved to {OUTPUT_FILE}. Ready for HuggingFace AutoTrain or LoRA fine-tuning.")

if __name__ == "__main__":
    asyncio.run(main())