routers/searchterm.py

import os
import re
import random
import asyncio
import httpx
import aiohttp
import trafilatura
import json
import uuid
import time
from pathlib import Path
from urllib.parse import urlparse
from typing import List, Dict, Any, Optional, Set, Tuple
from fastapi import APIRouter, HTTPException, Body
from fastapi.responses import FileResponse
from newspaper import Article
from threading import Timer
from google import genai
from google.genai import types
from asyncio import Queue, create_task, gather
from concurrent.futures import ThreadPoolExecutor
import aiofiles
import ujson  # JSON mais rápido

router = APIRouter()

BRAVE_API_KEY = os.getenv("BRAVE_API_KEY")
if not BRAVE_API_KEY:
    raise ValueError("BRAVE_API_KEY não está definido!")

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
if not GEMINI_API_KEY:
    raise ValueError("GEMINI_API_KEY não está definido!")

BRAVE_SEARCH_URL = "https://api.search.brave.com/res/v1/web/search"
BRAVE_HEADERS = {
    "Accept": "application/json",
    "Accept-Encoding": "gzip",
    "x-subscription-token": BRAVE_API_KEY
}

USER_AGENTS = [
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.0 Safari/605.1.15",
    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
]

BLOCKED_DOMAINS = frozenset({  # frozenset é mais rápido para lookup
    "reddit.com", "www.reddit.com", "old.reddit.com",
    "quora.com", "www.quora.com"
})

MAX_TEXT_LENGTH = 4000
MAX_CONCURRENT_SEARCHES = 30  # Aumentado
MAX_CONCURRENT_EXTRACTIONS = 80  # Aumentado significativamente
EXTRACTION_TIMEOUT = 8  # Reduzido
HTTP_TIMEOUT = 10  # Reduzido

# Diretório para arquivos temporários
TEMP_DIR = Path("/tmp")
TEMP_DIR.mkdir(exist_ok=True)

# Dicionário para controlar arquivos temporários
temp_files = {}

# Pool de threads para operações CPU-intensive
thread_pool = ThreadPoolExecutor(max_workers=20)

# Cache de domínios bloqueados para evitar verificações repetidas
domain_cache = {}

def is_blocked_domain(url: str) -> bool:
    try:
        host = urlparse(url).netloc.lower()
        
        # Cache lookup
        if host in domain_cache:
            return domain_cache[host]
        
        is_blocked = any(host == b or host.endswith("." + b) for b in BLOCKED_DOMAINS)
        domain_cache[host] = is_blocked
        return is_blocked
    except Exception:
        return False

def clamp_text(text: str) -> str:
    if not text or len(text) <= MAX_TEXT_LENGTH:
        return text
    return text[:MAX_TEXT_LENGTH]

def get_realistic_headers() -> Dict[str, str]:
    return {
        "User-Agent": random.choice(USER_AGENTS),
        "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
        "Accept-Language": "en-US,en;q=0.7,pt-BR;q=0.6",
        "Connection": "keep-alive",
        "Accept-Encoding": "gzip, deflate, br",
    }

def delete_temp_file(file_id: str, file_path: Path):
    """Remove arquivo temporário após expiração"""
    try:
        if file_path.exists():
            file_path.unlink()
        temp_files.pop(file_id, None)
        print(f"Arquivo temporário removido: {file_path}")
    except Exception as e:
        print(f"Erro ao remover arquivo temporário: {e}")

async def create_temp_file(data: Dict[str, Any]) -> Dict[str, str]:
    """Cria arquivo temporário assíncrono e agenda sua remoção"""
    file_id = str(uuid.uuid4())
    file_path = TEMP_DIR / f"fontes_{file_id}.txt"
    
    # Salva o JSON no arquivo de forma assíncrona
    async with aiofiles.open(file_path, 'w', encoding='utf-8') as f:
        await f.write(ujson.dumps(data, ensure_ascii=False, indent=2))
    
    # Agenda remoção em 24 horas
    timer = Timer(86400, delete_temp_file, args=[file_id, file_path])
    timer.start()
    
    # Registra o arquivo temporário
    temp_files[file_id] = {
        "path": file_path,
        "created_at": time.time(),
        "timer": timer
    }
    
    return {
        "file_id": file_id,
        "download_url": f"/download-temp/{file_id}",
        "expires_in_hours": 24
    }

