import requests


class AICosmosClient:
    def __init__(
        self,
        base_url: str,
        username: str = None,
        password: str = None,
        api_key: str = None,
    ):
        self.session = requests.Session()
        self.base_url = base_url
        self.username: str = username
        self.password: str = password
        self.api_key: str = api_key

        if self.username and self.password:
            if self.api_key:
                raise RuntimeError("You cannot use both account and API key")
        else:
            if not self.api_key:
                raise RuntimeError("You must provide account or API key")

        self.access_token: str = None

        self._login()

    def _login(self):
        if self.username and self.password:
            login_data = {
                "username": self.username,
                "password": self.password,
                "grant_type": "password",
            }
            login_url = f"{self.base_url}/user/login"
        else:
            login_data = {"api_key": self.api_key}
            login_url = f"{self.base_url}/user/api_login"
        headers = {"Content-Type": "application/x-www-form-urlencoded"}
        try:
            response = self.session.post(
                url=login_url, data=login_data, headers=headers
            )
            if response.status_code == 200:
                token_data = response.json()
                self.access_token = token_data["access_token"]
                return
            else:
                raise ValueError(f"Status code: {response.status_code}")
        except Exception as e:
            raise ValueError(f"Error: {e}")

    def _get_auth_headers(self):
        if not self.access_token:
            raise ValueError("Not logged in")
        return {
            "Authorization": f"Bearer {self.access_token}",
            "Content-Type": "application/json",
        }

    def _get_session_history(self, session_id):
        try:
            response = self.session.get(
                f"{self.base_url}/sessions/{session_id}/history",
                headers=self._get_auth_headers(),
            )
            success = response.status_code == 200
            if success:
                return response.json()
            else:
                raise ValueError(f"Status code: {response.status_code}")
        except Exception as e:
            raise ValueError(f"Error: {e}")

    def create_session(self):
        if not self.access_token:
            raise ValueError("Not logged in")
        try:
            response = self.session.post(
                f"{self.base_url}/sessions/create", headers=self._get_auth_headers()
            )
            if response.status_code == 200:
                response_json = response.json()
                return response_json["session_id"]
            else:
                raise ValueError(f"Status code: {response.status_code}")
        except Exception as e:
            raise ValueError(f"Error: {e}")

    def delete_session(self, session_id: str):
        if not self.access_token:
            raise ValueError("Not logged in")
        try:
            response = self.session.delete(
                f"{self.base_url}/sessions/{session_id}",
                headers=self._get_auth_headers(),
            )
            if response.status_code == 200:
                return
            else:
                raise ValueError(f"Status code: {response.status_code}")
        except Exception as e:
            raise ValueError(f"Error: {e}")

    def get_my_sessions(self):
        """
        Rettrieves session ids and titles.
        """
        if not self.access_token:
            raise ValueError("Not logged in")
        try:
            response = self.session.get(
                f"{self.base_url}/sessions/my_sessions",
                headers=self._get_auth_headers(),
            )
            if response.status_code == 200:
                sessions = response.json()
                self.active_sessions = sessions
                return [
                    {
                        "session_id": session["session_id"],
                        "title": session["environment_info"].get("title", None),
                    }
                    for session in sessions
                ]
            else:
                raise ValueError(f"Status code: {response.status_code}")
        except Exception as e:
            raise ValueError(f"Error: {e}")

    def get_session_history(self, session_id: str):
        """
        Retrieves the conversation history of the given session
        """
        return self.get_session_history(session_id)

    def chat(self, session_id: str, prompt: str, mode: str = "base"):
        """
        Returns conversation history. Avaible modes: base, code, lean
        """
        if not self.access_token:
            raise ValueError("Not logged in")
        data = {
            "user_input": prompt,
            "session_id": session_id,
            "mode": mode,
        }
        try:
            response = self.session.post(
                f"{self.base_url}/chat",
                json=data,
                headers=self._get_auth_headers(),
            )
            success = response.status_code == 200
            if success:
                return response.json()["conversation_history"]
            else:
                raise ValueError(f"Status code: {response.status_code}")
        except Exception as e:
            raise ValueError(f"Error: {e}")

    def get_browser_url(self, session_id: str):
        """
        Retrieves the url that can be opened in a browser.
        """
        try:
            response = self.session.get(
                f"{self.base_url}/sessions/{session_id}/status",
                headers=self._get_auth_headers(),
            )
            success = response.status_code == 200
            if success:
                return response.json().get("browser_url", None)
            else:
                raise ValueError(f"Status code: {response.status_code}")
        except Exception as e:
            raise ValueError(f"Error: {e}")
