import json
import os
import shutil
import tempfile
import re
import webbrowser

import gradio as gr
from git import Repo
from github import Github
from loguru import logger

from sweepai.app.api_client import APIClient, create_pr_function, create_pr_function_call
from sweepai.app.config import State, SweepChatConfig
from sweepai.core.entities import Snippet
from sweepai.utils.config.client import SweepConfig

def main():
    config = SweepChatConfig.load()

    api_client = APIClient(config=config)

    pr_summary_template = '''⏳ I'm creating the following PR...

    **{title}**
    {summary}

    Here is my plan:
    {plan}

    Reply to propose changes to the plan.'''

    print("Getting list of repos...")
    github_client = Github(config.github_pat)
    repos = list(github_client.get_user().get_repos())
    print("Done.")

    css = '''
    footer {
        visibility: hidden;
    }
    pre, code {
        white-space: pre-wrap !important;
        word-break: break-all !important;
    }
    #snippets {
        height: 600px;
        overflow-y: scroll;
    }
    #message_box > label > span {
        display: none;
    }
    '''


    def get_files_recursively(root_path, path=''):
        files = []
        path_to_contents = {}

        if path == '.git':
            return files, path_to_contents

        current_dir = os.path.join(root_path, path)
        entries = os.listdir(current_dir)

        for entry in entries:
            entry_path = os.path.join(current_dir, entry)

            if os.path.isfile(entry_path):
                try:
                    with open(entry_path, 'r', encoding="utf-8", errors="ignore") as file:
                        contents = file.read()
                    path_to_contents[entry_path[len(root_path) + 1:]] = contents
                    files.append(entry_path[len(root_path) + 1:])
                except UnicodeDecodeError as e:
                    logger.warning(f"Received warning {e}, skipping...")
                    continue
            elif os.path.isdir(entry_path):
                subfiles, subpath_to_contents = get_files_recursively(root_path, os.path.join(path, entry))
                files.extend(subfiles)
                path_to_contents.update(subpath_to_contents)

        return files, path_to_contents


    def get_installation_id(repo_full_name):
        config.repo_full_name = repo_full_name
        api_client.config = config
        installation_id = api_client.get_installation_id()
        return installation_id


    path_to_contents = {}


    def get_files(repo_full_name):
        global path_to_contents
        global repo
        if repo_full_name is None:
            all_files = []
        else:
            # Make sure repo is added to Sweep before checking all recursive files
            try:
                installation_id = get_installation_id(repo_full_name)
                assert installation_id
            except:
                return []
            repo = github_client.get_repo(repo_full_name)
            branch_name = SweepConfig.get_branch(repo)
            repo_url = f"https://x-access-token:{config.github_pat}@github.com/{repo_full_name}.git"
            try:
                repo_dir = os.path.join(tempfile.gettempdir(), repo_full_name)
                if os.path.exists(repo_dir):
                    git_repo = Repo(repo_dir)
                else:
                    git_repo = Repo.clone_from(repo_url, repo_dir)
                git_repo.git.checkout(branch_name)
                git_repo.remotes.origin.pull()
            except Exception as e:
                logger.warning(f"Git pull failed with error {e}, deleting cache and recloning...")
                shutil.rmtree(repo_dir)
                git_repo = Repo.clone_from(repo_url, repo_dir)
                git_repo.git.checkout(branch_name)
                git_repo.remotes.origin.pull()
            all_files, path_to_contents = get_files_recursively(repo_dir)
        return all_files


    def get_files_update(*args):
        global repo
        if len(args) > 0:
            repo = args[0]
        else:
            repo = config.repo_full_name
        return gr.Dropdown.update(choices=get_files(repo))


    def parse_response(raw_response: str) -> tuple[str, list[tuple[str, str]]]:
        if "Plan:" not in raw_response:
            response, raw_plan = raw_response, ""
        else:
            response, raw_plan = raw_response.split("Plan:", 1)
        if response.startswith("Response:"):
            response = response[len("Response:"):]
        plan = [(line[:line.find(":")].strip(), line[line.find(":") + 1:].strip()) for line in raw_plan.split("\n*") if
                line]
        return response, plan

    try:
        user_info = api_client.get_user_info()
    except Exception as e:
        logger.warning(e)
        user_info = {"is_paying_user": False, "remaining_tickets": 0}
    global_state = config.state

    with gr.Blocks(theme=gr.themes.Soft(), title="Sweep Chat", css=css) as demo:
        print("Launching gradio!")
        with gr.Row():
            with gr.Column(scale=2):
                repo_full_name = gr.Dropdown(choices=[repo.full_name for repo in repos], label="Repo full name",
                                            value=lambda: config.repo_full_name or "")
            print("Indexing files...")
            with gr.Column(scale=4):
                file_names = gr.Dropdown(choices=get_files(config.repo_full_name), multiselect=True, label="Files",
                                        value=lambda: global_state.file_paths)
            print("Indexed files!")
            repo_full_name.change(fn=get_files_update, inputs=repo_full_name, outputs=file_names)
            with gr.Column(scale=1):
                restart_button = gr.Button("Restart")

        with gr.Row():
            with gr.Column(scale=2):
                chatbot = gr.Chatbot(height=600, value=lambda: global_state.chat_history)
            with gr.Column():
                with gr.Row():
                    snippets_text = gr.Markdown(value=lambda: global_state.snippets_text, elem_id="snippets")

        with gr.Row():
            plan = gr.List(
                value=[[filename + ": " + instructions] for filename, instructions in global_state.plan] or [[""]],
                headers=["Proposed Plan"],
                interactive=True,
                col_count=(1, "static"),
                wrap=True,
                visible=global_state.plan_toggle,
            )

        with gr.Row():
            with gr.Column(scale=8):
                msg = gr.Textbox(placeholder="Send a message to Sweep", label=None, elem_id="message_box")
            create_pr_button_column = gr.Column(scale=0.5, visible=bool(global_state.plan))
            with create_pr_button_column:
                create_pr_button = gr.Button(value="Create PR", interactive=bool(global_state.plan))
        
        with gr.Row():
            with gr.Column():
                plan_toggle = gr.Checkbox(label="Propose plan (Experimental)", value=lambda: global_state.plan_toggle, interactive=True, elem_id="plan_toggle")
            with gr.Column(scale=5):
                tier_message = "⚡ Free Trial" if not user_info["is_paying_user"] else "💎 Pro Tier"
                remaining_tickets = user_info["remaining_tickets"]
                user_status = gr.Markdown(value=f"Logged in as {config.github_username} ({tier_message}), {remaining_tickets} tickets remaining.")
        
        def on_plan_toggle_click(plan_toggle):
            global global_state
            global_state.plan_toggle = plan_toggle
            config.state = global_state
            config.save()
            return gr.update(visible=plan_toggle), gr.update(visible=plan_toggle)
        
        plan_toggle.change(on_plan_toggle_click, [plan_toggle], [create_pr_button_column, plan])

        def clear_inputs():
            global global_state
            global_state = State()
            config.state = global_state
            config.save()
            return [], [], [[""]]


        restart_button.click(clear_inputs, None, [file_names, chatbot, plan])

        file_names.change(fn=get_files_update, inputs=repo_full_name, outputs=file_names)

        global searched
        global selected_snippets
        global file_to_str

        searched = False
        selected_snippets = []
        file_to_str = {}


        def repo_name_change(repo_full_name):
            global installation_id
            try:
                try:
                    installation_id = get_installation_id(repo_full_name)
                except Exception as e:
                    raise gr.Error(str(e))
                assert installation_id
                config.installation_id = installation_id
                api_client.config = config
                config.save()
                return ""
            except Exception as e:
                config.repo_full_name = None
                config.installation_id = None
                config.save()
                api_client.config = config
                raise e


        def build_string():
            global selected_snippets
            global file_to_str
            for snippet in selected_snippets:
                file_name = snippet.file_path
                if file_name not in file_to_str:
                    add_file_to_dict(file_name)
            snippets_text = "### Relevant snippets:\n" + "\n\n".join(
                [file_to_str[snippet.file_path] for snippet in selected_snippets])
            return snippets_text


        repo_full_name.change(fn=repo_name_change, inputs=[repo_full_name], outputs=[msg])


        def add_file_to_dict(file_name):
            global file_to_str
            global path_to_contents
            global repo
            if file_name in path_to_contents:
                file_contents = path_to_contents[file_name]
            else:
                file_contents = repo.get_contents(file_name, ref=SweepConfig.get_branch(repo)).decoded_content.decode(
                    'utf-8')
            file_contents_split = file_contents.split("\n")
            length = len(file_contents_split)
            backtick, escaped_backtick = "`", "\\`"
            preview = "\n".join(file_contents_split[:3]).replace(backtick, escaped_backtick)
            file_to_str[file_name] = f'{file_name}:0:{length}\n```\n{preview}\n...\n```'


        def file_names_change(file_names):
            global selected_snippets
            global file_to_str
            global path_to_contents
            selected_snippets = []
            for file_name in file_names:
                if file_name in path_to_contents:
                    selected_snippets.append(
                        Snippet(content=path_to_contents[file_name], start=0, end=path_to_contents[file_name].count('\n'),
                                file_path=file_name))
            return file_names, build_string()


        file_names.change(fn=file_names_change, inputs=[file_names], outputs=[file_names, snippets_text])

        def handle_message_submit(repo_full_name: str, user_message: str, history: list[tuple[str | None, str | None]]):
            if not repo_full_name:
                raise gr.Error("Set the repository name first")
            return gr.update(value="", interactive=False), history + [[user_message, None]], gr.Button.update(
                interactive=False)

        def _handle_message_stream(
            chat_history: list[tuple[str | None, str | None]], 
            snippets_text, 
            file_names, 
            plan,
            plan_toggle
        ):
            global selected_snippets
            global searched
            if plan is None or plan == [[]] or plan == [[""]] or plan == [["", ""]]:
                plan = [["", ""]]
            message = chat_history[-1][0]
            yield chat_history, snippets_text, file_names, plan
            if not selected_snippets:
                searched = True
                # Searching for relevant snippets
                chat_history[-1][1] = "Searching for relevant snippets..."
                snippets_text = build_string()
                yield chat_history, snippets_text, file_names, plan
                logger.info("Fetching relevant snippets...")
                try:
                    selected_snippets += api_client.search(chat_history[-1][0], 3)
                except Exception as e:
                    raise gr.Error(str(e))
                snippets_text = build_string()
                file_names = [snippet.file_path for snippet in selected_snippets]
                yield chat_history, snippets_text, file_names, plan
                logger.info("Fetched relevant snippets.")
                chat_history[-1][1] = "Found relevant snippets."
                # Update using chat_history
                snippets_text = build_string()
                yield chat_history, snippets_text, file_names, plan

            # Generate response
            logger.info("...")
            chat_history.append([None, "..."])
            yield chat_history, snippets_text, file_names, plan
            chat_history[-1][1] = ""
            logger.info("Starting to generate response...")
            if len(chat_history) > 1 and "create pr" in message.lower():
                try:
                    stream = api_client.stream_chat(
                        chat_history,
                        selected_snippets,
                        do_add_plan=plan_toggle,
                        functions=[create_pr_function],
                        function_call=create_pr_function_call,
                    )
                except Exception as e:
                    raise gr.Error(str(e))
            else:
                stream = api_client.stream_chat(chat_history, selected_snippets, do_add_plan=plan_toggle)
            function_name = ""
            raw_arguments = ""
            raw_response = ""
            parsed_response = ""
            try:
                for chunk in stream:
                    if chunk.get("content"):
                        token = chunk["content"]
                        raw_response += token
                        if plan_toggle:
                            parsed_response, plan = parse_response(raw_response)
                            chat_history[-1][1] = parsed_response
                            yield chat_history, snippets_text, file_names, plan
                        else:
                            chat_history[-1][1] = raw_response
                            yield chat_history, snippets_text, file_names, plan
                    if chunk.get("function_call"):
                        function_call = chunk["function_call"]
                        function_name = function_name or function_call.get("name")
                        raw_arguments += function_call.get("arguments")
                        chat_history[-1][1] = f"Calling function: `{function_name}`\n```json\n{raw_arguments}\n```"
                        yield chat_history, snippets_text, file_names, plan
            except Exception as e:
                logger.error(e)
                raise gr.Error(str(e))
            if function_name:
                arguments = json.loads(raw_arguments)
                if function_name == "create_pr":
                    assert "title" in arguments
                    assert "summary" in arguments
                    assert "plan" in arguments
                    if "branch" not in arguments:
                        arguments["branch"] = arguments["title"].lower().replace(" ", "_").replace("-", "_")[:50]
                    chat_history[-1][1] = pr_summary_template.format(
                        title=arguments["title"],
                        summary=arguments["summary"],
                        plan="\n".join([f"* `{item['file_path']}`: {item['instructions']}" for item in arguments["plan"]])
                    )
                    yield chat_history, snippets_text, file_names, plan
                    plan = [(item["file_path"], item["instructions"]) for item in arguments["plan"]]
                    yield chat_history, snippets_text, file_names, plan
                else:
                    raise NotImplementedError


        def handle_message_stream(
            chat_history: list[tuple[str | None, str | None]], 
            snippets_text, 
            file_paths, 
            plan, 
            plan_toggle
        ):
            global global_state
            for chat_history, snippets_text, file_paths, plan in _handle_message_stream(chat_history, snippets_text,
                                                                                        file_paths, plan, plan_toggle):
                if plan is None or plan == [[]] or plan == [[""]] or plan == [["", ""]]:
                    plan = [["", ""]]
                if plan and isinstance(plan[0], list):
                    plan = [item.split(":") for item in plan if ":" in item]
                global_state = State(
                    chat_history=chat_history,
                    snippets_text=snippets_text,
                    file_paths=file_paths,
                    plan=plan,
                )
                config.state = global_state
                config.save()
                if plan == [] or plan == [["", ""]]:
                    yield chat_history, snippets_text, file_paths, [[""]]
                else:
                    yield chat_history, snippets_text, file_paths, [(file_path + ": " + instructions,) for
                                                                file_path, instructions in plan]


        response = msg \
            .submit(handle_message_submit, [repo_full_name, msg, chatbot], [msg, chatbot, create_pr_button], queue=False) \
            .then(handle_message_stream, [chatbot, snippets_text, file_names, plan, plan_toggle],
                [chatbot, snippets_text, file_names, plan]) \
            .then(lambda: [gr.update(interactive=True), gr.update(interactive=True)], None, [msg, create_pr_button], queue=False)


        def validate_branch_name(branch_name):
            # Replace any characters that are not alphanumeric or '-' or '_' with '_'
            valid_branch_name = re.sub('[^0-9a-zA-Z_-]', '_', branch_name)
            return valid_branch_name

        def on_create_pr_button_click(chat_history: list[tuple[str | None, str | None]], plan: list[tuple[str]]):
            title = chat_history[0][0][:240]
            content = chat_history[-1][1]
            valid_branch_name = validate_branch_name(title.lower().replace(" ", "_").replace("-", "_")[:50])
            chat_history.append((None, "⌛ Creating PR..."))
            yield chat_history, gr.Button.update(interactive=False)
            try:
                pull_request = api_client.create_pr(
                    file_change_requests=[(item[:item.find(":")], item[item.find(":") + 1:]) for item, *_ in plan],
                    pull_request={
                        "title": title,
                        "content": content,
                        "branch_name": valid_branch_name
                    },
                    messages=chat_history,
                )
                chat_history.append((None, f"✅ PR created at {pull_request['html_url']}"))
                webbrowser.open_new_tab(pull_request["html_url"])
                yield chat_history, gr.Button.update(interactive=True)
            except Exception as e:
                yield chat_history, gr.Button.update(interactive=True)
                raise gr.Error(str(e))


        create_pr_button.click(on_create_pr_button_click, [chatbot, plan], [chatbot, create_pr_button])

    demo.queue()
    demo.launch(inbrowser=True)

if __name__ == "__main__":
    main()
