from huggingface_hub import hf_hub_download, snapshot_download, login
import pandas as pd
import importlib
import importlib.util
import streamlit as st
import sys
import os
import shutil
from pathlib import Path
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# #___________ cache issue 03_08
# if os.path.exists("private_space_cache"):
#     shutil.rmtree("private_space_cache")  # Forcefully remove corrupted cache folder
# incomplete_files = Path(".").rglob("*.incomplete")
# for file in incomplete_files:
#     try:
#         file.unlink()
#     except Exception as e:
#         print(f"Could not delete {file}: {e}")
# #____________________________________

# HF_TOKEN_LLAMA = os.environ.get("HF_TOKEN_LLAMA")
# login(token=HF_TOKEN_LLAMA)

# HF_TOKEN = os.environ.get("HF_TOKEN") #get HF_TOKEN 
# login(token=HF_TOKEN)


# USER_NAME = os.getenv("USER_NAME", "").strip().strip('"')
# PRIVATE_SPACE_NAME = os.getenv("PRIVATE_SPACE_NAME", "").strip().strip('"')

# #Construct the repo ID
# REPO_ID = f"{USER_NAME}/{PRIVATE_SPACE_NAME}"

# REPO_TYPE = "space"

# # sys.path.append(repo_dir)

# # Download the entire space, including the fine-tuned model folder
# repo_dir = snapshot_download(
#     repo_id=REPO_ID,
#     repo_type=REPO_TYPE,
#     token=HF_TOKEN,
#     cache_dir="private_space_cache",
#     force_download=True  # Forces redownload
# )


# # Change the working directory to the downloaded snapshot directory
# # This step is very imporptant
# os.chdir(repo_dir)

# # # Add repo directory to sys.path so Python can find modules inside it
# sys.path.append(repo_dir)


# # Download specific files (if snapshot_download wasn't used)
# app_path = hf_hub_download(
#     repo_id=REPO_ID,
#     filename="app.py",
#     repo_type=REPO_TYPE
# )

# # Load and execute `app.py`
# spec_app = importlib.util.spec_from_file_location("*", app_path)
# app_module = importlib.util.module_from_spec(spec_app)
# spec_app.loader.exec_module(app_module)
# # Now you can use functions from utils_module
# result = app_module.main()


# Set page config as the FIRST command
#st.set_page_config(page_title="📐 Math Assignment Optimizer", layout="wide")

# Helper function to clean up cache
def clear_cache(cache_dir="private_space_cache"):
    """Remove cache directory and all .incomplete files."""
    if os.path.exists(cache_dir):
        shutil.rmtree(cache_dir, ignore_errors=True)
    for file in Path(".").rglob("*.incomplete"):
        try:
            file.unlink()
        except Exception as e:
            st.warning(f"Could not delete {file}: {e}")

# Get environment variables
HF_TOKEN_LLAMA = os.environ.get("HF_TOKEN_LLAMA")
HF_TOKEN = os.environ.get("HF_TOKEN")
USER_NAME = os.getenv("USER_NAME", "").strip().strip('"')
PRIVATE_SPACE_NAME = os.getenv("PRIVATE_SPACE_NAME", "").strip().strip('"')

# Log in to Hugging Face
login(token=HF_TOKEN_LLAMA)
login(token=HF_TOKEN)

# Construct the repo ID
REPO_ID = f"{USER_NAME}/{PRIVATE_SPACE_NAME}"
REPO_TYPE = "space"

# Always clear cache before downloading
with st.spinner("Clearing old cache..."):
    clear_cache()

# Download the private space
with st.spinner("Downloading private space..."):
    try:
        repo_dir = snapshot_download(
            repo_id=REPO_ID,
            repo_type=REPO_TYPE,
            token=HF_TOKEN,
            cache_dir="private_space_cache",
            force_download=True,
            local_dir_use_symlinks=False
        )
    except Exception as e:
        st.error(f"Failed to download private space: {e}")
        st.stop()

# Change working directory and update sys.path
os.chdir(repo_dir)
sys.path.append(repo_dir)

# Download and load app.py from the private space
try:
    app_path = hf_hub_download(
        repo_id=REPO_ID,
        filename="app.py",
        repo_type=REPO_TYPE,
        token=HF_TOKEN,
        force_download=True
    )
except Exception as e:
    st.error(f"Failed to download app.py: {e}")
    st.stop()

# Load and execute app.py
spec_app = importlib.util.spec_from_file_location("app_module", app_path)
app_module = importlib.util.module_from_spec(spec_app)
spec_app.loader.exec_module(app_module)

# Run the private app's main function
try:
    result = app_module.main()
except Exception as e:
    st.error(f"Error running private app: {e}")
    st.stop()

# Optional: Clean up after execution
with st.spinner("Cleaning up..."):
    clear_cache()