import json
from datetime import datetime
from typing import Dict, List, Optional

import requests
from loguru import logger
from pydantic import BaseModel, Field
from swarms import OpenAIChat

# Ensure loguru logs are saved to a file
logger.add("medinsight_pro_logs.log", rotation="500 MB")


# Define the Pydantic schema for logging metadata
class MedInsightMetadata(BaseModel):
    query: str = Field(
        ..., description="The task or query sent to the agent"
    )
    pubmed_results: Optional[Dict] = Field(
        None, description="Results fetched from PubMed"
    )
    semantic_scholar_results: Optional[Dict] = Field(
        None, description="Results fetched from Semantic Scholar"
    )
    combined_summary: Optional[str] = Field(
        None, description="Final summarized output from the agent"
    )
    timestamp: datetime = Field(
        default_factory=datetime.utcnow,
        description="Time the request was processed",
    )
    status: str = Field(
        ...,
        description="Status of the agent task, e.g., success or failure",
    )

    class Config:
        schema_extra = {
            "example": {
                "query": "COVID-19 treatments",
                "pubmed_results": {
                    "paper_id": {"title": "COVID-19 vaccine trials"}
                },
                "semantic_scholar_results": {
                    "paper_id": {
                        "title": "Effectiveness of mRNA vaccines"
                    }
                },
                "combined_summary": "Recent studies highlight the efficacy of mRNA vaccines in preventing COVID-19...",
                "timestamp": "2023-09-10T12:00:00Z",
                "status": "success",
            }
        }


# Define the MedInsightPro class with customizable options and logging
class MedInsightPro:
    def __init__(
        self,
        openai_api_key: str,
        pubmed_api_key: str = None,
        semantic_scholar_api_key: str = None,
    ):
        self.openai_api_key = openai_api_key
        self.pubmed_api_key = pubmed_api_key
        self.semantic_scholar_api_key = semantic_scholar_api_key

        # Initialize the GPT-4 model
        self.model = OpenAIChat(
            api_key=self.openai_api_key,
            model_name="gpt-4",
            temperature=0.1,  # Low temperature for precise summarization
        )

        # Define the system prompt
        self.system_prompt = """
        You are a highly knowledgeable Medical Research Summarization Agent. 
        Your task is to read large volumes of medical research papers and generate concise summaries. 
        Highlight potential treatments, ongoing clinical trials, medical breakthroughs, and important medical insights.
        You will focus on clarity, relevance, and precision, ensuring the summaries are actionable for doctors and researchers.
        """

        # Initialize the metadata history log
        self.metadata_log: List[MedInsightMetadata] = []

    # Function to access PubMed data
    def fetch_pubmed_data(self, query, max_results=10):
        logger.info(f"Fetching data from PubMed for query: {query}")
        url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
        params = {
            "db": "pubmed",
            "term": query,
            "retmax": max_results,
            "api_key": self.pubmed_api_key,
            "retmode": "json",
        }
        response = requests.get(url, params=params)
        data = response.json()
        ids = data.get("esearchresult", {}).get("idlist", [])

        if ids:
            id_str = ",".join(ids)
            fetch_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi"
            fetch_params = {
                "db": "pubmed",
                "id": id_str,
                "retmode": "json",
            }
            fetch_response = requests.get(
                fetch_url, params=fetch_params
            )
            return fetch_response.json()
        return {}

    # Function to access Semantic Scholar data
    def fetch_semantic_scholar_data(self, query, max_results=10):
        logger.info(
            f"Fetching data from Semantic Scholar for query: {query}"
        )
        url = "https://api.semanticscholar.org/graph/v1/paper/search"
        headers = {"x-api-key": self.semantic_scholar_api_key}
        params = {"query": query, "limit": max_results}
        response = requests.get(url, headers=headers, params=params)
        return response.json()

    # Method to run the agent with a given task
    def run(self, task: str):
        logger.info(f"Running MedInsightPro agent for task: {task}")
        status = "success"
        pubmed_data, semantic_scholar_data = {}, {}
        combined_summary = ""

        try:
            # Fetch data from PubMed
            if self.pubmed_api_key:
                pubmed_data = self.fetch_pubmed_data(task)

            # Fetch data from Semantic Scholar
            if self.semantic_scholar_api_key:
                semantic_scholar_data = (
                    self.fetch_semantic_scholar_data(task)
                )

            # Summarize data with GPT-4
            combined_summary_input = f"PubMed Data: {pubmed_data}\nSemantic Scholar Data: {semantic_scholar_data}"
            combined_summary = self.model.run(combined_summary_input)
            logger.info(f"Summarization completed for task: {task}")
        except Exception as e:
            logger.error(
                f"Error during processing task: {task}. Error: {e}"
            )
            status = "failure"

        # Log metadata
        metadata = MedInsightMetadata(
            query=task,
            pubmed_results=pubmed_data,
            semantic_scholar_results=semantic_scholar_data,
            combined_summary=combined_summary,
            status=status,
        )
        self.metadata_log.append(metadata)

        # Save log to a JSON file
        self.save_metadata_log()

        return combined_summary

    # Method to save the metadata log to a JSON file
    def save_metadata_log(self):
        log_file = "medinsight_pro_history.json"
        with open(log_file, "w") as f:
            json.dump(
                [metadata.dict() for metadata in self.metadata_log],
                f,
                indent=4,
            )
        logger.info(f"Metadata log saved to {log_file}")
