# ruff: noqa: E402

"""
Medical Research Assistant - Non-instrumented Version

A baseline implementation without MCA SDK instrumentation for comparison.
"""

import logging
import os
import sys
import time
import warnings
from typing import Dict

# Suppress expected Pydantic warnings from Litellm mock responses
warnings.filterwarnings("ignore", category=UserWarning, module="pydantic")

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


# Add both sdk-examples (for shared.py) and mca-prototype root (for mca_sdk) to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))

from shared import (
    print_gcp_auth_status,
    print_collector_status,
)

import litellm

from tools import MockPubMedSearch, MockDrugDatabase, MockCalculator


class MedicalResearchAgent:
    """Agentic AI that researches clinical questions using multiple tools."""

    def __init__(self):
        self.pubmed = MockPubMedSearch()
        self.drug_db = MockDrugDatabase()
        self.calculator = MockCalculator()

    def research_question(self, question: str) -> Dict:
        """Research a clinical question using multi-step reasoning and tools."""
        start_time = time.time()

        # Check for policy violations (content safety)
        prohibited_terms = ["experimental_treatment", "off_label_use"]
        if any(term in question.lower() for term in prohibited_terms):
            logger.warning("Policy violation detected: Question contains prohibited terms")
            result = {"status": "blocked", "reason": "policy_violation"}
            return result

        # Generate goal ID
        goal_id = f"goal-{int(time.time() * 1000)}"
        logger.info(f"Goal started: {goal_id}")

        try:
            # Planning
            logger.info("Planning search strategy...")
            search_terms = ["diabetes type 2 cardiovascular"]
            logger.info(f"Search terms: {search_terms}")

            # Literature search
            papers = []
            for term in search_terms:
                logger.info(f"Searching PubMed for '{term}'...")
                results = self.pubmed.search(term)
                papers.extend(results)
                logger.info(f"Found {len(results)} papers for '{term}'")

            # Analysis
            logger.info("Analyzing papers...")
            drug_classes = ["sglt2_inhibitors", "glp1_agonists"]
            logger.info(f"Identified: {drug_classes}")

            # Query drug database
            drug_info = {}
            for drug_class in drug_classes:
                logger.info(f"Querying drug database for {drug_class}...")
                info = self.drug_db.query(drug_class)
                drug_info[drug_class] = info
                logger.info(f"Retrieved info for {drug_class}")

            # Calculate dosage
            logger.info("Calculating dosage...")
            ratio = self.calculator.calculate("10 * 1.5")
            logger.info(f"Calculated dosage: {ratio}mg")

            # Synthesis
            logger.info("Creating answer...")
            prompt = f"Summarize the following findings into a clinical answer: Papers analyzed: {len(papers)}, Drugs evaluated: {list(drug_info.keys())}"

            # Real LiteLLM call
            logger.info("Calling litellm with prompt...")
            try:
                response = litellm.completion(
                    model="gpt-3.5-turbo",
                    messages=[{"role": "user", "content": prompt}],
                    mock_response="Mocked clinical answer based on search terms."
                    if not os.environ.get("OPENAI_API_KEY")
                    else None,
                )
                answer = response.choices[0].message.content
            except Exception as e:
                logger.error(f"LiteLLM call failed: {e}")
                answer = f"Research complete: {len(papers)} papers analyzed, {len(drug_info)} drug classes evaluated"

            # Human review
            logger.info("Human review requested")

            # Complete goal
            latency = time.time() - start_time
            logger.info(f"Goal completed: {goal_id} (latency: {latency:.3f}s)")

            result = {"goal_id": goal_id, "answer": answer, "status": "success"}
            return result

        except Exception as e:
            logger.error(f"Goal failed: {goal_id} - {str(e)}")
            result = {"goal_id": goal_id, "error": str(e), "status": "failure"}
            return result


def main():
    logger.info("Medical Research Assistant - Non-instrumented Version")

    # Step 1-2: Verify prerequisites
    if not print_gcp_auth_status() or not print_collector_status():
        logger.warning("Prerequisites check failed (but continuing anyway for baseline)")

    logger.info("=== Agent Initialization ===")
    agent = MedicalResearchAgent()
    logger.info("Agent initialized (no instrumentation)")

    try:
        # Execute agentic workflow
        logger.info("=== Agentic Workflow ===")
        question = (
            "What are the latest treatments for Type 2 Diabetes with cardiovascular complications?"
        )

        result = agent.research_question(question)

        if result["status"] == "success":
            logger.info(f"Answer: {result['answer']}")
        elif result["status"] == "blocked":
            logger.warning(f"Blocked: {result['reason']}")

        # Test policy violation
        logger.info("=== Policy Violation Example ===")
        blocked_question = "experimental_treatment for rare disease"
        agent.research_question(blocked_question)

        logger.info("=== Execution Complete ===")
        logger.info("✓ Telemetry sent to collector (auto-instrumented version)\n")
        logger.info("=== Verify in GCP ===")
        logger.info("Cloud Logging: https://console.cloud.google.com/logs/query")
        logger.info('  Filter: logName="projects/bhsf-ai-team-projects/logs/emms-model-telemetry"')
        logger.info("\nCloud Trace: https://console.cloud.google.com/traces/list")
        logger.info("  Look for spans: agent.goal, agent.tool.*, agent.reasoning.*")
        logger.info("\nMetrics: agent.goals_*, agent.tool_calls_*, agent.human_interventions_*")

    except Exception as e:
        logger.error(f"Error: {e}")
        raise


if __name__ == "__main__":
    main()
