import gradio as gr
import os
import shutil
import tempfile
import zipfile
import subprocess
import uuid

from ask_agent import ask_agent
from doc_generator import generate_documented_code, generate_requirements_txt
from readme_generator import generate_readme_from_zip

last_processed_repo_path = ""

def process_repo(repo_path, zip_output_name="AutoDocs"):
    with tempfile.TemporaryDirectory() as temp_output_dir:
        # Document .py files
        for root, _, files in os.walk(repo_path):
            for file in files:
                if file.endswith(".py"):
                    file_path = os.path.join(root, file)
                    generate_documented_code(file_path, file_path)

        # requirements.txt
        requirements_path = os.path.join(repo_path, "requirements.txt")
        generate_requirements_txt(repo_path, requirements_path)

        # Create a temporary .zip for README/index
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp_zip:
            zip_path = tmp_zip.name
            with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
                for root, _, files in os.walk(repo_path):
                    for file in files:
                        full_path = os.path.join(root, file)
                        rel_path = os.path.relpath(full_path, repo_path)
                        zipf.write(full_path, rel_path)

        # README + index.md
        readme_path, index_path = generate_readme_from_zip(zip_path, temp_output_dir)

        # Copy the processed repo
        for item in os.listdir(repo_path):
            s = os.path.join(repo_path, item)
            d = os.path.join(temp_output_dir, item)
            if os.path.isdir(s):
                shutil.copytree(s, d, dirs_exist_ok=True)
            else:
                shutil.copy2(s, d)

        dest_readme = os.path.join(temp_output_dir, "README.md")
        dest_index = os.path.join(temp_output_dir, "index.md")

        if os.path.abspath(readme_path) != os.path.abspath(dest_readme):
            shutil.copy2(readme_path, dest_readme)
        if os.path.abspath(index_path) != os.path.abspath(dest_index):
            shutil.copy2(index_path, dest_index)

        # Output zip file with consistent name
        output_zip_path = os.path.join(
            tempfile.gettempdir(), f"{zip_output_name}.zip"
        )
        with zipfile.ZipFile(output_zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for root, _, files in os.walk(temp_output_dir):
                for file in files:
                    full_path = os.path.join(root, file)
                    arcname = os.path.relpath(full_path, temp_output_dir)
                    zipf.write(full_path, arcname)
        global last_processed_repo_path
        last_processed_repo_path = output_zip_path
        return output_zip_path

def process_zip_upload(uploaded_zip_file):
    zip_path = uploaded_zip_file.name
    zip_name = os.path.splitext(os.path.basename(zip_path))[0]  # e.g., my_project.zip → my_project

    with tempfile.TemporaryDirectory() as temp_input_dir:
        input_zip_path = os.path.join(temp_input_dir, "input_repo.zip")
        shutil.copy(zip_path, input_zip_path)
        with zipfile.ZipFile(input_zip_path, "r") as zip_ref:
            zip_ref.extractall(temp_input_dir)

        extracted_dirs = [d for d in os.listdir(temp_input_dir) if os.path.isdir(os.path.join(temp_input_dir, d))]
        repo_root = os.path.join(temp_input_dir, extracted_dirs[0]) if extracted_dirs else temp_input_dir

        return process_repo(repo_root, zip_name)

def process_github_clone(github_url):
    with tempfile.TemporaryDirectory() as clone_dir:
        try:
            subprocess.check_call(["git", "clone", github_url, clone_dir])
            return process_repo(clone_dir)
        except subprocess.CalledProcessError:
            return "❌ Error cloning the GitHub repository. Please check the URL."

# Wrapper for process_zip_upload that also returns the path for the state
def process_zip_and_update_state(uploaded_zip_file):
    zip_path = process_zip_upload(uploaded_zip_file)
    return zip_path, zip_path  # (output for gr.File, output for gr.State)

# Wrapper for process_github_clone as well
def process_git_and_update_state(github_url):
    zip_path = process_github_clone(github_url)
    return zip_path, zip_path

# Gradio user interface
with gr.Blocks() as demo:
    gr.Markdown("# 🤖 AutoDocs – Smart Documentation Generator")
    last_processed_repo_path_state = gr.State(value="")
    with gr.Tab("📦 Upload .zip"):
        zip_file_input = gr.File(label="Drop your repo .zip file here", file_types=['.zip'])
        generate_btn_zip = gr.Button("📄 Generate from ZIP")
        output_zip_zip = gr.File(label="⬇️ Download your documented repo")
    
    with gr.Tab("🌐 GitHub URL"):
        github_url_input = gr.Text(label="Link to GitHub repository", placeholder="https://github.com/user/repo.git")
        generate_btn_git = gr.Button("📄 Generate from GitHub")
        output_zip_git = gr.File(label="⬇️ Download your documented repo")

    with gr.Tab("🧠 Ask the agent about the repo"):
        chatbot = gr.Chatbot()
        user_input = gr.Textbox(placeholder="Ask your question here...")
        send_btn = gr.Button("Send")

        send_btn.click(
            fn=ask_agent,
            inputs=[chatbot, user_input, last_processed_repo_path_state],
            outputs=[chatbot, user_input]
        )

    generate_btn_zip.click(
        fn=process_zip_and_update_state,
        inputs=[zip_file_input],
        outputs=[output_zip_zip, last_processed_repo_path_state]
    )

    generate_btn_git.click(
        fn=process_git_and_update_state,
        inputs=[github_url_input],
        outputs=[output_zip_git, last_processed_repo_path_state]
    )

if __name__ == "__main__":
    demo.queue()
    demo.launch()