import asyncio
import logging
import random
import time
from asyncio import CancelledError
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Literal

import jwt
import yandexcloud
from aiohttp import ClientSession, ClientTimeout
from pydantic import BaseModel, RootModel
from yandex.cloud.iam.v1.iam_token_service_pb2 import CreateIamTokenRequest
from yandex.cloud.iam.v1.iam_token_service_pb2_grpc import IamTokenServiceStub

from mcp_tracker.tracker.custom.errors import IssueNotFound
from mcp_tracker.tracker.proto.common import YandexAuth
from mcp_tracker.tracker.proto.fields import GlobalDataProtocol
from mcp_tracker.tracker.proto.issues import IssueProtocol
from mcp_tracker.tracker.proto.queues import QueuesProtocol
from mcp_tracker.tracker.proto.types.fields import GlobalField, LocalField
from mcp_tracker.tracker.proto.types.inputs import (
    IssueUpdateFollower,
    IssueUpdateParent,
    IssueUpdatePriority,
    IssueUpdateProject,
    IssueUpdateSprint,
    IssueUpdateType,
)
from mcp_tracker.tracker.proto.types.issue_types import IssueType
from mcp_tracker.tracker.proto.types.issues import (
    ChecklistItem,
    Issue,
    IssueAttachment,
    IssueComment,
    IssueLink,
    IssueTransition,
    Worklog,
)
from mcp_tracker.tracker.proto.types.priorities import Priority
from mcp_tracker.tracker.proto.types.queues import (
    Queue,
    QueueExpandOption,
    QueueVersion,
)
from mcp_tracker.tracker.proto.types.resolutions import Resolution
from mcp_tracker.tracker.proto.types.statuses import Status
from mcp_tracker.tracker.proto.types.users import User
from mcp_tracker.tracker.proto.users import UsersProtocol

QueueList = RootModel[list[Queue]]
LocalFieldList = RootModel[list[LocalField]]
QueueTagList = RootModel[list[str]]
VersionList = RootModel[list[QueueVersion]]
IssueLinkList = RootModel[list[IssueLink]]
IssueList = RootModel[list[Issue]]
IssueCommentList = RootModel[list[IssueComment]]
WorklogList = RootModel[list[Worklog]]
IssueAttachmentList = RootModel[list[IssueAttachment]]
ChecklistItemList = RootModel[list[ChecklistItem]]
GlobalFieldList = RootModel[list[GlobalField]]
StatusList = RootModel[list[Status]]
IssueTypeList = RootModel[list[IssueType]]
PriorityList = RootModel[list[Priority]]
ResolutionList = RootModel[list[Resolution]]
UserList = RootModel[list[User]]
IssueTransitionList = RootModel[list[IssueTransition]]


logger = logging.getLogger(__name__)


class ServiceAccountSettings(BaseModel):
    key_id: str
    service_account_id: str
    private_key: str

    def to_yandexcloud_dict(self) -> dict[str, str]:
        return {
            "id": self.key_id,
            "service_account_id": self.service_account_id,
            "private_key": self.private_key,
        }


class IAMTokenInfo(BaseModel):
    token: str


