from dataclasses import dataclass
import json
import os

import tiktoken
from pydantic import BaseModel, Field
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider

from gpt_pr.gitutil import BranchInfo, fetch_nearest_git_dir
from gpt_pr.config import config, CONFIG_PROJECT_REPO_URL
import gpt_pr.consolecolor as cc


class PRTemplateModel(BaseModel):
    title: str = Field(description="Title of the pull request")
    description: str = Field(description="Description of the pull request")


TOKENIZER_RATIO = 4

DEFAULT_PR_TEMPLATE = (
    "### Ref. [Link]\n\n## What was done?\n[Fill here]\n\n"
    "## How was it done?\n[Fill here]\n\n"
    "## How was it tested?\n[Fill here with test information from diff content or commits]"
)

SYSTEM_PROMPT = '''You are a generator of pull request data based on diff changes.

By analyzing the diff content and commit messages between two branches, you must strictly adhere to
the provided Pull Request template and produce a complete, ready-to-use PR output.

Your response must include:
- A clear and concise PR title.
- A PR description that:
  - Details the work accomplished.
  - Describes the methodology used, including testing procedures.
  - Lists significant changes in bullet points.

Rules:
- Do not include raw diff content of any size.
- Do not add any explanations, suggestions, or messages directed to the user.

Pull Request Template:
---
{pr_template}
---
'''


def _get_pr_template():
    pr_template = DEFAULT_PR_TEMPLATE
    git_dir = fetch_nearest_git_dir(os.getcwd())

    try:
        github_dir = os.path.join(git_dir, ".github")
        github_files = os.listdir(github_dir)
        pr_template_file = [
            f for f in github_files if f.lower().startswith("pull_request_template")
        ][0]
        pr_template_file_path = os.path.join(github_dir, pr_template_file)

        with open(pr_template_file_path, "r") as f:
            local_pr_template = f.read()

            if local_pr_template.strip() != "":
                print("Found PR template at:", pr_template_file_path)
                pr_template = local_pr_template
            else:
                print(
                    "Empty PR template at:",
                    pr_template_file_path,
                    "using default template.",
                )
    except Exception:
        print("PR template not found in .github dir. Using default template.")

    return pr_template


def _get_open_ai_key():
    api_key = config.get_user_config("OPENAI_API_KEY")

    if not api_key:
        api_key = os.environ.get("OPENAI_API_KEY")

    if not api_key:
        print(
            'Please set "openai_api_key" config, just run:',
            cc.yellow("gpt-pr-config set openai_api_key [open ai key]"),
        )
        raise SystemExit(1)

    return api_key


def _count_tokens(text: str) -> int:
    """Returns the number of tokens in a text string."""
    openai_model = config.get_user_config("OPENAI_MODEL")
    try:
        encoding = tiktoken.encoding_for_model(openai_model)
    except KeyError:
        encoding = tiktoken.get_encoding("cl100k_base")

    return len(encoding.encode(text))


@dataclass
class PrData:
    branch_info: BranchInfo
    title: str
    body: str

    def to_display(self):
        return "\n".join(
            [
                f"{cc.bold('Repository')}: {cc.yellow(self.branch_info.owner)}/{cc.yellow(self.branch_info.repo)}",
                f"{cc.bold('Title')}: {cc.yellow(self.title)}",
                f"{cc.bold('Branch name')}: {cc.yellow(self.branch_info.branch)}",
                f"{cc.bold('Base branch')}: {cc.yellow(self.branch_info.base_branch)}",
                f"{cc.bold('PR Description')}:\n{self.create_body()}",
            ]
        )

    def create_body(self):
        body = self.body

        if config.get_user_config("ADD_TOOL_SIGNATURE") == "true":
            pr_signature = f"Generated by [GPT-PR]({CONFIG_PROJECT_REPO_URL})"
            body += "\n\n---\n\n" + pr_signature

        return body


def get_pr_data(branch_info):
    system_prompt, messages = _get_messages(branch_info)

    openai_model = config.get_user_config("OPENAI_MODEL")
    model = OpenAIChatModel(openai_model, provider=OpenAIProvider(api_key=_get_open_ai_key()))

    support_agent = Agent(
        model=model,  # TODO: make configurable for other providers
        output_type=PRTemplateModel,
        instructions=system_prompt,
    )

    print("Generating changes description using OpenAI model", cc.yellow(openai_model), '. This may take time...')

    result = support_agent.run_sync(messages)
    output = result.output

    return PrData(
        branch_info=branch_info, title=output.title.strip(), body=output.description.strip()
    )


def _get_messages(branch_info):
    system_prompt = SYSTEM_PROMPT.format(pr_template=_get_pr_template())

    messages = []

    if len(branch_info.highlight_commits) > 0:
        messages.append("main commits:\n" + "\n".join(branch_info.highlight_commits))
        messages.append("---")
        messages.append("secondary commits:\n" + "\n".join(branch_info.commits))
    else:
        messages.append("git commits:\n" + "\n".join(branch_info.commits))

    joined_messages = "\n".join([m for m in messages])
    current_total_tokens = _count_tokens(joined_messages) + _count_tokens(SYSTEM_PROMPT)

    input_max_tokens = int(config.get_user_config("INPUT_MAX_TOKENS"))

    if current_total_tokens > input_max_tokens:
        exp_message = (
            f"Length of {current_total_tokens} tokens for basic prompt "
            f"(description and commits) is greater than max tokens {input_max_tokens} "
            "(config 'input_max_tokens')"
        )
        raise Exception(exp_message)

    total_tokens_with_diff = current_total_tokens + _count_tokens(branch_info.diff)
    if total_tokens_with_diff > input_max_tokens:
        print_msg = (
            f"Length git changes with diff is too big (total is {total_tokens_with_diff}, "
            f"'input_max_tokens' config is {input_max_tokens})."
        )
        print(print_msg, cc.red("Skipping changes diff content..."))
    else:
        messages.append("Diff changes:\n" + branch_info.diff)

    return system_prompt, '\n'.join(messages)


def _parse_json(content):
    """
    A bit of a hack to parse the json content from the chat completion
    Sometimes it returns a string with invalid json content (line breaks) that
    makes it hard to parse.
    example:

    content = '{\n"title": "feat(dependencies): pin dependencies versions",\n"description":
                "### Ref. [Link]\n\n## What was done? ..."\n}'
    """

    try:
        content = content.replace('{\n"title":', '{"title":')
        content = content.replace(',\n"description":', ',"description":')
        content = content.replace("\n}", "}")
        content = content.replace("\n", "\\n")

        return json.loads(content)
    except Exception as e:
        print("Error to decode message:", e)
        print("Content:", content)
        raise e
