import os
import time
import subprocess
import torch
from huggingface_hub import HfApi

# ✅ Handle ImportError for older `huggingface_hub` versions
try:
    from huggingface_hub.utils import RepositoryNotFoundError as HfRepoNotFoundError
except ImportError:
    class HfRepoNotFoundError(Exception):
        pass  # Fallback if the module doesn't exist

# ✅ Configuration Toggles
HF_REPO_NAME = "mikahniehaus/ReviewAI"  # ✅ Change your Hugging Face repo name
GIT_BRANCH_NAME = "main"  # ✅ Change your Git branch name
CONTINUOUS_MODE = True  # ✅ Set to False to disable continuous mode
USE_AI_COMMIT_MESSAGES = True  # ✅ Set to False to disable AI-generated commit messages
CHECK_INTERVAL = 1800  # ✅ 30 minutes (adjustable)
AI_MAX_INPUT_LENGTH = 2000  # ✅ Prevent AI from breaking on large inputs

# ✅ Directories to ignore **ONLY for Hugging Face uploads** (Allows other AI models to be uploaded)
HF_IGNORE_DIRS = {"local_ai_model"}

# ✅ Directories to ignore for **Git commits & pushes**
GIT_IGNORE_DIRS = {".git", "__pycache__", ".venv", "venv", "env", ".env", "hf_env", "node_modules"}

# ✅ Initialize Hugging Face API
api = HfApi()

# ✅ Load AI Model for Summarization (Ensures it runs **completely offline**)
summarizer = None
if USE_AI_COMMIT_MESSAGES:
    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

    # ✅ Use a lightweight, locally-runnable model
    model_name = "sshleifer/distilbart-cnn-12-6"
    local_model_path = os.path.join(os.getcwd(), "local_ai_model")  # Local storage

    if not os.path.exists(local_model_path):
        print("📥 Downloading AI model for offline use...")
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        os.makedirs(local_model_path, exist_ok=True)
        tokenizer.save_pretrained(local_model_path)
        model.save_pretrained(local_model_path)
    else:
        print("✅ Loading AI model from local storage...")
        tokenizer = AutoTokenizer.from_pretrained(local_model_path)
        model = AutoModelForSeq2SeqLM.from_pretrained(local_model_path)

    summarizer = pipeline("summarization", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)

# ✅ Ensure the correct Git branch exists and switch to it
def ensure_git_branch():
    try:
        subprocess.run(["git", "rev-parse", "--is-inside-work-tree"], check=True, stdout=subprocess.DEVNULL)
        subprocess.run(["git", "checkout", "-B", GIT_BRANCH_NAME], check=True)
        print(f"✅ Switched to branch '{GIT_BRANCH_NAME}'")
    except subprocess.CalledProcessError as e:
        print(f"❌ Error switching to branch: {e}")

# ✅ Ensure the Hugging Face repo exists
def ensure_hf_repo():
    try:
        api.repo_info(repo_id=HF_REPO_NAME)  # Check if repo exists
        print(f"✅ Hugging Face repo '{HF_REPO_NAME}' exists.")
    except HfRepoNotFoundError:
        print(f"🚀 Hugging Face repo '{HF_REPO_NAME}' not found. Creating it...")
        api.create_repo(repo_id=HF_REPO_NAME, exist_ok=True)

# ✅ Ensure upstream branch exists before checking changes
def ensure_upstream_branch():
    try:
        subprocess.run(["git", "fetch", "--all"], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        result = subprocess.run(["git", "branch", "-r"], capture_output=True, text=True)
        remote_branches = result.stdout.strip().split("\n")
        return f"origin/{GIT_BRANCH_NAME}" in remote_branches
    except subprocess.CalledProcessError:
        return False

# ✅ Get actual code changes from `git diff` and limit size
def get_code_changes():
    if not ensure_upstream_branch():
        return "Initial commit - Added new files."

    try:
        diff_output = subprocess.check_output(
            ["git", "diff", f"origin/{GIT_BRANCH_NAME}"],
            text=True,
            encoding="utf-8",
            errors="replace",
        )

        cleaned_diff = []
        for line in diff_output.split("\n"):
            if not line.startswith("diff --git") and not line.startswith("--- ") and not line.startswith("+++ "):
                cleaned_diff.append(line)

        full_code_changes = "\n".join(cleaned_diff)

        # ✅ Prevent AI from breaking due to too many changes
        if len(full_code_changes) > AI_MAX_INPUT_LENGTH:
            print(f"⚠️ Code changes exceed {AI_MAX_INPUT_LENGTH} characters. Truncating and summarizing...")
            return full_code_changes[:AI_MAX_INPUT_LENGTH] + "\n[...Truncated...]"

        return full_code_changes

    except subprocess.CalledProcessError:
        return "No code changes detected."

# ✅ Generate AI-powered commit message with better context
def generate_commit_message():
    code_changes = get_code_changes()

    if not USE_AI_COMMIT_MESSAGES:
        return "Updated code with latest changes."

    input_text = f"""
    Summarize the following code changes into a short Git commit message:

    {code_changes}

    **Commit Message Guidelines:**
    - Be clear and concise (under 50 words).
    - Use present tense (e.g., "Fixes issue with login", "Adds API endpoint").
    - Summarize **what changed and why**.
    
    Generate a concise commit message:
    """

    try:
        summary = summarizer(input_text[:AI_MAX_INPUT_LENGTH], max_length=50, min_length=15, do_sample=False)
        return summary[0]["summary_text"]
    except Exception as e:
        print(f"⚠️ AI Commit Message Failed: {e}")
        return "Updated code with latest changes."

# ✅ Fix: Upload to Hugging Face **excluding only the AI model for comments**
def push_to_huggingface(commit_message):
    try:
        print("🚀 Uploading files to Hugging Face...")

        # ✅ Upload all files except those in `HF_IGNORE_DIRS`
        for root, _, files in os.walk("."):
            if any(ignored in root for ignored in HF_IGNORE_DIRS):
                continue  # ✅ Skip AI model directory

            for file in files:
                file_path = os.path.join(root, file)
                repo_path = os.path.relpath(file_path, ".").replace("\\", "/")

                api.upload_file(
                    path_or_fileobj=file_path,
                    path_in_repo=repo_path,
                    repo_id=HF_REPO_NAME,
                    repo_type="model",
                    commit_message=commit_message,
                )

        print("✅ Pushed changes to Hugging Face.")
    except Exception as e:
        print(f"❌ Hugging Face push failed: {e}")

# ✅ Function to push changes
def push_changes():
    ensure_git_branch()
    ensure_hf_repo()

    commit_message = generate_commit_message()
    print(f"📝 Commit message: {commit_message}")

    # ✅ Push to Hugging Face (Excludes only the AI model for commit comments)
    push_to_huggingface(commit_message)

# ✅ Run the script
if CONTINUOUS_MODE:
    while True:
        push_changes()
        time.sleep(CHECK_INTERVAL)
else:
    push_changes()