class ServiceAccountStore:
    DEFAULT_REFRESH_INTERVAL: float = 3500.0
    DEFAULT_RETRY_INTERVAL: float = 10.0

    def __init__(
        self,
        settings: ServiceAccountSettings,
        *,
        refresh_interval: float | None = None,
        retry_interval: float | None = None,
    ):
        self._settings = settings
        self._refresh_interval = refresh_interval or self.DEFAULT_REFRESH_INTERVAL
        self._retry_interval = retry_interval or self.DEFAULT_RETRY_INTERVAL

        self._yc_sdk = yandexcloud.SDK(
            service_account_key=self._settings.to_yandexcloud_dict()
        )
        self._iam_service = self._yc_sdk.client(IamTokenServiceStub)
        self._executor = ThreadPoolExecutor(max_workers=2)
        self._iam_token: IAMTokenInfo | None = None
        self._lock = asyncio.Lock()
        self._refresh_task: asyncio.Task[None] | None = None

    async def prepare(self):
        self._refresh_task = asyncio.create_task(self._refresher())

    async def close(self):
        try:
            if self._refresh_task is not None:
                self._refresh_task.cancel()
                await self._refresh_task
                self._refresh_task = None
        except CancelledError:
            return
        except Exception as e:  # pragma: no cover
            logger.error("error while closing ServiceAccountStore: %s", e)

    async def get_iam_token(self, *, force_refresh: bool = False) -> str:
        if force_refresh or self._iam_token is None:
            async with self._lock:
                if not force_refresh and self._iam_token is not None:
                    return self._iam_token.token

                iam_token = await asyncio.get_running_loop().run_in_executor(
                    self._executor, self._fetch_iam_token, self._settings
                )

                self._iam_token = iam_token
                logger.info("Successfully fetched new IAM token.")

        return self._iam_token.token

    async def _refresher(self):
        while True:
            try:
                await self.get_iam_token(force_refresh=True)
                interval = self._refresh_interval
            except asyncio.CancelledError:  # pragma: no cover
                return
            except Exception as e:
                logger.error("Error refreshing IAM token: %s", e)
                interval = self._retry_interval

            jitter = random.random() * min(interval * 0.1, 100)
            await asyncio.sleep(interval + jitter)

    def _fetch_iam_token(self, service_account: ServiceAccountSettings) -> IAMTokenInfo:
        now = int(time.time())
        payload = {
            "aud": "https://iam.api.cloud.yandex.net/iam/v1/tokens",
            "iss": service_account.service_account_id,
            "iat": now,
            "exp": now + 3600,
        }

        jwt_token = jwt.encode(
            payload=payload,
            key=service_account.private_key,
            algorithm="PS256",
            headers={"kid": service_account.key_id},
        )

        iam_token = self._iam_service.Create(CreateIamTokenRequest(jwt=jwt_token))
        return IAMTokenInfo(token=iam_token.iam_token)


