import asyncio
from contextlib import asynccontextmanager
from typing import Optional
from fastapi import APIRouter, FastAPI, HTTPException
from fastapi.routing import APIRouter
import httpx
from httpx import HTTPStatusError
from pydantic import BaseModel, Field
from playwright.async_api import async_playwright, Browser, BrowserContext, Page
import logging
import uvicorn

from scrap import PatentScrapBulkResponse, PatentScrapResult, scrap_patent_async, scrap_patent_bulk_async
from serp import SerpQuery, SerpResults, query_arxiv, query_bing_search, query_brave_search, query_ddg_search, query_google_patents, query_google_scholar
from ops import OPSBulkResponse, OPSNotConfigured, ops_scrap_patent, ops_scrap_patent_bulk, ops_search, token_manager as ops_token_manager
from utils import log_gathered_exceptions

logging.basicConfig(
    level=logging.INFO,
    format='[%(asctime)s][%(levelname)s][%(filename)s:%(lineno)d]: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)

# playwright global context
playwright = None
pw_browser: Optional[Browser] = None

# httpx client
httpx_client = httpx.AsyncClient(timeout=30, limits=httpx.Limits(
    max_connections=30, max_keepalive_connections=20))


@asynccontextmanager
async def api_lifespan(app: FastAPI):
    global playwright, pw_browser
    try:
        playwright = await async_playwright().start()
        pw_browser = await playwright.chromium.launch(headless=True)
        logging.info("Playwright browser started.")
    except Exception as e:
        logging.warning(f"Playwright unavailable, browser-based endpoints will fail: {e}")
    yield

    if pw_browser:
        await pw_browser.close()
    if playwright:
        await playwright.stop()

app = FastAPI(lifespan=api_lifespan, docs_url="/",
              title="SERPent", description=open("docs/docs.md").read())

# Router for scrapping related endpoints
scrap_router = APIRouter(prefix="/scrap", tags=["scrapping"])
# Router for SERP-scrapping related endpoints
serp_router = APIRouter(prefix="/serp", tags=["serp scrapping"])
# Router for EPO OPS (official patent API) endpoints
ops_router = APIRouter(prefix="/ops", tags=["EPO OPS"])

# ===================== Search endpoints =====================


@serp_router.post("/search_scholar")
async def search_google_scholar(params: SerpQuery):
    """Queries google scholar for the specified query"""
    logging.info(f"Searching Google Scholar for queries: {params.queries}")
    results = await asyncio.gather(*[query_google_scholar(pw_browser, q, params.n_results) for q in params.queries], return_exceptions=True)
    log_gathered_exceptions(results, "google scholar search", params)

    # Filter out exceptions and flatten the results
    filtered_results = [r for r in results if not isinstance(r, Exception)]
    flattened_results = [
        item for sublist in filtered_results for item in sublist]

    # all queries failed, return the last exception
    if len(filtered_results) == 0:
        return SerpResults(results=[], error=str(results[-1]))

    return SerpResults(results=flattened_results, error=None)


@serp_router.post("/search_arxiv")
async def search_arxiv(params: SerpQuery):
    """Searches arxiv for the specified queries and returns the found documents."""
    logging.info(f"Searching Arxiv for queries: {params.queries}")
    results = await asyncio.gather(*[query_arxiv(httpx_client, q, params.n_results) for q in params.queries], return_exceptions=True)
    log_gathered_exceptions(results, "arxiv search", params)

    filtered_results = [r for r in results if not isinstance(r, Exception)]
    flattened_results = [
        item for sublist in filtered_results for item in sublist]

    if len(filtered_results) == 0:
        return SerpResults(results=[], error=str(results[-1]))

    return SerpResults(results=flattened_results, error=None)


@serp_router.post("/search_patents")
async def search_patents(params: SerpQuery) -> SerpResults:
    """Searches google patents for the specified queries and returns the found documents.

    Falls back to the EPO OPS API for any query Google Patents returns nothing
    for, when OPS credentials are configured.
    """
    logging.info(f"Searching Google Patents for queries: {params.queries}")
    results = await asyncio.gather(*[query_google_patents(pw_browser, q, params.n_results) for q in params.queries], return_exceptions=True)
    log_gathered_exceptions(results, "google patent search", params)

    # Fall back to OPS for queries that errored or returned no results.
    if ops_token_manager.configured:
        for i, res in enumerate(results):
            if isinstance(res, Exception) or not res:
                try:
                    logging.info(
                        f"Google Patents empty for `{params.queries[i]}`, trying OPS.")
                    results[i] = await ops_search(httpx_client, params.queries[i], params.n_results)
                except Exception as e:
                    logging.warning(f"OPS fallback failed: {e}")

    # Filter out exceptions and flatten the results
    filtered_results = [r for r in results if not isinstance(r, Exception)]
    flattened_results = [
        item for sublist in filtered_results for item in sublist]

    # all queries failed, return the last exception
    if len(filtered_results) == 0:
        return SerpResults(results=[], error=str(results[-1]))

    return SerpResults(results=flattened_results, error=None)


@serp_router.post("/search_brave")
async def search_brave(params: SerpQuery) -> SerpResults:
    """Searches brave search for the specified queries and returns the found documents."""
    logging.info(f"Searching Brave Search for queries: {params.queries}")
    results = await asyncio.gather(*[query_brave_search(pw_browser, q, params.n_results) for q in params.queries], return_exceptions=True)
    log_gathered_exceptions(results, "brave search", params)

    # Filter out exceptions and flatten the results
    filtered_results = [r for r in results if not isinstance(r, Exception)]
    flattened_results = [
        item for sublist in filtered_results for item in sublist]

    # all queries failed, return the last exception
    if len(filtered_results) == 0:
        return SerpResults(results=[], error=str(results[-1]))

    return SerpResults(results=flattened_results, error=None)


@serp_router.post("/search_bing")
async def search_bing(params: SerpQuery) -> SerpResults:
    """Searches Bing search for the specified queries and returns the found documents."""
    logging.info(f"Searching Bing Search for queries: {params.queries}")
    results = await asyncio.gather(*[query_bing_search(pw_browser, q, params.n_results) for q in params.queries], return_exceptions=True)
    log_gathered_exceptions(results, "bing search", params)

    # Filter out exceptions and flatten the results
    filtered_results = [r for r in results if not isinstance(r, Exception)]
    flattened_results = [
        item for sublist in filtered_results for item in sublist]

    # all queries failed, return the last exception
    if len(filtered_results) == 0:
        return SerpResults(results=[], error=str(results[-1]))

    return SerpResults(results=flattened_results, error=None)


@serp_router.post("/search_duck")
async def search_duck(params: SerpQuery) -> SerpResults:
    """Searches duckduckgo for the specified queries and returns the found documents"""
    logging.info(f"Searching DuckDuckGo for queries: {params.queries}")
    results = await asyncio.gather(*[query_ddg_search(q, params.n_results) for q in params.queries], return_exceptions=True)
    log_gathered_exceptions(results, "duckduckgo search", params)

    # Filter out exceptions and flatten the results
    filtered_results = [r for r in results if not isinstance(r, Exception)]
    flattened_results = [
        item for sublist in filtered_results for item in sublist]

    # all queries failed, return the last exception
    if len(filtered_results) == 0:
        return SerpResults(results=[], error=str(results[-1]))

    return SerpResults(results=flattened_results, error=None)


@serp_router.post("/search")
async def search(params: SerpQuery):
    """Attempts to search the specified queries using ALL backends"""
    results = []

    for q in params.queries:
        try:
            logging.info(f"Querying DDG with query: `{q}`")
            res = await query_ddg_search(q, params.n_results)
            results.extend(res)
            continue
        except Exception as e:
            logging.error(f"Failed to query DDG with query `{q}`: {e}")
            logging.info("Trying with next browser backend.")

        try:
            logging.info(f"Querying Brave Search with query: `{q}`")
            res = await query_brave_search(pw_browser, q, params.n_results)
            results.extend(res)
            continue
        except Exception as e:
            logging.error(
                f"Failed to query Brave Search with query `{q}`: {e}")
            logging.info("Trying with next browser backend.")

        try:
            logging.info(f"Querying Bing with query: `{q}`")
            res = await query_bing_search(pw_browser, q, params.n_results)
            results.extend(res)
            continue
        except Exception as e:
            logging.error(f"Failed to query Bing search with query `{q}`: {e}")
            logging.info("Trying with next browser backend.")

        if len(results) == 0:
            return SerpResults(results=[], error="All backends are rate-limited.")

    return SerpResults(results=results, error=None)

# =========================== Scrapping endpoints ===========================


@scrap_router.get("/scrap_patent/{patent_id}")
async def scrap_patent(patent_id: str) -> PatentScrapResult:
    """Scraps the specified patent from Google Patents.

    Falls back to the EPO OPS API (which covers patents missing from Google
    Patents) when the scrape fails and OPS credentials are configured.
    """
    try:
        return await scrap_patent_async(httpx_client, f"https://patents.google.com/patent/{patent_id}/en")
    except HTTPStatusError as e:
        logging.warning(
            f"Google Patents returned {e.response.status_code} for {patent_id}.")
    except Exception as e:
        logging.warning(f"Failed to scrap patent {patent_id}: {e}")

    if not ops_token_manager.configured:
        raise HTTPException(
            status_code=404,
            detail=f"Patent '{patent_id}' not found on Google Patents (EPO OPS fallback not configured).")

    try:
        logging.info(f"Trying OPS for patent {patent_id}.")
        return await ops_scrap_patent(httpx_client, patent_id)
    except HTTPStatusError as e:
        if e.response.status_code == 404:
            raise HTTPException(
                status_code=404,
                detail=f"Patent '{patent_id}' not found on Google Patents or EPO OPS.")
        raise HTTPException(
            status_code=502,
            detail=f"EPO OPS returned {e.response.status_code} for '{patent_id}'.")
    except Exception as e:
        logging.warning(f"OPS fallback failed for {patent_id}: {e}")
        raise HTTPException(
            status_code=502,
            detail=f"EPO OPS request failed for '{patent_id}': {e}")


class ScrapPatentsRequest(BaseModel):
    """Request model for scrapping multiple patents."""
    patent_ids: list[str] = Field(...,
                                  description="List of patent IDs to scrap")


@scrap_router.post("/scrap_patents_bulk", response_model=PatentScrapBulkResponse)
async def scrap_patents(params: ScrapPatentsRequest) -> PatentScrapBulkResponse:
    """Scraps multiple patents from Google Patents."""
    patents = await scrap_patent_bulk_async(httpx_client, params.patent_ids)
    return patents

# =========================== EPO OPS endpoints ===========================


@ops_router.post("/search")
async def ops_keyword_search(params: SerpQuery) -> SerpResults:
    """Keyword-searches patents via the official EPO OPS API."""
    logging.info(f"Searching EPO OPS for queries: {params.queries}")
    results = await asyncio.gather(*[ops_search(httpx_client, q, params.n_results) for q in params.queries], return_exceptions=True)
    log_gathered_exceptions(results, "OPS search", params)

    filtered_results = [r for r in results if not isinstance(r, Exception)]
    flattened_results = [
        item for sublist in filtered_results for item in sublist]

    if len(filtered_results) == 0:
        return SerpResults(results=[], error=str(results[-1]))

    return SerpResults(results=flattened_results, error=None)


@ops_router.get("/scrap_patent/{patent_id}")
async def ops_get_patent(patent_id: str) -> PatentScrapResult:
    """Retrieves a patent (biblio + abstract + claims + description) via EPO OPS."""
    if not ops_token_manager.configured:
        raise HTTPException(
            status_code=503,
            detail="EPO OPS is not configured (OPS_CONSUMER_KEY / OPS_CONSUMER_SECRET missing).")
    try:
        return await ops_scrap_patent(httpx_client, patent_id)
    except OPSNotConfigured:
        raise HTTPException(status_code=503, detail="EPO OPS is not configured.")
    except HTTPStatusError as e:
        if e.response.status_code == 404:
            raise HTTPException(
                status_code=404, detail=f"Patent '{patent_id}' not found in EPO OPS.")
        raise HTTPException(
            status_code=502,
            detail=f"EPO OPS returned {e.response.status_code} for '{patent_id}'.")
    except Exception as e:
        logging.warning(f"Failed to retrieve patent {patent_id} from OPS: {e}")
        raise HTTPException(
            status_code=502, detail=f"EPO OPS request failed for '{patent_id}': {e}")


@ops_router.post("/scrap_patents_bulk", response_model=OPSBulkResponse)
async def ops_get_patents_bulk(params: ScrapPatentsRequest) -> OPSBulkResponse:
    """Retrieves multiple patents via EPO OPS."""
    return await ops_scrap_patent_bulk(httpx_client, params.patent_ids)

# ===============================================================================

app.include_router(serp_router)
app.include_router(scrap_router)
app.include_router(ops_router)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)