import os
import subprocess
import sys
import zipfile
import urllib.request
import shutil
from pathlib import Path

def setup_indiclid():
    """Setup IndicLID with proper class file"""
    print("🚀 Setting up IndicLID...")
    
    # Create directory structure
    os.makedirs("ai4bharat", exist_ok=True)
    os.makedirs("models", exist_ok=True)
    
    # Create __init__.py
    with open("ai4bharat/__init__.py", "w") as f:
        f.write("# AI4Bharat package\n")
    
    # The IndicLID.py file should be created separately (above)
    # Download model files
    print("🤖 Downloading IndicLID models...")
    model_urls = [
        "https://github.com/AI4Bharat/IndicLID/releases/download/v1.0/indiclid-bert.zip",
        "https://github.com/AI4Bharat/IndicLID/releases/download/v1.0/indiclid-ftn.zip",
        "https://github.com/AI4Bharat/IndicLID/releases/download/v1.0/indiclid-ftr.zip"
    ]
    
    current_dir = os.getcwd()
    os.chdir("models")
    
    for url in model_urls:
        filename = url.split("/")[-1]
        print(f"📥 Downloading {filename}...")
        try:
            urllib.request.urlretrieve(url, filename)
            print(f"✅ Downloaded {filename}")
            
            print(f"📦 Extracting {filename}...")
            with zipfile.ZipFile(filename, 'r') as zip_ref:
                zip_ref.extractall('.')
            os.remove(filename)
            print(f"✅ Extracted {filename}")
            
        except Exception as e:
            print(f"❌ Error with {filename}: {e}")
    
    os.chdir(current_dir)
    print("🎉 IndicLID setup completed!")
    return True

if __name__ == "__main__":
    setup_indiclid()