import subprocess
from pathlib import Path

import pyperclip
from dotenv import find_dotenv
from dotenv import load_dotenv
from ruamel.yaml import YAML

from mdbt.ai_core import AiCore
from mdbt.prompts import Prompts

load_dotenv(find_dotenv("../.env"))
load_dotenv(find_dotenv(".env"))


class BuildDBTDocs(AiCore):
    """
    # Make sure you have ANTHROPIC_API_KEY set in your environment variables.
    """

    def __init__(self):
        super().__init__()

    def main(self, select, sys_context, is_new=False):
        if not is_new:
            print(
                """
            1) Build new DBT documentation.
            2) Check existing DBT documentation against model for missing definitions.
            """
            )
            mode = int(input())
        else:
            mode = 1
        print("Getting file.")
        models = self.get_models(select)
        if not models:
            raise ValueError(f"No models found for select '{select}'")

        for model in models:
            original_file_path = model["original_file_path"]
            model_name = model["name"]
            if "l4" in original_file_path.lower() or "l3" in original_file_path.lower():
                system_instructions = Prompts().dbt_docs_gte_l3_prompt
            else:
                system_instructions = Prompts().dbt_docs_lte_l2_prompt

            if sys_context:
                system_instructions += (
                    f"\nContext about system docs are generated for: \n{sys_context}\n"
                )

            sample_data = self._get_sample_data_from_snowflake([model_name])

            system_instructions = system_instructions + sample_data[model_name]

            # Get existing column descriptions from neighboring YML files
            existing_descriptions = self.get_existing_column_descriptions(
                original_file_path
            )
            if existing_descriptions:
                system_instructions += f"Here are some existing column descriptions from neighboring models. These may or may not be relevant to the current model. Assess, and use as a source if it seems to match the sample data: \n{existing_descriptions}"

            # Might bring this back in the future.
            extra_info = ""

            if mode == 1:
                # Build new documentation
                user_input = self.build_user_msg_mode_1(original_file_path, extra_info)
                yml_file_path = original_file_path.replace(".sql", ".yml")
            elif mode == 2:
                # Check existing documentation
                yml_file_path = original_file_path[:-4] + ".yml"
                user_input = self.build_user_msg_mode_2(
                    original_file_path, yml_file_path, extra_info
                )
            else:
                print(mode)
                raise ValueError("Invalid mode")

            messages = [
                {"role": "user", "content": system_instructions + "\n" + user_input}
            ]

            assistant_responses = []
            result = self.send_message(messages)
            assistant_responses.append(result)

            messages.append({"role": "assistant", "content": assistant_responses[0]})
            print(assistant_responses[0])
            output = assistant_responses[0]
            # Check for ``` at end of output (str) and remove
            # Remove trailing markdown code fences if present
            lines = output.split("\n")
            new_lines = []
            for line in lines:
                if "```" not in line:
                    new_lines.append(line)
            output = "\n".join(new_lines)
            if not is_new:
                clip_or_file = input(
                    f"1. to copy to clipboard\n2, to write to file ({yml_file_path}\n:"
                )
            else:
                clip_or_file = "2"

            if clip_or_file == "1":
                print("Output copied to clipboard")
                pyperclip.copy(output)
            elif clip_or_file == "2":
                if mode == 2:
                    # Make a backup of the current YML file.
                    self.backup_existing_yml_file(yml_file_path)
                with open(yml_file_path, "w") as file:
                    file.write(output)
                if not self.is_file_committed(yml_file_path):
                    if not is_new:
                        commit_file = input(
                            "Press 1 to add to git, any other key to byapss: "
                        )
                    else:
                        commit_file = "1"

                    if commit_file == "1":
                        subprocess.run(["git", "add", yml_file_path])

    @staticmethod
    def backup_existing_yml_file(yml_file_path):
        with open(yml_file_path, "r") as file:
            yml_content = file.read()
        with open(yml_file_path + ".bak", "w") as file:
            file.write(yml_content)

    def get_existing_column_descriptions(self, model_file_path: str) -> str:
        """
        Search for YML files in the same folder as the model and extract
        existing column descriptions to provide context for consistent naming.

        Args:
            model_file_path: Path to the SQL model file.

        Returns:
            A formatted string containing existing column descriptions found
            in neighboring YML files.
        """
        folder_path = Path(model_file_path).parent
        current_yml_path = model_file_path.replace(".sql", ".yml")
        column_descriptions = {}

        yaml = YAML()
        yaml.preserve_quotes = True

        # Find all .yml files in the same folder
        yml_files = list(folder_path.glob("*.yml"))

        for yml_file in yml_files:
            # Skip the current model's yml file if it exists
            if str(yml_file) == current_yml_path:
                continue

            try:
                with open(yml_file, "r") as f:
                    yml_content = yaml.load(f)

                if not yml_content:
                    continue

                # Extract column descriptions from models section
                models = yml_content.get("models", [])
                for model in models:
                    columns = model.get("columns", [])
                    for column in columns:
                        col_name = column.get("name")
                        col_desc = column.get("description")
                        if col_name and col_desc:
                            # Only store if not already seen, or prefer longer descriptions
                            if col_name not in column_descriptions or len(
                                str(col_desc)
                            ) > len(str(column_descriptions[col_name])):
                                column_descriptions[col_name] = str(col_desc).strip()

            except Exception as e:
                print(f"Warning: Could not parse {yml_file}: {e}")
                continue

        if not column_descriptions:
            return ""

        # Filter out ID columns to reduce noise (columns ending in _id or named id)
        non_id_descriptions = {
            col_name: desc
            for col_name, desc in column_descriptions.items()
            if not col_name.lower().endswith("_id") and col_name.lower() != "id"
        }

        # Limit to first 300 unique columns to avoid oversized prompts
        max_columns = 300
        limited_descriptions = dict(
            sorted(non_id_descriptions.items())[:max_columns]
        )

        if not limited_descriptions:
            return ""

        # Format the output
        lines = [
            "\n--- Existing Column Descriptions from Neighboring Models ---",
            "Use these existing descriptions for consistency when the same "
            "or similar columns appear in the new model:\n",
        ]
        for col_name, description in limited_descriptions.items():
            # Truncate very long descriptions for context
            desc_preview = (
                description[:200] + "..." if len(description) > 200 else description
            )
            lines.append(f"  - {col_name}: {desc_preview}")

        lines.append("--- End of Existing Column Descriptions ---\n")
        return "\n".join(lines)

    def build_user_msg_mode_1(self, _sql_file_path: str, extra_info: str) -> str:
        self.read_file(_sql_file_path)
        model_name = _sql_file_path.split("/")[-1].split(".")[0]
        prompt_str = f"Build new DBT documentation for the following SQL query with model name {model_name}"
        if len(extra_info):
            prompt_str += f"\n{extra_info}"

        return prompt_str

    def build_user_msg_mode_2(
        self, _sql_file_path: str, _yml_file_path: str, extra_info: str
    ) -> str:
        self.read_file(_sql_file_path)
        yml = self.read_file(_yml_file_path)
        model_name = _sql_file_path.split("/")[-1].split(".")[0]
        prompt_str = f"Check for missing columns in the following DBT documentation for the following SQL query with model name {model_name}. Identify any columns in the DBT documentation that do not exist in the SQL and comment them out."
        if len(extra_info):
            prompt_str += f"\n {extra_info}"
        prompt_str += f"\nYML File Contents:\n{yml}"

        return prompt_str


if __name__ == "__main__":
    BuildDBTDocs().main("revenue_by_dvm")