async def generate_search_terms(context: str) -> List[str]:
    """Gera termos de pesquisa usando o modelo Gemini"""
    try:
        client = genai.Client(api_key=GEMINI_API_KEY)
        model = "gemini-2.5-flash-lite"
        
        system_prompt = """Com base num contexto inicial, gere termos de pesquisa (até 20 termos, no máximo), em um JSON. Esses textos devem ser interpretados como termos que podem ser usados por outras inteligências artificiais pra pesquisar no google e retornar resultados mais refinados e completos pra busca atual. Analise muito bem o contexto, por exemplo, se está falando de uma série coreana, gere os termos em coreano por que obviamente na mídia coreana terá mais cobertura que a americana, etc.

Deve seguir esse formato: "terms": []

Retorne apenas o JSON, sem mais nenhum texto."""

        contents = [
            types.Content(
                role="user",
                parts=[
                    types.Part.from_text(text="Contexto: Taylor Sheridan's 'Landman' Announces Season 2 Premiere Date"),
                ],
            ),
            types.Content(
                role="model",
                parts=[
                    types.Part.from_text(text='{"terms": [ "imdb landman episodes season 2", "imdb landman series", "landman season 2 release date", "taylor sheridan landman series", "landman season 2 cast sam elliott", "billy bob thornton returns landman", "demi moore landman new season", "andy garcia ali larter landman season 2", "landman texas oil drama", "taylor sheridan tv series schedule", "landman 10 month turnaround new episodes", "landman season 2 november 16 premiere", "sam elliott joins taylor sheridan show", "landman streaming platform premiere", "landman season 2 filming details", "landman new cast and returning actors", "taylor sheridan quick tv show production" ]}'),
                ],
            ),
            types.Content(
                role="user",
                parts=[
                    types.Part.from_text(text="Contexto: Pixar's latest animated feature will arrive on digital (via platforms like Apple TV, Amazon Prime Video, and Fandango at Home) on Aug. 19 and on physical media (4K UHD, Blu-ray and DVD) on Sept. 9. The film has not yet set a Disney+ streaming release date, but that will likely come after the Blu-ray release."),
                ],
            ),
            types.Content(
                role="model",
                parts=[
                    types.Part.from_text(text='{ "terms": [ "pixar elio 2024 movie details", "disney pixar new release elio", "elio animated film august 19 digital", "pixar sci-fi comedy elio home release", "elio movie blu-ray dvd release september", "where to watch elio online", "elio streaming disney plus release date", "elio digital release apple tv amazon prime", "elio physical media 4k uhd blu-ray dvd", "elio movie bonus features", "elio cast voice actors", "elio behind the scenes making of", "elio deleted scenes blu-ray", "elio soundtrack and score", "elio merchandise release date", "upcoming disney pixar movies 2024", "pixar elio critical reviews", "elio movie box office results" ] }'),
                ],
            ),
            types.Content(
                role="user",
                parts=[
                    types.Part.from_text(text=f"Contexto: {context}"),
                ],
            ),
        ]
        
        generate_content_config = types.GenerateContentConfig(
            thinking_config=types.ThinkingConfig(thinking_budget=0),
        )
        
        # Coletamos toda a resposta em stream
        full_response = ""
        for chunk in client.models.generate_content_stream(
            model=model,
            contents=contents,
            config=generate_content_config,
        ):
            if chunk.text:
                full_response += chunk.text
        
        # Tenta extrair o JSON da resposta
        try:
            clean_response = full_response.strip()
            if clean_response.startswith("```json"):
                clean_response = clean_response[7:]
            if clean_response.endswith("```"):
                clean_response = clean_response[:-3]
            clean_response = clean_response.strip()
            
            response_data = ujson.loads(clean_response)
            terms = response_data.get("terms", [])
            
            if not isinstance(terms, list):
                raise ValueError("Terms deve ser uma lista")
            
            return terms[:20]
            
        except (ujson.JSONDecodeError, ValueError) as e:
            print(f"Erro ao parsear resposta do Gemini: {e}")
            return []
        
    except Exception as e:
        print(f"Erro ao gerar termos de pesquisa: {str(e)}")
        return []

