"""
Meta AI OpenAI-compatible provider.
Uses Meta AI's chat API via web authentication.
"""

import json
import time
import urllib
import uuid
from typing import Any, Dict, Generator, List, Optional, Union

from curl_cffi import CurlError
from curl_cffi.requests import Session

from webscout.Provider.OPENAI.base import BaseChat, BaseCompletions, OpenAICompatibleProvider
from webscout.Provider.OPENAI.utils import (
    ChatCompletion,
    ChatCompletionChunk,
    ChatCompletionMessage,
    Choice,
    ChoiceDelta,
    CompletionUsage,
)

try:
    from webscout.litagent import LitAgent
except ImportError:
    LitAgent = None


def generate_offline_threading_id() -> str:
    """Generates an offline threading ID."""
    import random
    max_int = (1 << 64) - 1
    mask22_bits = (1 << 22) - 1
    timestamp = int(time.time() * 1000)
    random_value = random.getrandbits(64)
    shifted_timestamp = timestamp << 22
    masked_random = random_value & mask22_bits
    return str((shifted_timestamp | masked_random) & max_int)


def extract_value(text: str, start_str: str, end_str: str) -> str:
    """Helper function to extract a specific value from the given text."""
    start = text.find(start_str) + len(start_str)
    end = text.find(end_str, start)
    return text[start:end]


def format_response(response: dict) -> str:
    """Formats the response from Meta AI."""
    text = ""
    for content in (
        response.get("data", {})
        .get("node", {})
        .get("bot_response_message", {})
        .get("composed_text", {})
        .get("content", [])
    ):
        text += content.get("text", "") + "\n"
    return text.strip()


