from __future__ import annotations

import requests
import urllib.parse
from typing import List, Dict, Any, Optional, Literal

class WebSearchJinaComponent:
    def __init__(
        self,
        *,
        api_key: str,
        
        query: str,
        
        top_k: int = 5,
        site_limit: Optional[str] = None,  
        timeout: int = 60,
        
        no_cache: bool = False,
        
        response_format: Literal['markdown', 'html', 'text'] = 'markdown',
        
        retain_images: Literal['none', 'all', 'occluded'] = 'none',
        
        with_generated_alt: bool = True,
        
        with_links_summary: bool = False,
    ):
        self.api_key = api_key
        self.query = query
        self.top_k = top_k
        self.site_limit = site_limit
        self.timeout = timeout
        
        self.no_cache = no_cache
        self.response_format = response_format
        self.retain_images = retain_images
        self.with_generated_alt = with_generated_alt
        self.with_links_summary = with_links_summary
        
        self.base_url = "https://s.jina.ai/"

    def _construct_query(self) -> str:
        final_query = self.query
        if self.site_limit:
            final_query = f"{final_query} site:{self.site_limit}"
        return final_query

    def _get_headers(self) -> Dict[str, str]:
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Accept": "application/json",
            
            "X-Retain-Images": self.retain_images,
            "X-Return-Format": self.response_format,
            "X-With-Generated-Alt": "true" if self.with_generated_alt else "false",
            "X-With-Links-Summary": "true" if self.with_links_summary else "false",
        }
        
        if self.no_cache:
            headers["X-No-Cache"] = "true"
            
        return headers

    def __call__(self) -> List[Dict[str, Any]]:
        query_str = self._construct_query()
        encoded_query = urllib.parse.quote(query_str)
        url = f"{self.base_url}{encoded_query}"
        
        headers = self._get_headers()

        try:
            response = requests.get(url, headers=headers, timeout=self.timeout)
            
            if response.status_code != 200:
                print(f"Error: API returned status {response.status_code}")
                print(f"Message: {response.text}")
                return []
            json_resp = response.json()
            
            items = []
            if isinstance(json_resp, dict):
                if "data" in json_resp and isinstance(json_resp["data"], list):
                    items = json_resp["data"]
                elif "results" in json_resp:
                    items = json_resp["results"]
            elif isinstance(json_resp, list):
                items = json_resp
            if not items:
                return []
            results = items[:self.top_k]
            clean_results = []
            for item in results:
                entry = {
                    "title": item.get("title", "No Title"),
                    "url": item.get("url", ""),
                    "description": item.get("description", ""),
                    "content": item.get("content", ""), 
                }
                if self.with_links_summary and "links" in item:
                    entry["links"] = item["links"]
                    
                clean_results.append(entry)
                
            return clean_results

        except requests.exceptions.Timeout:
            print("Error: Search request timed out.")
            return []
        except requests.exceptions.RequestException as e:
            print(f"Error calling Jina API: {e}")
            return []
        except Exception as e:
            print(f"Unexpected error: {e}")
            return []

    def to_string_context(self) -> str:
        results = self.__call__()
        if not results:
            return "No search results found."

        parts = []
        for i, item in enumerate(results, 1):
            # limit content length to avoid overly long context
            content_preview = item['content'][:2000] if item['content'] else item['description']
            
            parts.append(
                f"### Result {i}\n"
                f"**Title**: {item['title']}\n"
                f"**Source**: {item['url']}\n"
                f"**Content**:\n{content_preview}\n"
            )
        return "\n\n".join(parts)
    
if __name__ == "__main__":
    import os
    JINA_API_KEY = os.getenv("JINA_API_KEY")
    search_query = "DeepSeek-V3 的技术架构特点"
    search_node = WebSearchJinaComponent(
        api_key=JINA_API_KEY,
        query=search_query,
        top_k=3
    )
    print(f"正在搜索: {search_query} ...")
    search_context = search_node.to_string_context()
    print("搜索完成，部分上下文预览：")
    print(search_context + "...\n")