async def search_brave_batch(client: httpx.AsyncClient, terms: List[str]) -> List[Tuple[str, List[Dict[str, str]]]]:
    """Busca múltiplos termos em paralelo com otimizações"""
    semaphore = asyncio.Semaphore(MAX_CONCURRENT_SEARCHES)
    
    async def search_single_term(term: str) -> Tuple[str, List[Dict[str, str]]]:
        async with semaphore:
            params = {"q": term, "count": 10, "safesearch": "off", "summary": "false"}
            
            try:
                resp = await client.get(BRAVE_SEARCH_URL, headers=BRAVE_HEADERS, params=params)
                if resp.status_code != 200:
                    return (term, [])

                data = resp.json()
                results = []
                
                if "web" in data and "results" in data["web"]:
                    for item in data["web"]["results"]:
                        url = item.get("url")
                        age = item.get("age", "Unknown")
                        
                        if url and not is_blocked_domain(url):
                            results.append({"url": url, "age": age})

                return (term, results)
            except Exception as e:
                print(f"Erro na busca do termo '{term}': {e}")
                return (term, [])
    
    # Executa todas as buscas em paralelo
    tasks = [search_single_term(term) for term in terms]
    return await gather(*tasks, return_exceptions=False)

def extract_with_trafilatura(html: str) -> str:
    """Extração CPU-intensive executada em thread pool"""
    try:
        extracted = trafilatura.extract(html)
        return extracted.strip() if extracted else ""
    except Exception:
        return ""

def extract_with_newspaper(url: str) -> str:
    """Extração com newspaper executada em thread pool"""
    try:
        art = Article(url)
        art.config.browser_user_agent = random.choice(USER_AGENTS)
        art.config.request_timeout = 6
        art.config.number_threads = 1
        art.download()
        art.parse()
        return (art.text or "").strip()
    except Exception:
        return ""

async def extract_article_text_optimized(url: str, session: aiohttp.ClientSession) -> str:
    """Extração de artigo otimizada com paralelização de métodos"""
    
    # Tentativa 1: Newspaper em thread pool (paralelo com download HTTP)
    newspaper_task = asyncio.create_task(
        asyncio.get_event_loop().run_in_executor(thread_pool, extract_with_newspaper, url)
    )
    
    # Tentativa 2: Download HTTP e trafilatura
    try:
        headers = get_realistic_headers()
        async with session.get(url, headers=headers, timeout=EXTRACTION_TIMEOUT) as resp:
            if resp.status != 200:
                # Aguarda newspaper se HTTP falhou
                newspaper_result = await newspaper_task
                return clamp_text(newspaper_result) if newspaper_result and len(newspaper_result) > 100 else ""
                
            html = await resp.text()
            
            # Verifica paywall rapidamente
            if re.search(r"(paywall|subscribe|metered|registration|captcha|access denied)", 
                        html[:2000], re.I):  # Verifica apenas o início
                newspaper_result = await newspaper_task
                return clamp_text(newspaper_result) if newspaper_result and len(newspaper_result) > 100 else ""

            # Extração com trafilatura em thread pool
            trafilatura_task = asyncio.create_task(
                asyncio.get_event_loop().run_in_executor(thread_pool, extract_with_trafilatura, html)
            )
            
            # Aguarda ambos os métodos e pega o melhor resultado
            newspaper_result, trafilatura_result = await gather(newspaper_task, trafilatura_task)
            
            # Escolhe o melhor resultado
            best_result = ""
            if trafilatura_result and len(trafilatura_result) > 100:
                best_result = trafilatura_result
            elif newspaper_result and len(newspaper_result) > 100:
                best_result = newspaper_result
                
            return clamp_text(best_result) if best_result else ""
                
    except Exception:
        # Se tudo falhar, tenta pelo menos o newspaper
        try:
            newspaper_result = await newspaper_task
            return clamp_text(newspaper_result) if newspaper_result and len(newspaper_result) > 100 else ""
        except Exception:
            return ""

