from abc import abstractmethod
from dataclasses import dataclass
from functools import wraps
from typing import Optional, Iterable, Any

from rhodecode.apps.ai_agents.ai_settings import AISettings
from rhodecode.lib.codeblocks import DiffSet
from rhodecode.lib.vcs.nodes import FileNode


FILE = "FILE:"


class AIServiceError(Exception):
    pass


@dataclass
class Response:
    model: str
    message: str
    error: bool = False


@dataclass
class Request:
    content: str


@dataclass
class PingRequest(Request):
    function_name: str
    role: str


@dataclass
class Review:
    role: str
    msg: str


@dataclass
class CodeReviewRequest(Request):
    content: list[Review]
    function_name: str


@dataclass
class CustomFunctions:
    Ping: str = "ping_pong"
    Code_review: str = "return_code_review"


def wrap_ai_exceptions(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        try:
            return f(*args, **kwargs)
        except Exception as e:
            exc = AIServiceError(str(e))
            # to not lose the original traceback
            exc.__cause__ = e
            raise exc.with_traceback(e.__traceback__)

    return wrapper


class AIService:
    DEFAULT_BASIC_REVIEW_POINTS = [
        "Correctness and edge cases (logic errors, boundary conditions, invalid inputs).",
        "Error handling and resilience (fail-fast where appropriate, clear propagation, retries/backoff, cleanup).",
        "Security (input validation/sanitization, injection risks, unsafe eval/exec, authn/z, secret handling, serialization).",
        "Readability and maintainability (clear names, comments/docs where helpful, remove dead code, consistent formatting).",
        "Interface & API design (encapsulation, stable contracts, backward compatibility, minimal surface area).",
        "Performance & complexity (hot paths, unnecessary allocations/work, algorithmic complexity, scalability).",
        "Concurrency & asynchrony (race conditions, synchronization, thread/process safety, async/await or equivalents).",
        "Resource management (files, network, DBs; timeouts; quotas; connection pooling; deterministic cleanup).",
        "Observability (useful logging, metrics, tracing; avoid sensitive data in logs; actionable error messages).",
        "Testability & testing (deterministic seams, unit/integration tests, fixtures/mocks, meaningful coverage).",
        "Dependency & supply-chain hygiene (version constraints, provenance, minimal deps, portability).",
        "Portability & interoperability (standards compliance, platform differences, encoding/locale issues).",
    ]
    REQUIRED_FUNCTIONS = [
        CustomFunctions.Ping,
        CustomFunctions.Code_review,
    ]

    def __init__(self, model_settings: AISettings):
        self.model_settings = model_settings
        self._validate_mandatory_settings()
        self._client = None

    @wrap_ai_exceptions
    def _validate_mandatory_settings(self):
        api_key = self.model_settings.api_key
        assert api_key is not None and api_key, "API key is required"

        name = self.model_settings.model_name
        assert name is not None and name, "Model name is required"

        version = self.model_settings.model_version
        assert version is not None and version, "Model version is required"

    @wrap_ai_exceptions
    def ping(self) -> Response:
        resp = self._get_response(self._get_ping_request())
        return self._transform(resp)

    @wrap_ai_exceptions
    def code_review(self, pr_diffset: DiffSet, *args, **kwargs) -> Response:
        resp = self._get_response(self._get_review_requests(pr_diffset, *args, **kwargs))
        return self._transform(resp)

    def get_code_review_instructions(self, basic_points, review_content, custom_instructions=None):
        user_msg = (
            "BASIC REVIEW INSTRUCTIONS:\n"
            + "\n".join(f"- {point}" for point in basic_points)
            + ("\n\nADDITIONAL CUSTOM INSTRUCTIONS:\n" + custom_instructions if custom_instructions else "")
            + "\n\nIMPORTANT OUTPUT RULES:\n"
            f"- Use the function call to return `{CustomFunctions.Code_review}`: a list of tuples "
            " [line_number:int (1-based, exactly same index as it was in the request per file), line_text:str, suggestion:str].\n"
            "- Only include tuples for lines with a concrete, actionable suggestion.\n"
            "- Keep suggestions concise and specific.\n"
            "- If many lines share the same issue, include a representative subset and list other line numbers.\n"
            "- Do not include overall prose; only return via the function.\n"
            + "\nFILES (numbered per file):\n"
            + self.numbered_code_block(review_content)
        )
        return user_msg

    def numbered_code_block(self, review_content: list[dict]) -> str:
        """
        Render code with explicit 1-based line numbers so the model can reference them precisely.
        """
        parts: list[str] = []
        for content in review_content:
            full_code = content["content"].splitlines()
            only_include_lines: list[int] = content["changed_lines"]
            width = max(2, len(str(len(full_code)))) if full_code else 2
            part_template = f"{FILE} %s\nTYPE: %s"
            parts.append(part_template % (content["file_name"], content["file_type"]))
            if full_code:
                for i, line in enumerate(full_code, start=1):
                    if i in only_include_lines:
                        # include only changes without losing index
                        parts.append(f"{str(i).rjust(width)} | {line}")
            else:
                parts.append("(empty file)")
            parts.append("")  # blank line between files
        return "\n".join(parts).rstrip()

    def get_api_model_name(self):
        return self.model_settings.model_version

    def list_model_api_and_display_names(self) -> list[tuple[str, str]]:
        models = self.list_models()
        res = []
        for m in models:
            if getattr(m, "id") and getattr(m, "display_name"):
                res.append((m.id, m.display_name))
            else:
                self.log.warning(f"Model {m} has no id or display_name")
        return self._filter_api_models(res)

    def _filter_api_models(self, models: list[tuple[str, str]]) -> list[tuple[str, str]]:
        res = []
        filter_models = self._list_non_text_models()
        if not filter_models:
            return models

        for api_name, display_name in models:
            if not any(p in api_name.lower() for p in filter_models):
                res.append((api_name, display_name))

        return res

    def _list_non_text_models(self) -> list[str]:
        return []

    def list_models(self):
        assert self._client is not None, "Client is not initialized"
        return self._client.models.list()

    def _get_response_input(self, request: PingRequest | CodeReviewRequest) -> list[dict[str, str]]:
        """
        It extracts data from internal data classes and transforms it into an input value suitable for the model’s SDK.
        """
        _input = None
        if isinstance(request, CodeReviewRequest):
            return [
                {
                    "role": review.role,
                    "content": review.msg,
                }
                for review in request.content
            ]

        if isinstance(request, PingRequest):
            return [
                {
                    "role": request.role,
                    "content": request.content,
                }
            ]

        raise ValueError("Unknown request type: %s" % type(request))

    def _extract_changes(self, pr_diffset: DiffSet) -> list[Any]:
        review_content = []
        for f in pr_diffset.files:
            target_file_node: FileNode = f.get("target_filenode", None)
            if not target_file_node:
                self.log.warning("No target file node, skipping it")
                continue

            file_name = f.get("patch", {}).get("filename", "unknown")
            if not target_file_node.str_content:
                self.log.debug("No content, skipping file: %s", file_name)
                continue

            modified_lines_nums = []
            for hunk in f.get("hunks", []):
                for line in hunk.get("lines", []):
                    action = line.get("modified", {}).get("action", "").strip()
                    if action and action != "-":
                        modified_lines_nums.append(line.get("modified", {}).get("lineno"))

            file_content = {
                "content": target_file_node.str_content,
                "changed_lines": modified_lines_nums,
                "file_name": file_name,
                "file_type": f.target_file_type,
            }

            review_content.append(file_content)
        return review_content

    def _get_system_message(self) -> str:
        return (
            "You are a meticulous senior code reviewer. "
            "You must ONLY return your findings by calling the provided function "
            f"`{CustomFunctions.Code_review}` with the required JSON shape. "
            "Each tuple = [line_number:int (1-based, GLOBAL across all files), line_text:str, suggestion:str]. "
            "Only include lines that genuinely warrant a suggestion. "
            "The input may mix languages; use each file's TYPE header and the line's syntax to tailor advice. "
            "If a suggestion spans lines or files, attach it to the most relevant line and mention the related line numbers. "
            "Treat 'Plain Text' as docs/logs/config and suggest clarity/safety where relevant."
        )

    def _get_function(self, name):
        """
        this is a GPT/Claude-specific function - AKA instruction for GPT
        doc: https://platform.openai.com/docs/guides/function-calling

        Claude uses the same tools method
        """
        match name:
            case CustomFunctions.Ping:
                return {
                    "type": "function",
                    "name": CustomFunctions.Ping,
                    "description": "Always return pong",
                    "strict": True,
                    "parameters": {
                        "type": "object",
                        "properties": {"response": {"type": "string", "enum": ["pong"]}},
                        "required": ["response"],
                        "additionalProperties": False,
                    },
                }
            case CustomFunctions.Code_review:
                return {
                    "type": "function",
                    "name": CustomFunctions.Code_review,
                    "description": (
                        "Return structured code review per file. Each file uses per-file 1-based "
                        "line numbers. Only include lines that warrant a concrete suggestion."
                    ),
                    "strict": True,
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "response": {
                                "type": "array",
                                "items": {
                                    "type": "object",
                                    "properties": {
                                        "file_name": {"type": "string"},
                                        "review": {
                                            "type": "array",
                                            "items": {
                                                "type": "object",
                                                "properties": {
                                                    "line_number": {"type": "integer", "minimum": 1},
                                                    "line_code": {"type": "string"},
                                                    "suggestion": {"type": "string"},
                                                },
                                                "required": ["line_number", "line_code", "suggestion"],
                                                "additionalProperties": False,
                                            },
                                        },
                                    },
                                    "required": ["file_name", "review"],
                                    "additionalProperties": False,
                                },
                            }
                        },
                        "required": ["response"],
                        "additionalProperties": False,
                    },
                }
            case _:
                raise ValueError("Unknown function name: %s" % name)

    def _get_ping_request(self) -> Request:
        return PingRequest(content="ping", function_name=CustomFunctions.Ping, role="user")

    def _get_review_requests(
        self,
        pr_diffset: DiffSet,
        basic_instructions: Optional[Iterable[str]] = None,
        custom_instructions: Optional[str] = None,
        *args,
        **kwargs,
    ) -> CodeReviewRequest:
        review_content = self._extract_changes(pr_diffset)

        system_msg = self._get_system_message()

        basic_points = list(basic_instructions) if basic_instructions else self.DEFAULT_BASIC_REVIEW_POINTS

        user_msg = self.get_code_review_instructions(basic_points, review_content, custom_instructions)

        return CodeReviewRequest(
            content=[Review(role="system", msg=system_msg), Review(role="user", msg=user_msg)],
            function_name=CustomFunctions.Code_review,
        )

    @abstractmethod
    def _get_response(self, request: PingRequest | CodeReviewRequest) -> Response:
        raise NotImplementedError

    @abstractmethod
    def _transform(self, resp, model=None) -> Response:
        if model is None:
            model = self.get_api_model_name()
        return Response(model=model, message=resp)
