import asyncio
import httpx
from typing import List, Optional

from .schemes import (
    MessageResponse,
    MessageRequest,
    User,
    UserCreate,
    UserUpdate,
    UserData,
    Message,
    Agent,
    AgentCreate,
    AgentUpdate,
    Unit,
    UnitCreate,
    UnitUpdate,
    RAG,
    RAGCreate,
    RAGUpdate,
)
from .exceptions import (
    NotFound,
    NoAccess,
    Conflict,
    ConnectionError as UnlinkConnectionError,
)


def _raise_for_status(r: httpx.Response) -> None:
    if r.status_code == 401 or r.status_code == 403:
        raise NoAccess()
    if r.status_code == 404:
        raise NotFound()
    if r.status_code == 409:
        raise Conflict()
    r.raise_for_status()


class Client:
    token: str
    base_url: str
    headers: dict

    def __init__(self, token: str, base_url: str) -> None:
        self.token = token
        _base = base_url.rstrip("/")
        self.base_url = f"{_base}/api/v1/"
        self.admin_url = f"{_base}/api/v1/admin/"

        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.token}",
        }

        self.check_connection()

    def check_connection(self) -> None:
        with httpx.Client(headers=self.headers) as client:
            try:
                r = client.get(self.base_url)
                if r.status_code == 401 or r.status_code == 403:
                    raise NoAccess()
                r.raise_for_status()
            except httpx.HTTPError as e:
                raise UnlinkConnectionError(
                    f"Connection error, check your base url: {e}"
                )

    async def send_message(self, user_id: str, message: str, retries: int = 3) -> MessageResponse:
        data = MessageRequest(user_id=str(user_id), request=message)
        for attempt in range(retries):
            try:
                async with httpx.AsyncClient(headers=self.headers, timeout=120) as client:
                    r = await client.post(
                        f"{self.base_url}request",
                        json=data.model_dump(),
                    )
                    _raise_for_status(r)
                    return MessageResponse(**r.json())
            except httpx.ReadError:
                if attempt == retries - 1:
                    raise
                await asyncio.sleep(1)

    # Users

    async def create_user(self, external_id: str) -> User:
        data = UserCreate(external_id=str(external_id))
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.post(
                f"{self.base_url}users",
                json=data.model_dump(),
            )
            _raise_for_status(r)
            return User(**r.json())

    async def get_user(self, user_id: str) -> User:
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.get(f"{self.base_url}users/{user_id}")
            _raise_for_status(r)
            return User(**r.json())

    async def update_user(
        self,
        user_id: str,
        external_id: Optional[str] = None,
        content: Optional[dict] = None,
    ) -> User:
        data = UserUpdate(
            external_id=external_id,
            content=UserData(**content) if content is not None else None,
        )
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.patch(
                f"{self.base_url}users/{user_id}",
                json=data.model_dump(),
            )
            _raise_for_status(r)
            return User(**r.json())

    async def delete_user(self, user_id: str) -> User:
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.delete(f"{self.base_url}users/{user_id}")
            _raise_for_status(r)
            return User(**r.json())

    # History

    async def get_history(self, user_id: str) -> List[Message]:
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.get(f"{self.base_url}users/{user_id}/history")
            _raise_for_status(r)
            return [Message(**m) for m in r.json()]

    async def clear_history(self, user_id: str) -> None:
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.delete(f"{self.base_url}users/{user_id}/history")
            _raise_for_status(r)

    # Admin API
    # Agents

    async def get_agents(self) -> List[Agent]:
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.get(f"{self.admin_url}agents")
            _raise_for_status(r)
            return [Agent(**a) for a in r.json()]

    async def create_agent(self, prompt: str, model: str) -> Agent:
        data = AgentCreate(prompt=prompt, model=model)
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.post(
                f"{self.admin_url}agents",
                json=data.model_dump(),
            )
            _raise_for_status(r)
            return Agent(**r.json())

    async def get_self_agent(self) -> Agent:
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.get(f"{self.admin_url}agents/self")
            _raise_for_status(r)
            return Agent(**r.json())

    async def get_agent(self, agent_id: int) -> Agent:
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.get(f"{self.admin_url}agents/{agent_id}")
            _raise_for_status(r)
            return Agent(**r.json())

    async def update_agent(
        self,
        agent_id: int,
        prompt: Optional[str] = None,
        model: Optional[str] = None,
    ) -> Agent:
        data = AgentUpdate(prompt=prompt, model=model)
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.patch(
                f"{self.admin_url}agents/{agent_id}",
                json=data.model_dump(),
            )
            _raise_for_status(r)
            return Agent(**r.json())

    async def delete_agent(self, agent_id: int) -> Agent:
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.delete(f"{self.admin_url}agents/{agent_id}")
            _raise_for_status(r)
            return Agent(**r.json())

    # Units

    async def get_units(self) -> List[Unit]:
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.get(f"{self.admin_url}units")
            _raise_for_status(r)
            return [Unit(**u) for u in r.json()]

    async def create_unit(
        self,
        agent_id: Optional[int] = None,
        ip: Optional[str] = None,
    ) -> Unit:
        data = UnitCreate(agent_id=agent_id, ip=ip)
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.post(
                f"{self.admin_url}units",
                json=data.model_dump(),
            )
            _raise_for_status(r)
            return Unit(**r.json())

    async def get_self_unit(self) -> Unit:
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.get(f"{self.admin_url}units/self")
            _raise_for_status(r)
            return Unit(**r.json())

    async def get_unit(self, unit_id: int) -> Unit:
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.get(f"{self.admin_url}units/{unit_id}")
            _raise_for_status(r)
            return Unit(**r.json())

    async def update_unit(
        self,
        unit_id: int,
        agent_id: Optional[int] = None,
        ip: Optional[str] = None,
    ) -> Unit:
        data = UnitUpdate(agent_id=agent_id, ip=ip)
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.patch(
                f"{self.admin_url}units/{unit_id}",
                json=data.model_dump(),
            )
            _raise_for_status(r)
            return Unit(**r.json())

    async def delete_unit(self, unit_id: int) -> Unit:
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.delete(f"{self.admin_url}units/{unit_id}")
            _raise_for_status(r)
            return Unit(**r.json())

    # RAG

    async def search_rag(
        self,
        query: str,
        limit: int = 10,
        include_embedding: bool = False,
        agent_id: Optional[int] = None,
    ) -> List[RAG]:
        params = {
            "query": query,
            "limit": limit,
            "include_embedding": include_embedding,
        }
        if agent_id is not None:
            params["agent_id"] = agent_id
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.get(
                f"{self.admin_url}rag",
                params=params,
            )
            _raise_for_status(r)
            return [RAG(**item) for item in r.json()]

    async def create_rag(
        self,
        content: str,
        embedding_content: Optional[str] = None,
        metadata: Optional[dict] = None,
        agent_id: Optional[int] = None,
        retries: int = 3,
    ) -> RAG:
        data = RAGCreate(
            content=content, embedding_content=embedding_content, metadata=metadata
        )

        params = {}
        if agent_id is not None:
            params["agent_id"] = agent_id
        for attempt in range(retries):
            try:
                async with httpx.AsyncClient(headers=self.headers, timeout=120) as client:
                    r = await client.post(
                        f"{self.admin_url}rag",
                        params=params,
                        json=data.model_dump(),
                    )
                    _raise_for_status(r)
                    return RAG(**r.json())
            except httpx.ReadError:
                if attempt == retries - 1:
                    raise
                await asyncio.sleep(1)

    async def get_rag(
        self,
        rag_id: int,
        include_embedding: bool = False,
        agent_id: Optional[int] = None,
    ) -> RAG:
        params = {"include_embedding": include_embedding}
        if agent_id is not None:
            params["agent_id"] = agent_id
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.get(
                f"{self.admin_url}rag/{rag_id}",
                params=params,
            )
            _raise_for_status(r)
            return RAG(**r.json())

    async def update_rag(
        self,
        rag_id: int,
        content: str,
        embedding_content: Optional[str] = None,
        metadata: Optional[dict] = None,
        agent_id: Optional[int] = None,
        retries: int = 3,
    ) -> RAG:
        data = RAGUpdate(
            content=content, embedding_content=embedding_content, metadata=metadata
        )
        params = {}
        if agent_id is not None:
            params["agent_id"] = agent_id
        for attempt in range(retries):
            try:
                async with httpx.AsyncClient(headers=self.headers, timeout=120) as client:
                    r = await client.patch(
                        f"{self.admin_url}rag/{rag_id}",
                        params=params,
                        json=data.model_dump(),
                    )
                    _raise_for_status(r)
                    return RAG(**r.json())
            except httpx.ReadError:
                if attempt == retries - 1:
                    raise
                await asyncio.sleep(1)

    async def delete_rag(self, rag_id: int) -> RAG:
        async with httpx.AsyncClient(headers=self.headers) as client:
            r = await client.delete(f"{self.admin_url}rag/{rag_id}")
            _raise_for_status(r)
            return RAG(**r.json())