class Completions(BaseCompletions):
    def __init__(self, client: 'Meta'):
        self._client = client

    def create(
        self,
        *,
        model: str = "meta-ai",
        messages: List[Dict[str, str]],
        max_tokens: Optional[int] = 4096,
        stream: bool = False,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        timeout: Optional[int] = None,
        proxies: Optional[Dict[str, str]] = None,
        **kwargs: Any
    ) -> Union[ChatCompletion, Generator[ChatCompletionChunk, None, None]]:
        """
        Creates a model response for the given chat conversation.
        Mimics openai.chat.completions.create
        """
        request_id = f"chatcmpl-{uuid.uuid4()}"
        created_time = int(time.time())

        # Get just the last user message content for the prompt
        prompt = ""
        for msg in reversed(messages):
            if msg.get("role") == "user":
                prompt = msg.get("content", "")
                break

        if stream:
            return self._create_stream(request_id, created_time, model, prompt, timeout, proxies)
        else:
            return self._create_non_stream(request_id, created_time, model, prompt, timeout, proxies)

    def _create_stream(
        self,
        request_id: str,
        created_time: int,
        model: str,
        prompt: str,
        timeout: Optional[int] = None,
        proxies: Optional[Dict[str, str]] = None
    ) -> Generator[ChatCompletionChunk, None, None]:
        """Implementation for streaming chat completions."""
        try:
            # Get access token if not authenticated
            if not self._client.is_authed:
                self._client.access_token = self._client.get_access_token()
                auth_payload = {"access_token": self._client.access_token}
                url = "https://graph.meta.ai/graphql?locale=user"
            else:
                auth_payload = {"fb_dtsg": self._client.cookies["fb_dtsg"]}
                url = "https://www.meta.ai/api/graphql/"

            if not self._client.external_conversation_id:
                self._client.external_conversation_id = str(uuid.uuid4())

            payload = {
                **auth_payload,
                "fb_api_caller_class": "RelayModern",
                "fb_api_req_friendly_name": "useAbraSendMessageMutation",
                "variables": json.dumps({
                    "message": {"sensitive_string_value": prompt},
                    "externalConversationId": self._client.external_conversation_id,
                    "offlineThreadingId": generate_offline_threading_id(),
                    "suggestedPromptIndex": None,
                    "flashVideoRecapInput": {"images": []},
                    "flashPreviewInput": None,
                    "promptPrefix": None,
                    "entrypoint": "ABRA__CHAT__TEXT",
                    "icebreaker_type": "TEXT",
                    "__relay_internal__pv__AbraDebugDevOnlyrelayprovider": False,
                    "__relay_internal__pv__WebPixelRatiorelayprovider": 1,
                }),
                "server_timestamps": "true",
                "doc_id": "7783822248314888",
            }
            payload = urllib.parse.urlencode(payload)

            headers = {
                "content-type": "application/x-www-form-urlencoded",
                "x-fb-friendly-name": "useAbraSendMessageMutation",
            }
            if self._client.is_authed:
                headers["cookie"] = f'abra_sess={self._client.cookies["abra_sess"]}'

            response = self._client.session.post(
                url,
                headers=headers,
                data=payload,
                stream=True,
                timeout=timeout or self._client.timeout
            )

            if not response.ok:
                raise IOError(f"Meta AI request failed: {response.status_code} - {response.text}")

            prompt_tokens = 0
            completion_tokens = 0
            full_response = ""

            lines = response.iter_lines()
            # Check first line for errors
            first_line = next(lines, None)
            if first_line:
                try:
                    is_error = json.loads(first_line)
                    if len(is_error.get("errors", [])) > 0:
                        raise IOError(f"Meta AI returned error: {first_line}")
                except json.JSONDecodeError:
                    pass

            for line in lines:
                if line:
                    try:
                        json_line = json.loads(line)
                        message_text = format_response(json_line)

                        if not message_text:
                            continue

                        # Yield only the new content (delta)
                        if message_text and len(message_text) > len(full_response):
                            new_content = message_text[len(full_response):]
                            full_response = message_text
                            completion_tokens += 1

                            delta = ChoiceDelta(
                                content=new_content,
                                role="assistant" if completion_tokens == 1 else None,
                                tool_calls=None
                            )
                            choice = Choice(
                                index=0,
                                delta=delta,
                                finish_reason=None,
                                logprobs=None
                            )
                            chunk = ChatCompletionChunk(
                                id=request_id,
                                choices=[choice],
                                created=created_time,
                                model=model,
                                system_fingerprint=None
                            )
                            chunk.usage = {
                                "prompt_tokens": prompt_tokens,
                                "completion_tokens": completion_tokens,
                                "total_tokens": prompt_tokens + completion_tokens,
                                "estimated_cost": None
                            }
                            yield chunk

                    except json.JSONDecodeError:
                        continue

            # Final chunk with finish_reason="stop"
            delta = ChoiceDelta(content=None, role=None, tool_calls=None)
            choice = Choice(index=0, delta=delta, finish_reason="stop", logprobs=None)
            chunk = ChatCompletionChunk(
                id=request_id,
                choices=[choice],
                created=created_time,
                model=model,
                system_fingerprint=None
            )
            chunk.usage = {
                "prompt_tokens": prompt_tokens,
                "completion_tokens": completion_tokens,
                "total_tokens": prompt_tokens + completion_tokens,
                "estimated_cost": None
            }
            yield chunk

        except CurlError as e:
            raise IOError(f"Meta AI stream request failed (CurlError): {e}") from e
        except Exception as e:
            raise IOError(f"Meta AI stream request failed: {e}") from e

    def _create_non_stream(
        self,
        request_id: str,
        created_time: int,
        model: str,
        prompt: str,
        timeout: Optional[int] = None,
        proxies: Optional[Dict[str, str]] = None
    ) -> ChatCompletion:
        """Implementation for non-streaming chat completions."""
        try:
            # Get access token if not authenticated
            if not self._client.is_authed:
                self._client.access_token = self._client.get_access_token()
                auth_payload = {"access_token": self._client.access_token}
                url = "https://graph.meta.ai/graphql?locale=user"
            else:
                auth_payload = {"fb_dtsg": self._client.cookies["fb_dtsg"]}
                url = "https://www.meta.ai/api/graphql/"

            if not self._client.external_conversation_id:
                self._client.external_conversation_id = str(uuid.uuid4())

            payload = {
                **auth_payload,
                "fb_api_caller_class": "RelayModern",
                "fb_api_req_friendly_name": "useAbraSendMessageMutation",
                "variables": json.dumps({
                    "message": {"sensitive_string_value": prompt},
                    "externalConversationId": self._client.external_conversation_id,
                    "offlineThreadingId": generate_offline_threading_id(),
                    "suggestedPromptIndex": None,
                    "flashVideoRecapInput": {"images": []},
                    "flashPreviewInput": None,
                    "promptPrefix": None,
                    "entrypoint": "ABRA__CHAT__TEXT",
                    "icebreaker_type": "TEXT",
                    "__relay_internal__pv__AbraDebugDevOnlyrelayprovider": False,
                    "__relay_internal__pv__WebPixelRatiorelayprovider": 1,
                }),
                "server_timestamps": "true",
                "doc_id": "7783822248314888",
            }
            payload = urllib.parse.urlencode(payload)

            headers = {
                "content-type": "application/x-www-form-urlencoded",
                "x-fb-friendly-name": "useAbraSendMessageMutation",
            }
            if self._client.is_authed:
                headers["cookie"] = f'abra_sess={self._client.cookies["abra_sess"]}'

            response = self._client.session.post(
                url,
                headers=headers,
                data=payload,
                timeout=timeout or self._client.timeout
            )

            if not response.ok:
                raise IOError(f"Meta AI request failed: {response.status_code} - {response.text}")

            # Extract the last response from the streamed data
            raw_response = response.text
            last_response = None
            full_message = ""

            for line in raw_response.split("\n"):
                try:
                    json_line = json.loads(line)
                    bot_response_message = (
                        json_line.get("data", {})
                        .get("node", {})
                        .get("bot_response_message", {})
                    )
                    streaming_state = bot_response_message.get("streaming_state")
                    if streaming_state == "OVERALL_DONE":
                        last_response = json_line
                except json.JSONDecodeError:
                    continue

            if last_response:
                full_message = format_response(last_response)

            message = ChatCompletionMessage(
                role="assistant",
                content=full_message
            )
            choice = Choice(
                index=0,
                message=message,
                finish_reason="stop"
            )
            usage = CompletionUsage(
                prompt_tokens=0,
                completion_tokens=len(full_message.split()),
                total_tokens=len(full_message.split())
            )

            completion = ChatCompletion(
                id=request_id,
                choices=[choice],
                created=created_time,
                model=model,
                usage=usage,
            )
            return completion

        except CurlError as e:
            raise IOError(f"Meta AI non-stream request failed (CurlError): {e}") from e
        except Exception as e:
            raise IOError(f"Meta AI non-stream request failed: {e}") from e