async def process_urls_batch(session: aiohttp.ClientSession, urls_data: List[Tuple[str, str, str]]) -> List[Dict[str, Any]]:
    """Processa URLs em lotes otimizados"""
    semaphore = asyncio.Semaphore(MAX_CONCURRENT_EXTRACTIONS)
    results = []
    used_urls: Set[str] = set()
    
    async def process_single_url(term: str, url: str, age: str) -> Optional[Dict[str, Any]]:
        async with semaphore:
            if url in used_urls:
                return None
                
            text = await extract_article_text_optimized(url, session)
            if text:
                used_urls.add(url)
                return {
                    "term": term,
                    "age": age,
                    "url": url,
                    "text": text
                }
            return None
    
    # Cria todas as tasks de uma vez
    tasks = []
    for term, url, age in urls_data:
        tasks.append(process_single_url(term, url, age))
    
    # Processa tudo em paralelo
    processed_results = await gather(*tasks, return_exceptions=True)
    
    # Filtra resultados válidos
    return [r for r in processed_results if r is not None and not isinstance(r, Exception)]

@router.post("/search-terms")
async def search_terms(payload: Dict[str, str] = Body(...)) -> Dict[str, Any]:
    start_time = time.time()
    
    context = payload.get("context")
    if not context or not isinstance(context, str):
        raise HTTPException(status_code=400, detail="Campo 'context' é obrigatório e deve ser uma string.")
    
    if len(context.strip()) == 0:
        raise HTTPException(status_code=400, detail="Campo 'context' não pode estar vazio.")
    
    print(f"Iniciando geração de termos...")
    # Gera os termos de pesquisa usando o Gemini
    terms = await generate_search_terms(context)
    
    if not terms:
        raise HTTPException(status_code=500, detail="Não foi possível gerar termos de pesquisa válidos.")

    print(f"Termos gerados em {time.time() - start_time:.2f}s. Iniciando buscas...")
    
    # Configurações otimizadas para conexões
    connector = aiohttp.TCPConnector(
        limit=200,  # Aumentado
        limit_per_host=30,  # Aumentado
        ttl_dns_cache=300,
        use_dns_cache=True,
        enable_cleanup_closed=True
    )
    timeout = aiohttp.ClientTimeout(total=HTTP_TIMEOUT, connect=5)
    
    # Cliente HTTP otimizado
    http_client = httpx.AsyncClient(
        timeout=HTTP_TIMEOUT,
        limits=httpx.Limits(
            max_connections=200,  # Aumentado
            max_keepalive_connections=50  # Aumentado
        ),
        http2=True  # Ativa HTTP/2
    )
    
    try:
        async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
            # Fase 1: Busca todos os termos em paralelo
            search_results = await search_brave_batch(http_client, terms)
            print(f"Buscas concluídas em {time.time() - start_time:.2f}s. Iniciando extrações...")
            
            # Fase 2: Prepara dados para extração em lote
            urls_data = []
            for term, results in search_results:
                for result in results:
                    urls_data.append((term, result["url"], result["age"]))
            
            print(f"Processando {len(urls_data)} URLs...")
            
            # Fase 3: Processa todas as URLs em paralelo
            final_results = await process_urls_batch(session, urls_data)
            
            print(f"Extração concluída em {time.time() - start_time:.2f}s. Salvando arquivo...")
    
    finally:
        await http_client.aclose()

    # Fase 4: Cria arquivo temporário assíncrono
    result_data = {"results": final_results}
    temp_file_info = await create_temp_file(result_data)
    
    total_time = time.time() - start_time
    print(f"Processo completo em {total_time:.2f}s")
    
    return {
        "message": "Dados salvos em arquivo temporário",
        "total_results": len(final_results),
        "context": context,
        "generated_terms": terms,
        "file_info": temp_file_info,
        "processing_time": f"{total_time:.2f}s"
    }

@router.get("/download-temp/{file_id}")
async def download_temp_file(file_id: str):
    """Endpoint para download do arquivo temporário"""
    if file_id not in temp_files:
        raise HTTPException(status_code=404, detail="Arquivo não encontrado ou expirado")
    
    file_info = temp_files[file_id]
    file_path = file_info["path"]
    
    if not file_path.exists():
        temp_files.pop(file_id, None)
        raise HTTPException(status_code=404, detail="Arquivo não encontrado")
    
    return FileResponse(
        path=str(file_path),
        filename="fontes.txt",
        media_type="text/plain",
        headers={"Content-Disposition": "attachment; filename=fontes.txt"}
    )