gpt4free

from __future__ import annotations

import os
import uuid
import json
from typing import AsyncIterator

from ..typing import AsyncResult, Messages, Cookies
from ..requests import StreamSession, raise_for_status, sse_stream
from ..cookies import get_cookies, get_cookies_dir
from ..providers.response import (
    ProviderInfo, JsonConversation, JsonRequest, JsonResponse, 
    Reasoning, Sources, SuggestedFollowups, ImageResponse, 
    VariantResponse, YouTubeResponse, TitleGeneration
)
from ..providers.base_provider import AsyncGeneratorProvider, ProviderModelMixin
from .. import debug

# Perplexity API endpoints
PERPLEXITY_URL = "https://www.perplexity.ai"
PERPLEXITY_DOMAIN = ".perplexity.ai"
AUTH_ENDPOINT = f"{PERPLEXITY_URL}/api/auth/session"
QUERY_ENDPOINT = f"{PERPLEXITY_URL}/rest/sse/perplexity_ask"

class Perplexity(AsyncGeneratorProvider, ProviderModelMixin):
    """
    Perplexity provider using browser emulation with HAR file support.
    
    This provider extends the base Perplexity implementation with HAR file support
    for easier authentication management. It uses curl_cffi's Chrome impersonation
    for realistic browser-like requests.
    """
    
    label = "Perplexity"
    url = PERPLEXITY_URL
    cookie_domain = PERPLEXITY_DOMAIN
    working = True
    active_by_default = True
    
    default_model = "auto"
    models = [
        default_model,
        "turbo",
        "gpt41",
        "gpt5",
        "gpt5_thinking",
        "o3",
        "o3pro",
        "claude2",
        "claude37sonnetthinking",
        "claude40opus",
        "claude40opusthinking",
        "claude41opusthinking",
        "claude45sonnet",
        "claude45sonnetthinking",
        "experimental",
        "grok",
        "grok4",
        "gemini2flash",
        "pplx_pro",
        "pplx_pro_upgraded",
        "pplx_alpha",
        "pplx_beta",
        "comet_max_assistant",
        "o3_research",
        "o3pro_research",
        "claude40sonnet_research",
        "claude40sonnetthinking_research",
        "claude40opus_research",
        "claude40opusthinking_research",
        "o3_labs",
        "o3pro_labs",
        "claude40sonnetthinking_labs",
        "claude40opusthinking_labs",
        "o4mini",
        "o1",
        "gpt4o",
        "gpt45",
        "gpt4",
        "o3mini",
        "claude35haiku",
        "llama_x_large",
        "mistral",
        "claude3opus",
        "gemini",
        "pplx_reasoning",
        "r1"
    ]
    fallback_models = ["perplexity", "pplx_pro"]
    model_aliases = {
        "gpt-5": "gpt5",
        "gpt-5-thinking": "gpt5_thinking",
        "r1-1776": "r1",
    }

    @classmethod
    async def get_quota(cls, **kwargs):
        chunks = []
        async for chunk in cls.create_async_generator(
            cls.default_model,
            [{"role": "user", "content": "say only okay"}]
        ):
            if isinstance(chunk, str):
                return {"content": chunk}
        raise RuntimeError("No response")
    
    @classmethod
    async def create_async_generator(
        cls,
        model: str,
        messages: Messages,
        cookies: Cookies = None,
        proxy: str = None,
        conversation: JsonConversation = None,
        prompt: str = None,
        **kwargs
    ) -> AsyncResult:
        """
        Create async generator for Perplexity requests with HAR file support.
        
        Authentication priority:
        1. HAR file cookies (har_and_cookies/perplexity*.har)
        2. Cookie jar from get_cookies()
        """
        if not model:
            model = cls.default_model
        
        # Try to get cookies from HAR file first
        if cookies is None:
            cookies = get_cookies(cls.cookie_domain, False)
            if cookies:
                debug.log(f"Perplexity: Using {len(cookies)} cookies from cookie jar")

        # Initialize conversation if needed
        if conversation is None:
            conversation = JsonConversation(
                frontend_uid=str(uuid.uuid4()),
                frontend_context_uuid=str(uuid.uuid4()),
                visitor_id=str(uuid.uuid4()),
                user_id=None,
                thread_url_slug=None,  # For conversation continuity via Referer header
            )
        
        request_id = str(uuid.uuid4())
        
        # Set referer based on thread_url_slug for conversation continuity
        referer = f"{cls.url}/"
        if hasattr(conversation, 'thread_url_slug') and conversation.thread_url_slug:
            referer = f"{cls.url}/search/{conversation.thread_url_slug}"
            # debug.log(f"Perplexity: Using conversation referer: {referer}")
        
        headers = {
            "accept": "text/event-stream",
            "accept-language": "en-US,en;q=0.9",
            "cache-control": "no-cache",
            "content-type": "application/json",
            "origin": cls.url,
            "referer": referer,
            "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/140.0.0.0 Safari/537.36",
            "x-perplexity-request-reason": "perplexity-query-state-provider",
            "x-request-id": request_id,
        }
        
        # Extract query from messages
        for message in reversed(messages):
            if message["role"] == "user":
                prompt = message["content"]
                if isinstance(prompt, list):
                    prompt = "\n".join([item.get("text", "") for item in prompt]) if prompt else ""
                break
        
        # Use StreamSession with Chrome impersonation
        async with StreamSession(headers=headers, cookies=cookies, proxy=proxy, impersonate="chrome") as session:
            # Get user info if needed
            if conversation.user_id is None:
                try:
                    async with session.get(f"{cls.url}/api/auth/session") as response:
                        await raise_for_status(response)
                        user = await response.json()
                        conversation.user_id = user.get("user", {}).get("id")
                        debug.log(f"Perplexity: User ID: {conversation.user_id}")
                except Exception as e:
                    debug.error(f"Perplexity: Failed to get user info: {e}")
            
            yield conversation
            
            # Determine model
            if model == "auto" or model == "perplexity":
                model = "pplx_pro" if conversation.user_id else "turbo"
            
            yield ProviderInfo(**cls.get_dict(), model=model)
            
            if model in cls.model_aliases:
                model = cls.model_aliases[model]
            
            # Build request data (same as original Perplexity)
            # Check if this is a followup request (has session tokens)
            is_followup = hasattr(conversation, 'last_backend_uuid') and conversation.last_backend_uuid
            
            # debug.log(f"Perplexity: is_followup={is_followup}")
            if is_followup:
                debug.log(f"Perplexity: followup with last_backend_uuid={conversation.last_backend_uuid}, read_write_token={getattr(conversation, 'read_write_token', None)}")
            
            # Generate new frontend_uuid for followup requests (browser does this)
            if is_followup:
                conversation.frontend_uid = str(uuid.uuid4())
            
            if not is_followup:
                data = {
                    "params": {
                        "attachments": [],
                        "language": "en-US",
                        "timezone": "America/Los_Angeles",
                        "search_focus": "internet",
                        "sources": ["web"],
                        "search_recency_filter": None,
                        "frontend_uuid": conversation.frontend_uid,
                        "mode": "copilot",  # Match HAR - use copilot mode
                        "model_preference": model,
                        "is_related_query": False,
                        "is_sponsored": False,
                        "frontend_context_uuid": conversation.frontend_context_uuid,
                        "prompt_source": "user",
                        "query_source": "home",
                        "is_incognito": False,
                        "time_from_first_type": 18361,  # Match HAR value
                        "local_search_enabled": False,
                        "use_schematized_api": True,
                        "send_back_text_in_streaming_api": False,
                        "supported_block_use_cases": [
                            "answer_modes",
                            "media_items",
                            "knowledge_cards",
                            "inline_entity_cards",
                            "place_widgets",
                            "finance_widgets",
                            "prediction_market_widgets",
                            "sports_widgets",
                            "flight_status_widgets",
                            "news_widgets",
                            "shopping_widgets",
                            "jobs_widgets",
                            "search_result_widgets",
                            "inline_images",
                            "inline_assets",
                            "placeholder_cards",
                            "diff_blocks",
                            "inline_knowledge_cards",
                            "entity_group_v2",
                            "refinement_filters",
                            "canvas_mode",
                            "maps_preview",
                            "answer_tabs",
                            "price_comparison_widgets",
                            "preserve_latex",
                            "generic_onboarding_widgets",
                            "in_context_suggestions",
                            "inline_claims"
                        ],
                        "client_coordinates": None,
                        "mentions": [],
                        "dsl_query": prompt,
                        "skip_search_enabled": True,
                        "is_nav_suggestions_disabled": False,
                        "source": "default",
                        "always_search_override": False,
                        "override_no_search": False,
                        "should_ask_for_mcp_tool_confirmation": True,
                        "browser_agent_allow_once_from_toggle": False,
                        "force_enable_browser_agent": False,
                        "supported_features": [
                            "browser_agent_permission_banner_v1.1"
                        ],
                        "version": "2.18"
                    },
                    "query_str": prompt
                }
            else:
                data = {
                    "params": {
                        "last_backend_uuid": getattr(conversation, 'last_backend_uuid', None),
                        "read_write_token": getattr(conversation, 'read_write_token', None),
                        "attachments": [],
                        "language": "en-US",
                        "timezone": "America/Los_Angeles",
                        "search_focus": "internet",
                        "sources": ["web"],
                        "search_recency_filter": None,
                        "frontend_uuid": conversation.frontend_uid,  # New UUID for followup
                        "mode": "copilot",  # Match HAR - use copilot mode
                        "model_preference": model,
                        "is_related_query": False,
                        "is_sponsored": False,
                        "prompt_source": "user",
                        "query_source": "followup",
                        "followup_source": "link",  # Critical for conversation continuity
                        "is_incognito": False,
                        "time_from_first_type": 8758,  # Match HAR value
                        "local_search_enabled": False,
                        "use_schematized_api": True,
                        "send_back_text_in_streaming_api": False,
                        "supported_block_use_cases": [
                            "answer_modes",
                            "media_items",
                            "knowledge_cards",
                            "inline_entity_cards",
                            "place_widgets",
                            "finance_widgets",
                            "prediction_market_widgets",
                            "sports_widgets",
                            "flight_status_widgets",
                            "news_widgets",
                            "shopping_widgets",
                            "jobs_widgets",
                            "search_result_widgets",
                            "inline_images",
                            "inline_assets",
                            "placeholder_cards",
                            "diff_blocks",
                            "inline_knowledge_cards",
                            "entity_group_v2",
                            "refinement_filters",
                            "canvas_mode",
                            "maps_preview",
                            "answer_tabs",
                            "price_comparison_widgets",
                            "preserve_latex",
                            "generic_onboarding_widgets",
                            "in_context_suggestions",
                            "inline_claims"
                        ],
                        "client_coordinates": None,
                        "mentions": [],
                        "dsl_query": prompt,
                        "skip_search_enabled": True,
                        "is_nav_suggestions_disabled": False,
                        "source": "default",
                        "always_search_override": False,
                        "override_no_search": False,
                        "should_ask_for_mcp_tool_confirmation": True,
                        "force_enable_browser_agent": False,
                        "supported_features": [
                            "browser_agent_permission_banner_v1.1"
                        ],
                        "version": "2.18"
                    },
                    "query_str": prompt
                }
            
            yield JsonRequest.from_dict(data)
            
            # Log full request data for debugging
            # debug.log(f"Perplexity: Request data: {json.dumps(data, indent=2, default=str)[:1000]}")
            
            # Send request
            # debug.log(f"Perplexity: Sending request to {QUERY_ENDPOINT}")
            
            async with session.post(QUERY_ENDPOINT, json=data) as response:
                # Process SSE stream
                # debug.log(f"Perplexity: Processing response...")
                await raise_for_status(response)
                
                full_response = ""
                full_reasoning = ""
                sources = []
                
                async for json_data in sse_stream(response):
                    yield JsonResponse.from_dict(json_data)
                    
                    # Capture session tokens for conversation continuity
                    # Note: The 'backend_uuid' field in responses is the backend UUID we need for followups
                    if 'backend_uuid' in json_data:
                        conversation.last_backend_uuid = json_data['backend_uuid']
                    
                    # Only capture read_write_token if we don't have one yet (like a session cookie)
                    if 'read_write_token' in json_data and not hasattr(conversation, 'read_write_token'):
                        conversation.read_write_token = json_data['read_write_token']
                    
                    # Capture thread_url_slug for conversation continuity via Referer header
                    if 'thread_url_slug' in json_data and (not hasattr(conversation, 'thread_url_slug') or not conversation.thread_url_slug):
                        conversation.thread_url_slug = json_data.get('thread_url_slug')

                    if 'thread_title' in json_data:
                        conversation.thread_title = json_data['thread_title']
                        yield TitleGeneration(json_data['thread_title'])
                    for block in json_data.get("blocks", []):
                        # Handle sources
                        if block.get("intended_usage") == "sources_answer_mode":
                            sources = block.get("sources_mode_block", {}).get("web_results", [])
                            continue
                        
                        # Handle media items
                        if block.get("intended_usage") == "media_items":
                            yield VariantResponse("".join([chunk.to_string() if hasattr(chunk, "to_string") else str(chunk) for chunk in [
                                ImageResponse(item.get("url"), item.get("name"), {
                                    "height": item.get("image_height"),
                                    "width": item.get("image_width"),
                                    **item
                                }) if item.get("medium") == "image" else YouTubeResponse(item.get("url").split("=").pop())
                                for item in block.get("media_block", {}).get("media_items", [])
                            ]]))
                            continue
                        
                        # Handle response text
                        for patch in block.get("diff_block", {}).get("patches", []):
                            if patch.get("path") == "/progress":
                                continue
                            
                            value = patch.get("value", "")
                            
                            # Handle reasoning
                            if isinstance(value, dict) and "chunks" in value:
                                value = "".join(value.get("chunks", []))
                            
                            if patch.get("path").startswith("/goals"):
                                if isinstance(value, str):
                                    if value.startswith(full_reasoning):
                                        value = value[len(full_reasoning):]
                                    if value:
                                        yield Reasoning(value)
                                        full_reasoning += value
                                else:
                                    yield Reasoning(status="")
                                continue
                            
                            # Handle regular response
                            if block.get("diff_block").get("field") != "markdown_block":
                                continue
                            
                            value = value.get("answer", "") if isinstance(value, dict) else value
                            
                            if value and isinstance(value, str):
                                if value.startswith(full_response):
                                    value = value[len(full_response):]
                                elif full_response.endswith(value):
                                    value = ""
                                if value:
                                    full_response += value
                                    yield value
                    
                    # Handle follow-ups
                    if "related_query_items" in json_data:
                        followups = []
                        for item in json_data["related_query_items"]:
                            followups.append(item.get("text", ""))
                        yield SuggestedFollowups(followups)
                if sources:
                    yield Sources([{"name": f"Perplexity - {conversation.thread_title}", "url": f"{cls.url}/search/{conversation.thread_url_slug}"}] + sources)
                yield conversation
                
                # debug.log("Perplexity: Request completed successfully")
                # debug.log(f"Perplexity: last_backend_uuid={getattr(conversation, 'last_backend_uuid', None)}, read_write_token={getattr(conversation, 'read_write_token', None)}")