class Chat(BaseChat):
    def __init__(self, client: 'Meta'):
        self.completions = Completions(client)


class Meta(OpenAICompatibleProvider):
    """
    OpenAI-compatible client for Meta AI.
    No API key required - uses web authentication.
    """
    required_auth = False

    AVAILABLE_MODELS = [
        "meta-ai",
        "llama-3"
    ]

    def __init__(
        self,
        fb_email: str = None,
        fb_password: str = None,
        timeout: int = 60,
        proxies: dict = None,
        browser: str = "chrome"
    ):
        """
        Initialize the Meta AI OpenAI-compatible client.

        Args:
            fb_email: Optional Facebook email for authenticated access
            fb_password: Optional Facebook password for authenticated access
            timeout: Request timeout in seconds
            proxies: Proxy settings
            browser: Browser type for fingerprinting
        """
        self.fb_email = fb_email
        self.fb_password = fb_password
        self.timeout = timeout
        self.proxies = proxies or {}

        self.session = Session()
        if LitAgent:
            agent = LitAgent()
            self.session.headers.update({
                "user-agent": agent.random()
            })

        self.access_token = None
        self.is_authed = fb_password is not None and fb_email is not None
        self.cookies = self._get_cookies()
        self.external_conversation_id = None
        self.offline_threading_id = None

        if proxies:
            self.session.proxies = proxies

        # Initialize chat interface
        self.chat = Chat(self)

    def _get_cookies(self) -> dict:
        """Extracts necessary cookies from the Meta AI main page."""
        headers = {}

        # Import Facebook login if needed
        if self.fb_email is not None and self.fb_password is not None:
            try:
                from webscout.Provider.meta import get_fb_session
                fb_session = get_fb_session(self.fb_email, self.fb_password, self.proxies)
                headers = {"cookie": f"abra_sess={fb_session['abra_sess']}"}
            except Exception:
                pass

        response = self.session.get(
            url="https://www.meta.ai/",
            headers=headers,
            proxies=self.proxies,
        )

        cookies = {
            "_js_datr": extract_value(
                response.text, start_str='_js_datr":{"value":"', end_str='",'
            ),
            "datr": extract_value(
                response.text, start_str='datr":{"value":"', end_str='",'
            ),
            "lsd": extract_value(
                response.text, start_str='"LSD",[],{"token":"', end_str='"}'
            ),
            "fb_dtsg": extract_value(
                response.text, start_str='DTSGInitData",[],{"token":"', end_str='"'
            ),
        }

        if len(headers) > 0:
            try:
                from webscout.Provider.meta import get_fb_session
                fb_session = get_fb_session(self.fb_email, self.fb_password, self.proxies)
                cookies["abra_sess"] = fb_session["abra_sess"]
            except Exception:
                pass
        else:
            cookies["abra_csrf"] = extract_value(
                response.text, start_str='abra_csrf":{"value":"', end_str='",'
            )
        return cookies

    def get_access_token(self) -> str:
        """Retrieves an access token using Meta's authentication API."""
        if self.access_token:
            return self.access_token

        url = "https://www.meta.ai/api/graphql/"
        payload = {
            "lsd": self.cookies["lsd"],
            "fb_api_caller_class": "RelayModern",
            "fb_api_req_friendly_name": "useAbraAcceptTOSForTempUserMutation",
            "variables": {
                "dob": "1999-01-01",
                "icebreaker_type": "TEXT",
                "__relay_internal__pv__WebPixelRatiorelayprovider": 1,
            },
            "doc_id": "7604648749596940",
        }
        payload = urllib.parse.urlencode(payload)
        headers = {
            "content-type": "application/x-www-form-urlencoded",
            "cookie": f'_js_datr={self.cookies["_js_datr"]}; '
            f'abra_csrf={self.cookies.get("abra_csrf", "")}; datr={self.cookies["datr"]};',
            "sec-fetch-site": "same-origin",
            "x-fb-friendly-name": "useAbraAcceptTOSForTempUserMutation",
        }

        response = self.session.post(url, headers=headers, data=payload)

        try:
            auth_json = response.json()
        except json.JSONDecodeError:
            raise IOError(
                "Unable to receive a valid response from Meta AI. "
                "This is likely due to your region being blocked."
            )

        access_token = auth_json["data"]["xab_abra_accept_terms_of_service"][
            "new_temp_user_auth"
        ]["access_token"]

        time.sleep(1)  # Meta needs time to register cookies
        return access_token

    @property
    def models(self):
        class _ModelList:
            def list(inner_self):
                return type(self).AVAILABLE_MODELS
        return _ModelList()


if __name__ == "__main__":
    # Example usage - no API key required
    client = Meta()
    print(f"Available models: {client.models.list()}")

    # Test non-streaming
    response = client.chat.completions.create(
        model="meta-ai",
        messages=[{"role": "user", "content": "Hello!"}],
        stream=False
    )
    print(f"Response: {response.choices[0].message.content}")

    # Test streaming
    print("\nStreaming response:")
    for chunk in client.chat.completions.create(
        model="meta-ai",
        messages=[{"role": "user", "content": "Tell me a short joke"}],
        stream=True
    ):
        if chunk.choices[0].delta.content:
            print(chunk.choices[0].delta.content, end="", flush=True)
    print()
