from typing import List, Optional
import requests
import json
from mcp.server.fastmcp import FastMCP

mcp = FastMCP("organizedprogrammers-mcp-server")

@mcp.tool()
def search_arxiv_papers(keyword: str, limit: int = 5) -> str:
    """
    Search papers from arXiv database with specified keywords [optional: a limit of papers the user wants]
    Args: keyword: string, [optional: limit: integer, set limit to 5 if not specified]
    """
    response = requests.post("https://om4r932-arxiv.hf.space/search", headers={
        "Content-Type": "application/json"
    }, data=json.dumps({
        "keyword": keyword,
        "limit": limit
    }), verify=False)

    if response.status_code != 200:
        return "Unable to find papers: error on post()"
    
    responseJson = response.json()

    if responseJson.get("error") or not isinstance(responseJson['message'], dict):
        return f"Unable to find papers: error on API -> {responseJson['message']}"

    if len(responseJson["message"].keys()) == 0:
        return "No papers has been found"
    
    return "\n".join([f"arXiv n°{paper_id} - {paper_meta['title']} by {paper_meta['authors']} : {paper_meta['abstract']}" for paper_id, paper_meta in responseJson['message'].items()])

@mcp.tool()
def locate_3gpp_document(doc_id: str) -> str:
    """
    Find 3GPP document location with the document's ID
    Args: doc_id: string
    """
    response = requests.post("https://organizedprogrammers-3gppdocfinder.hf.space/find", headers={
        "Content-Type": "application/json"
    }, data=json.dumps({
        "doc_id": doc_id
    }), verify=False)

    if response.status_code != 200:
        return f"Unable to find document: {response.status_code} - {response.content}"
    
    responseJson = response.json()

    if responseJson.get("detail"):
        return responseJson['detail']
    
    return f"Document ID {responseJson['doc_id']} version {responseJson['version']} is downloadable via this link: {responseJson['url']}.\n{responseJson['scope']}"

@mcp.tool()
def locate_multiple_3gpp_documents(doc_ids: List[str]) -> str:
    """
    Find 3GPP document location with the document's ID
    Args: doc_id: string
    """
    response = requests.post("https://organizedprogrammers-3gppdocfinder.hf.space/batch", headers={
        "Content-Type": "application/json"
    }, data=json.dumps({
        "doc_ids": doc_ids
    }), verify=False)

    if response.status_code != 200:
        return f"Unable to find document: {response.status_code} - {response.content}"
    
    responseJson = response.json()

    if responseJson.get("detail"):
        return responseJson['detail']
    
    return "\n".join([f"The document {doc_id} is downloadable via this link: {url}" for doc_id, url in responseJson['results']] + [f"We can't find document {doc_id}" for doc_id in responseJson['missing']])

@mcp.tool()
def locate_etsi_document(doc_id: str) -> str:
    """
    Find ETSI document location with the document's ID (starts with SET or SCP)
    Args: doc_id: string
    """
    response = requests.post("https://organizedprogrammers-etsidocfinder.hf.space/find", headers={
        "Content-Type": "application/json"
    }, data=json.dumps({
        "doc_id": doc_id
    }), verify=False)

    if response.status_code != 200:
        return f"Unable to find document: {response.status_code} - {response.content}"
    
    responseJson = response.json()

    if responseJson.get("detail"):
        return responseJson['detail']
    
    return f"Document ID {responseJson['doc_id']} is downloadable via this link: {responseJson['url']}"

@mcp.tool()
def search_3gpp_specifications(keywords: str, threshold: int, release: Optional[str] = "", working_group: Optional[str] = "", spec_type: Optional[str] = "") -> str:
    """
    Search 3GPP specifications with specified keywords and filters using BM25
    Args: keywords: string, threshold: integer 0-100 [default 60], release: optional filter, string [only the number Rel-19 -> '19'], working_group: optional filter, string [options: C1,C2,...,C6,CP or S1,S2,...,S6,SP], spec_type: optional filter, string [either TS (Technical Specification) or TR (Technical Report)]
    For each non-used optional filters, leave a empty string
    """
    body = {"keywords": keywords, "threshold": threshold}
    if release:
        body['release'] = release
    if working_group:
        body['working_group'] = working_group
    if spec_type:
        body['spec_type'] = spec_type

    response = requests.post("https://organizedprogrammers-3gppdocfinder.hf.space/search-spec/experimental", headers={
        "Content-Type": "application/json"
    }, data=json.dumps(body), verify=False)

    if response.status_code != 200:
        return f"Unable to find document: {response.status_code} - {response.content}"
    
    responseJson = response.json()

    if responseJson.get("detail"):
        return responseJson['detail']
    
    return "\n--\n".join([f"3GPP {spec['type']} {spec['id']} version {spec['version']} - {spec['title']} is downloadable via this link: {spec['url']}\n{spec['scope']}" for spec in responseJson['results']])

@mcp.tool()
def ask_questions_to_3gpp_database(question: str, threshold: int = 65, release: Optional[str] = "", working_group: Optional[str] = "", spec_type: Optional[str] = "") -> str:
    """
    Retrieve technical documents sections to help AI answer the user's technical question, if same topic already called, re-use the downloaded documents
    3GPP specifications are used as source documents, using BM25 to filter the documents
    Args: question: string, threshold: integer 0-100 [default 60], release: optional filter, string [only the number Rel-19 -> '19'], working_group: optional filter, string [options: C1,C2,...,C6,CP or S1,S2,...,S6,SP], spec_type: optional filter, string [either TS (Technical Specification) or TR (Technical Report)]
    For each non-used optional filters, leave a empty string
    After extracting the documents, answer to the question with a complete and detailed paragraph with the sources cited
    """
    body = {"question": question, "threshold": threshold}
    if release:
        body['release'] = release
    if working_group:
        body['working_group'] = working_group
    if spec_type:
        body['spec_type'] = spec_type

    response = requests.post("https://organizedprogrammers-3gppdocfinder.hf.space/list-rag-docs", headers={
        "Content-Type": "application/json"
    }, data=json.dumps(body), verify=False)

    if response.status_code != 200:
        return f"Unable to extract documents: {response.status_code}"

    docs = response.json()['output']
    return docs

if __name__ == "__main__":
    mcp.run(transport="stdio")