class TrackerClient(QueuesProtocol, IssueProtocol, GlobalDataProtocol, UsersProtocol):
    def __init__(
        self,
        *,
        token: str | None,
        iam_token: str | None = None,
        token_type: Literal["Bearer", "OAuth"] | None = None,
        service_account: ServiceAccountSettings | None = None,
        org_id: str | None = None,
        cloud_org_id: str | None = None,
        base_url: str = "https://api.tracker.yandex.net",
        timeout: float = 10,
    ):
        self._token = token
        self._token_type = token_type
        self._static_iam_token = iam_token
        self._service_account_store: ServiceAccountStore | None = (
            ServiceAccountStore(service_account) if service_account else None
        )
        self._org_id = org_id
        self._cloud_org_id = cloud_org_id

        self._session = ClientSession(
            base_url=base_url,
            timeout=ClientTimeout(total=timeout),
        )

    async def prepare(self):
        if self._service_account_store:
            await self._service_account_store.prepare()

    async def close(self):
        if self._service_account_store:
            await self._service_account_store.close()
        await self._session.close()

    async def _build_headers(self, auth: YandexAuth | None = None) -> dict[str, str]:
        # Priority: OAuth from auth > static OAuth > static IAM token > service account
        auth_header = None

        if auth and auth.token:
            token_type = self._token_type or "OAuth"
            auth_header = f"{token_type} {auth.token}"
        elif self._token:
            token_type = self._token_type or "OAuth"
            auth_header = f"{token_type} {self._token}"
        elif self._static_iam_token:
            auth_header = f"Bearer {self._static_iam_token}"
        elif self._service_account_store is not None:
            iam_token = await self._service_account_store.get_iam_token()
            auth_header = f"Bearer {iam_token}"

        if not auth_header:
            raise ValueError(
                "No authentication method provided. "
                "Provide either OAuth token, IAM token, or use OAuth flow."
            )

        headers = {"Authorization": auth_header}

        # Handle org_id logic
        org_id = auth.org_id if auth and auth.org_id else self._org_id
        cloud_org_id = (
            auth.cloud_org_id if auth and auth.cloud_org_id else self._cloud_org_id
        )

        if org_id and cloud_org_id:
            raise ValueError("Only one of org_id or cloud_org_id should be provided.")

        if org_id:
            headers["X-Org-ID"] = org_id
        elif cloud_org_id:
            headers["X-Cloud-Org-ID"] = cloud_org_id
        else:
            raise ValueError("Either org_id or cloud_org_id must be provided.")

        return headers

    async def queues_list(
        self, per_page: int = 100, page: int = 1, *, auth: YandexAuth | None = None
    ) -> list[Queue]:
        params = {
            "perPage": per_page,
            "page": page,
        }
        async with self._session.get(
            "v3/queues", headers=await self._build_headers(auth), params=params
        ) as response:
            response.raise_for_status()
            return QueueList.model_validate_json(await response.read()).root

    async def queues_get_local_fields(
        self, queue_id: str, *, auth: YandexAuth | None = None
    ) -> list[LocalField]:
        async with self._session.get(
            f"v3/queues/{queue_id}/localFields", headers=await self._build_headers(auth)
        ) as response:
            response.raise_for_status()
            return LocalFieldList.model_validate_json(await response.read()).root

    async def queues_get_tags(
        self, queue_id: str, *, auth: YandexAuth | None = None
    ) -> list[str]:
        async with self._session.get(
            f"v3/queues/{queue_id}/tags", headers=await self._build_headers(auth)
        ) as response:
            response.raise_for_status()
            return QueueTagList.model_validate_json(await response.read()).root

    async def queues_get_versions(
        self, queue_id: str, *, auth: YandexAuth | None = None
    ) -> list[QueueVersion]:
        async with self._session.get(
            f"v3/queues/{queue_id}/versions", headers=await self._build_headers(auth)
        ) as response:
            response.raise_for_status()
            return VersionList.model_validate_json(await response.read()).root

    async def queues_get_fields(
        self, queue_id: str, *, auth: YandexAuth | None = None
    ) -> list[GlobalField]:
        async with self._session.get(
            f"v3/queues/{queue_id}/fields", headers=await self._build_headers(auth)
        ) as response:
            response.raise_for_status()
            return GlobalFieldList.model_validate_json(await response.read()).root

    async def queue_get(
        self,
        queue_id: str,
        *,
        expand: list[QueueExpandOption] | None = None,
        auth: YandexAuth | None = None,
    ) -> Queue:
        params: dict[str, str] = {}
        if expand:
            params["expand"] = ",".join(expand)

        async with self._session.get(
            f"v3/queues/{queue_id}",
            headers=await self._build_headers(auth),
            params=params if params else None,
        ) as response:
            response.raise_for_status()
            return Queue.model_validate_json(await response.read())

    async def get_global_fields(
        self, *, auth: YandexAuth | None = None
    ) -> list[GlobalField]:
        async with self._session.get(
            "v3/fields", headers=await self._build_headers(auth)
        ) as response:
            response.raise_for_status()
            return GlobalFieldList.model_validate_json(await response.read()).root

    async def get_statuses(self, *, auth: YandexAuth | None = None) -> list[Status]:
        async with self._session.get(
            "v3/statuses", headers=await self._build_headers(auth)
        ) as response:
            response.raise_for_status()
            return StatusList.model_validate_json(await response.read()).root

    async def get_issue_types(
        self, *, auth: YandexAuth | None = None
    ) -> list[IssueType]:
        async with self._session.get(
            "v3/issuetypes", headers=await self._build_headers(auth)
        ) as response:
            response.raise_for_status()
            return IssueTypeList.model_validate_json(await response.read()).root

    async def get_priorities(self, *, auth: YandexAuth | None = None) -> list[Priority]:
        async with self._session.get(
            "v3/priorities", headers=await self._build_headers(auth)
        ) as response:
            response.raise_for_status()
            return PriorityList.model_validate_json(await response.read()).root

    async def get_resolutions(
        self, *, auth: YandexAuth | None = None
    ) -> list[Resolution]:
        async with self._session.get(
            "v3/resolutions", headers=await self._build_headers(auth)
        ) as response:
            response.raise_for_status()
            return ResolutionList.model_validate_json(await response.read()).root

    async def issue_get(
        self, issue_id: str, *, auth: YandexAuth | None = None
    ) -> Issue:
        async with self._session.get(
            f"v3/issues/{issue_id}", headers=await self._build_headers(auth)
        ) as response:
            if response.status == 404:
                raise IssueNotFound(issue_id)
            response.raise_for_status()
            return Issue.model_validate_json(await response.read())

    async def issues_get_links(
        self, issue_id: str, *, auth: YandexAuth | None = None
    ) -> list[IssueLink]:
        async with self._session.get(
            f"v3/issues/{issue_id}/links", headers=await self._build_headers(auth)
        ) as response:
            if response.status == 404:
                raise IssueNotFound(issue_id)
            response.raise_for_status()
            return IssueLinkList.model_validate_json(await response.read()).root

    async def issue_get_comments(
        self, issue_id: str, *, auth: YandexAuth | None = None
    ) -> list[IssueComment]:
        async with self._session.get(
            f"v3/issues/{issue_id}/comments", headers=await self._build_headers(auth)
        ) as response:
            if response.status == 404:
                raise IssueNotFound(issue_id)
            response.raise_for_status()
            return IssueCommentList.model_validate_json(await response.read()).root

    async def issues_find(
        self,
        query: str,
        *,
        per_page: int = 15,
        page: int = 1,
        auth: YandexAuth | None = None,
    ) -> list[Issue]:
        params = {
            "perPage": per_page,
            "page": page,
        }

        body: dict[str, Any] = {
            "query": query,
        }

        async with self._session.post(
            "v3/issues/_search",
            headers=await self._build_headers(auth),
            json=body,
            params=params,
        ) as response:
            response.raise_for_status()
            return IssueList.model_validate_json(await response.read()).root

    async def issue_get_worklogs(
        self, issue_id: str, *, auth: YandexAuth | None = None
    ) -> list[Worklog]:
        async with self._session.get(
            f"v3/issues/{issue_id}/worklog", headers=await self._build_headers(auth)
        ) as response:
            if response.status == 404:
                raise IssueNotFound(issue_id)
            response.raise_for_status()
            return WorklogList.model_validate_json(await response.read()).root

    async def issue_get_attachments(
        self, issue_id: str, *, auth: YandexAuth | None = None
    ) -> list[IssueAttachment]:
        async with self._session.get(
            f"v3/issues/{issue_id}/attachments", headers=await self._build_headers(auth)
        ) as response:
            if response.status == 404:
                raise IssueNotFound(issue_id)
            response.raise_for_status()
            return IssueAttachmentList.model_validate_json(await response.read()).root

    async def users_list(
        self, per_page: int = 50, page: int = 1, *, auth: YandexAuth | None = None
    ) -> list[User]:
        params: dict[str, str | int] = {
            "perPage": per_page,
            "page": page,
        }
        async with self._session.get(
            "v3/users", headers=await self._build_headers(auth), params=params
        ) as response:
            response.raise_for_status()
            return UserList.model_validate_json(await response.read()).root

    async def user_get(
        self, user_id: str, *, auth: YandexAuth | None = None
    ) -> User | None:
        async with self._session.get(
            f"v3/users/{user_id}", headers=await self._build_headers(auth)
        ) as response:
            if response.status == 404:
                return None
            response.raise_for_status()
            return User.model_validate_json(await response.read())

    async def user_get_current(self, *, auth: YandexAuth | None = None) -> User:
        async with self._session.get(
            "v3/myself", headers=await self._build_headers(auth)
        ) as response:
            response.raise_for_status()
            return User.model_validate_json(await response.read())

    async def issue_get_checklist(
        self, issue_id: str, *, auth: YandexAuth | None = None
    ) -> list[ChecklistItem]:
        async with self._session.get(
            f"v3/issues/{issue_id}/checklistItems",
            headers=await self._build_headers(auth),
        ) as response:
            if response.status == 404:
                raise IssueNotFound(issue_id)
            response.raise_for_status()
            return ChecklistItemList.model_validate_json(await response.read()).root

    async def issues_count(self, query: str, *, auth: YandexAuth | None = None) -> int:
        body: dict[str, Any] = {
            "query": query,
        }

        async with self._session.post(
            "v3/issues/_count", headers=await self._build_headers(auth), json=body
        ) as response:
            response.raise_for_status()
            return int(await response.text())

    async def issue_create(
        self,
        queue: str,
        summary: str,
        *,
        type: int | None = None,
        description: str | None = None,
        assignee: str | int | None = None,
        priority: str | int | None = None,
        parent: str | None = None,
        sprint: list[str] | None = None,
        auth: YandexAuth | None = None,
        **kwargs: dict[str, Any],
    ) -> Issue:
        body: dict[str, Any] = {
            "queue": queue,
            "summary": summary,
        }

        if type is not None:
            body["type"] = type
        if description is not None:
            body["description"] = description
        if assignee is not None:
            body["assignee"] = assignee
        if priority is not None:
            body["priority"] = priority
        if parent is not None:
            body["parent"] = parent
        if sprint is not None:
            body["sprint"] = sprint

        for k, v in kwargs.items():
            if k not in body:
                body[k] = v

        async with self._session.post(
            "v3/issues", headers=await self._build_headers(auth), json=body
        ) as response:
            response.raise_for_status()
            return Issue.model_validate_json(await response.read())

    async def issue_get_transitions(
        self, issue_id: str, *, auth: YandexAuth | None = None
    ) -> list[IssueTransition]:
        async with self._session.get(
            f"v2/issues/{issue_id}/transitions", headers=await self._build_headers(auth)
        ) as response:
            if response.status == 404:
                raise IssueNotFound(issue_id)
            response.raise_for_status()
            return IssueTransitionList.model_validate_json(await response.read()).root

    async def issue_execute_transition(
        self,
        issue_id: str,
        transition_id: str,
        *,
        comment: str | None = None,
        fields: dict[str, str | int | list[str]] | None = None,
        auth: YandexAuth | None = None,
    ) -> list[IssueTransition]:
        body: dict[str, Any] = {}
        if comment is not None:
            body["comment"] = comment
        if fields is not None:
            body.update(fields)

        async with self._session.post(
            f"v3/issues/{issue_id}/transitions/{transition_id}/_execute",
            headers=await self._build_headers(auth),
            json=body,
        ) as response:
            if response.status == 404:
                raise IssueNotFound(issue_id)
            response.raise_for_status()
            return IssueTransitionList.model_validate_json(await response.read()).root

    async def issue_close(
        self,
        issue_id: str,
        resolution_id: str,
        *,
        comment: str | None = None,
        fields: dict[str, str | int | list[str]] | None = None,
        auth: YandexAuth | None = None,
    ) -> list[IssueTransition]:
        # Fetch transitions and statuses in parallel
        async with asyncio.TaskGroup() as tg:
            transitions_task = tg.create_task(
                self.issue_get_transitions(issue_id, auth=auth)
            )
            statuses_task = tg.create_task(self.get_statuses(auth=auth))

        transitions = transitions_task.result()
        statuses = statuses_task.result()

        # Build a map of status key -> status type
        status_type_map: dict[str, str | None] = {
            status.key: status.type for status in statuses
        }

        # Find a transition to a status with type="done"
        done_transition: IssueTransition | None = None
        for transition in transitions:
            if transition.to and transition.to.key:
                status_type = status_type_map.get(transition.to.key)
                if status_type == "done":
                    done_transition = transition
                    break

        if done_transition is None:
            raise ValueError(
                f"No transition to a 'done' status found for issue {issue_id}. "
                f"Available transitions: {[t.id for t in transitions]}."
            )

        if fields is None:
            fields = {}

        fields["resolution"] = resolution_id

        return await self.issue_execute_transition(
            issue_id,
            done_transition.id,
            comment=comment,
            fields=fields,
            auth=auth,
        )

    async def issue_update(
        self,
        issue_id: str,
        *,
        summary: str | None = None,
        description: str | None = None,
        markup_type: str | None = None,
        parent: IssueUpdateParent | None = None,
        sprint: list[IssueUpdateSprint] | None = None,
        type: IssueUpdateType | None = None,
        priority: IssueUpdatePriority | None = None,
        followers: list[IssueUpdateFollower] | None = None,
        project: IssueUpdateProject | None = None,
        attachment_ids: list[str] | None = None,
        description_attachment_ids: list[str] | None = None,
        tags: list[str] | None = None,
        version: int | None = None,
        auth: YandexAuth | None = None,
        **kwargs: Any,
    ) -> Issue:
        body: dict[str, Any] = {}

        if summary is not None:
            body["summary"] = summary
        if description is not None:
            body["description"] = description
        if markup_type is not None:
            body["markupType"] = markup_type
        if parent is not None:
            body["parent"] = parent.model_dump(exclude_none=True)
        if sprint is not None:
            body["sprint"] = [s.model_dump(exclude_none=True) for s in sprint]
        if type is not None:
            body["type"] = type.model_dump(exclude_none=True)
        if priority is not None:
            body["priority"] = priority.model_dump(exclude_none=True)
        if followers is not None:
            body["followers"] = [f.model_dump(exclude_none=True) for f in followers]
        if project is not None:
            body["project"] = project.model_dump(exclude_none=True)
        if attachment_ids is not None:
            body["attachmentIds"] = attachment_ids
        if description_attachment_ids is not None:
            body["descriptionAttachmentIds"] = description_attachment_ids
        if tags is not None:
            body["tags"] = tags

        for k, v in kwargs.items():
            if k not in body:
                body[k] = v

        params: dict[str, int] = {}
        if version is not None:
            params["version"] = version

        async with self._session.patch(
            f"v3/issues/{issue_id}",
            headers=await self._build_headers(auth),
            json=body,
            params=params if params else None,
        ) as response:
            if response.status == 404:
                raise IssueNotFound(issue_id)
            response.raise_for_status()
            return Issue.model_validate_json(await response.read())
