import os
import requests
from dotenv import load_dotenv
import json

load_dotenv()

# Retrieve API key securely from .env
API_KEY = os.getenv("SWARMS_API_KEY")
BASE_URL = "https://api.swarms.world"

# Headers for secure API communication
headers = {"x-api-key": API_KEY, "Content-Type": "application/json"}


def create_medical_swarm(patient_case: str):
    """
    Constructs and triggers a full-stack medical swarm consisting of three agents:
    Diagnostic Specialist, Medical Coder, and Treatment Advisor.
    Each agent is provided with a comprehensive, detailed system prompt to ensure high reliability.
    """

    payload = {
        "swarm_name": "Enhanced Medical Diagnostic Swarm",
        "description": "A swarm of agents specialized in performing comprehensive medical diagnostics, analysis, and coding.",
        "agents": [
            {
                "agent_name": "Diagnostic Specialist",
                "description": "Agent specialized in analyzing patient history, symptoms, lab results, and imaging data to produce accurate diagnoses.",
                "system_prompt": (
                    "You are an experienced, board-certified medical diagnostician with over 20 years of clinical practice. "
                    "Your role is to analyze all available patient information—including history, symptoms, lab tests, and imaging results—"
                    "with extreme attention to detail and clinical nuance. Provide a comprehensive differential diagnosis considering "
                    "common, uncommon, and rare conditions. Always cross-reference clinical guidelines and evidence-based medicine. "
                    "Explain your reasoning step by step and provide a final prioritized list of potential diagnoses along with their likelihood. "
                    "Consider patient demographics, comorbidities, and risk factors. Your diagnosis should be reliable, clear, and actionable."
                ),
                "model_name": "openai/gpt-4o",
                "role": "worker",
                "max_loops": 2,
                "max_tokens": 4000,
                "temperature": 0.3,
                "auto_generate_prompt": False,
                "tools_dictionary": [
                    {
                        "type": "function",
                        "function": {
                            "name": "search_topic",
                            "description": "Conduct an in-depth search on a specified topic or subtopic, generating a comprehensive array of highly detailed search queries tailored to the input parameters.",
                            "parameters": {
                                "type": "object",
                                "properties": {
                                    "depth": {
                                        "type": "integer",
                                        "description": "Indicates the level of thoroughness for the search. Values range from 1 to 3, where 1 represents a superficial search and 3 signifies an exploration of the topic.",
                                    },
                                    "detailed_queries": {
                                        "type": "array",
                                        "description": "An array of highly specific search queries that are generated based on the input query and the specified depth. Each query should be designed to elicit detailed and relevant information from various sources.",
                                        "items": {
                                            "type": "string",
                                            "description": "Each item in this array should represent a unique search query that targets a specific aspect of the main topic, ensuring a comprehensive exploration of the subject matter.",
                                        },
                                    },
                                },
                                "required": [
                                    "depth",
                                    "detailed_queries",
                                ],
                            },
                        },
                    },
                ],
            },
            {
                "agent_name": "Medical Coder",
                "description": "Agent responsible for translating medical diagnoses and procedures into accurate standardized medical codes (ICD-10, CPT, etc.).",
                "system_prompt": (
                    "You are a certified and experienced medical coder, well-versed in ICD-10, CPT, and other coding systems. "
                    "Your task is to convert detailed medical diagnoses and treatment procedures into precise, standardized codes. "
                    "Consider all aspects of the clinical documentation including severity, complications, and comorbidities. "
                    "Provide clear explanations for the codes chosen, referencing the latest coding guidelines and payer policies where relevant. "
                    "Your output should be comprehensive, reliable, and fully compliant with current medical coding standards."
                ),
                "model_name": "openai/gpt-4o",
                "role": "worker",
                "max_loops": 1,
                "max_tokens": 3000,
                "temperature": 0.2,
                "auto_generate_prompt": False,
                "tools_dictionary": [],
            },
            {
                "agent_name": "Treatment Advisor",
                "description": "Agent dedicated to suggesting evidence-based treatment options, including pharmaceutical and non-pharmaceutical interventions.",
                "system_prompt": (
                    "You are a highly knowledgeable medical treatment specialist with expertise in the latest clinical guidelines and research. "
                    "Based on the diagnostic conclusions provided, your task is to recommend a comprehensive treatment plan. "
                    "Your suggestions should include first-line therapies, potential alternative treatments, and considerations for patient-specific factors "
                    "such as allergies, contraindications, and comorbidities. Explain the rationale behind each treatment option and reference clinical guidelines where applicable. "
                    "Your recommendations should be reliable, detailed, and clearly prioritized based on efficacy and safety."
                ),
                "model_name": "openai/gpt-4o",
                "role": "worker",
                "max_loops": 1,
                "max_tokens": 5000,
                "temperature": 0.3,
                "auto_generate_prompt": False,
                "tools_dictionary": [],
            },
        ],
        "max_loops": 3,
        "swarm_type": "SequentialWorkflow",
        "task": patient_case,
    }

    # Payload includes the patient case as the task to be processed by the swarm

    response = requests.post(
        f"{BASE_URL}/v1/swarm/completions",
        headers=headers,
        json=payload,
    )

    if response.status_code == 200:
        print("Swarm successfully executed!")
        return json.dumps(response.json(), indent=4)
    else:
        print(f"Error {response.status_code}: {response.text}")
        return None


# Example Patient Task for the Swarm to diagnose and analyze
if __name__ == "__main__":
    patient_case = (
        "Patient is a 55-year-old male presenting with severe chest pain, shortness of breath, elevated blood pressure, "
        "nausea, and a family history of cardiovascular disease. Blood tests show elevated troponin levels, and EKG indicates ST-segment elevations. "
        "The patient is currently unstable. Provide a detailed diagnosis, coding, and treatment plan."
    )

    diagnostic_output = create_medical_swarm(patient_case)
    print(diagnostic_output)
