import html
import json
import os
from textwrap import dedent
from typing import Any, Coroutine, List

from directory_tree import display_tree
from dotenv import load_dotenv
from pydantic import Field

from ...core.main import ChatMessage, Models, Step, step_to_json_schema
from ...core.sdk import ContinueSDK
from ...libs.util.devdata import dev_data_logger
from ...libs.util.strings import remove_quotes_and_escapes
from ...libs.util.telemetry import posthog_logger

load_dotenv()


def add_ellipsis(text: str, max_length: int = 200) -> str:
    if len(text) > max_length:
        return text[: max_length - 3] + "..."
    return text


class SimpleChatStep(Step):
    name: str = "Generating Response..."
    manage_own_chat_context: bool = True
    description: str = ""
    messages: List[ChatMessage] = None

    async def run(self, sdk: ContinueSDK):
        # Check if proxy server API key
        messages = self.messages or await sdk.get_chat_context()

        generator = sdk.models.chat.stream_chat(
            messages, temperature=sdk.config.temperature
        )

        posthog_logger.capture_event(
            "model_use",
            {
                "model": sdk.models.default.model,
                "provider": sdk.models.default.__class__.__name__,
            },
        )
        dev_data_logger.capture(
            "model_use",
            {
                "model": sdk.models.default.model,
                "provider": sdk.models.default.__class__.__name__,
            },
        )

        async for chunk in generator:
            if sdk.current_step_was_deleted():
                # So that the message doesn't disappear
                self.hide = False
                await sdk.update_ui()
                break

            if "content" in chunk:
                self.description += chunk["content"]

                # HTML unencode
                end_size = len(chunk["content"]) - 6
                if "&" in self.description[-end_size:]:
                    self.description = self.description[:-end_size] + html.unescape(
                        self.description[-end_size:]
                    )

                await sdk.update_ui()

        if sdk.config.disable_summaries:
            self.name = ""
        else:
            self.name = "Generating title..."
            await sdk.update_ui()
            self.name = add_ellipsis(
                remove_quotes_and_escapes(
                    await sdk.models.summarize.complete(
                        f'"{self.description}"\n\nPlease write a short title summarizing the message quoted above. Use no more than 10 words:',
                        max_tokens=20,
                        log=False,
                    )
                ),
                200,
            )

        await sdk.update_ui()

        self.chat_context.append(
            ChatMessage(role="assistant", content=self.description, summary=self.name)
        )

        # TODO: Never actually closing.
        await generator.aclose()


class AddFileStep(Step):
    name: str = "Add File"
    description = "Add a file to the workspace. Should always view the directory tree before this."
    filename: str
    file_contents: str

    async def describe(
        self, models: Models
    ) -> Coroutine[Any, Any, Coroutine[str, None, None]]:
        return f"Added a file named `{self.filename}` to the workspace."

    async def run(self, sdk: ContinueSDK):
        await sdk.add_file(self.filename, self.file_contents)

        await sdk.ide.setFileOpen(
            os.path.join(sdk.ide.workspace_directory, self.filename)
        )


class DeleteFileStep(Step):
    name: str = "Delete File"
    description = "Delete a file from the workspace."
    filename: str

    async def describe(
        self, models: Models
    ) -> Coroutine[Any, Any, Coroutine[str, None, None]]:
        return f"Deleted a file named `{self.filename}` from the workspace."

    async def run(self, sdk: ContinueSDK):
        await sdk.delete_file(self.filename)


class AddDirectoryStep(Step):
    name: str = "Add Directory"
    description = "Add a directory to the workspace."
    directory_name: str

    async def describe(
        self, models: Models
    ) -> Coroutine[Any, Any, Coroutine[str, None, None]]:
        return f"Added a directory named `{self.directory_name}` to the workspace."

    async def run(self, sdk: ContinueSDK):
        try:
            await sdk.add_directory(self.directory_name)
        except FileExistsError:
            self.description = f"Directory {self.directory_name} already exists."


class RunTerminalCommandStep(Step):
    name: str = "Run Terminal Command"
    description: str = "Run a terminal command."
    command: str

    async def run(self, sdk: ContinueSDK):
        self.description = f"Copy this command and run in your terminal:\n\n```bash\n{self.command}\n```"


class ViewDirectoryTreeStep(Step):
    name: str = "View Directory Tree"
    description: str = "View the directory tree to learn which folder and files exist. You should always do this before adding new files."

    async def describe(
        self, models: Models
    ) -> Coroutine[Any, Any, Coroutine[str, None, None]]:
        return "Viewed the directory tree."

    async def run(self, sdk: ContinueSDK):
        self.description = (
            f"```\n{display_tree(sdk.ide.workspace_directory, True, max_depth=2)}\n```"
        )


class EditFileStep(Step):
    name: str = "Edit File"
    description: str = "Edit a file in the workspace that is not currently open."
    filename: str = Field(..., description="The name of the file to edit.")
    instructions: str = Field(..., description="The instructions to edit the file.")
    hide: bool = True

    async def run(self, sdk: ContinueSDK):
        await sdk.edit_file(self.filename, self.instructions)