# ruff: noqa: E402

"""
MCA SDK Example: Agentic AI Medical Research Assistant

Demonstrates agentic AI-specific SDK features with Registry API integration.
"""

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")


# 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,
    COLLECTOR_HTTP_ENDPOINT,
    REGISTRY_URL,
    AGENTIC_MODEL_ID,
)

from mca_sdk import MCAClient
from mca_sdk.integrations import MCALiteLLMCallback
import litellm
from mca_sdk.registry import RegistryClient

from tools import MockPubMedSearch, MockDrugDatabase, MockCalculator


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

    def __init__(self, mca_client: MCAClient):
        self.client = mca_client
        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):
            with self.client.tracer.start_as_current_span("policy_check"):
                self.client.record_policy_violation(policy_type="content_safety")
            print("⚠️  Policy violation detected: Question contains prohibited terms")
            result = {"status": "blocked", "reason": "policy_violation"}
            self.client.record_prediction(
                input_data={"question": question}, output=result, latency=time.time() - start_time
            )
            return result

        # Start goal tracking
        goal_id = self.client.record_goal_started(
            goal_description=question[:100], goal_type="research_question", complexity="high"
        )
        print(f"Goal started: {goal_id}")

        try:
            # Planning
            with self.client.reasoning_step(
                step_type="planning", step_description="Planning search strategy"
            ) as span:
                span.set_attribute("tokens_used", 150)
                search_terms = ["diabetes type 2 cardiovascular"]
                print(f"  Search terms: {search_terms}")

            # Literature search
            papers = []
            for term in search_terms:
                with self.client.track_tool(
                    "pubmed_search", goal_id=goal_id, search_term=term
                ) as span:
                    results = self.pubmed.search(term)
                    span.set_attribute("papers_found", len(results))
                    papers.extend(results)
                    print(f"  Found {len(results)} papers for '{term}'")

            # Analysis
            with self.client.reasoning_step(
                step_type="evaluation", step_description="Analyzing papers"
            ) as span:
                span.set_attribute("tokens_used", 300)
                drug_classes = ["sglt2_inhibitors", "glp1_agonists"]
                print(f"  Identified: {drug_classes}")

            # Query drug database
            drug_info = {}
            for drug_class in drug_classes:
                with self.client.track_tool(
                    "drug_database", goal_id=goal_id, drug_class=drug_class
                ) as span:
                    info = self.drug_db.query(drug_class)
                    drug_info[drug_class] = info
                    print(f"  Retrieved info for {drug_class}")

            # Calculate dosage
            with self.client.track_tool("calculator", goal_id=goal_id) as span:
                ratio = self.calculator.calculate("10 * 1.5")
                print(f"  Calculated dosage: {ratio}mg")

            # Synthesis
            with self.client.reasoning_step(
                step_type="synthesis", step_description="Creating answer"
            ) as span:
                prompt = f"Summarize the following findings into a clinical answer: Papers analyzed: {len(papers)}, Drugs evaluated: {list(drug_info.keys())}"

                # Real LiteLLM call
                print("  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:
                    print(f"  LiteLLM call failed: {e}")
                    answer = f"Research complete: {len(papers)} papers analyzed, {len(drug_info)} drug classes evaluated"

                span.set_attribute(
                    "tokens_used",
                    getattr(response, "usage", {}).get("total_tokens", 250)
                    if "response" in locals()
                    else 250,
                )

            # Human review
            self.client.record_human_intervention(
                intervention_type="review_request", reason="clinical_decision_support"
            )
            print("  Human review requested")

            # Complete goal with metrics
            self.client.record_goal_completed(
                goal_id,
                status="success",
                tools_used=4,
                iterations=len(search_terms) + len(drug_classes),
                papers_analyzed=len(papers),
            )
            print(f"Goal completed: {goal_id}")

            result = {"goal_id": goal_id, "answer": answer, "status": "success"}
            self.client.record_prediction(
                input_data={"question": question}, output=result, latency=time.time() - start_time
            )
            return result

        except Exception as e:
            self.client.record_goal_completed(goal_id, status="failure")
            self.client.record_error(
                error=e,
                error_type="agent_error",
                severity="high",
                context={"goal_id": goal_id, "question": question[:100]},
            )
            result = {"goal_id": goal_id, "error": str(e), "status": "failure"}
            self.client.record_prediction(
                input_data={"question": question}, output=result, latency=time.time() - start_time
            )
            return result


def main():
    print("MCA SDK - Agentic AI Medical Research Assistant\n")

    # Step 1-2: Verify prerequisites
    if not print_gcp_auth_status() or not print_collector_status():
        sys.exit(1)

    # Step 3-5: Registry integration
    print("\n=== Registry Integration ===")
    model_id = os.environ.get("MODEL_ID", AGENTIC_MODEL_ID)

    registry_client = RegistryClient(url=REGISTRY_URL, use_gcp_auth=True, timeout=10.0)

    # Fetch specific model config
    model_config = registry_client.fetch_model_config(model_id)
    print(f"✓ Using: {model_config.service_name} ({model_config.team_name})\n")

    # Step 6: Initialize SDK

    # Initialize GenAI client
    print("=== SDK Initialization ===")
    client = MCAClient(
        service_name=model_config.service_name,
        model_id=model_id,
        model_type="agentic",
        team_name=model_config.team_name,
        llm_provider="openai",
        llm_model="gpt-4",
        capture_input_data=True,
        capture_output_data=True,
        collector_endpoint=COLLECTOR_HTTP_ENDPOINT,
        allow_insecure_collector=True,
        registry_url=REGISTRY_URL,
        use_gcp_auth=True,
    )

    # Register LiteLLM callback
    callback = MCALiteLLMCallback(client)
    litellm.callbacks = [callback]

    print("✓ Agentic SDK initialized\n")

    try:
        # Step 7: Execute agentic workflow
        print("=== Agentic Workflow ===")
        agent = MedicalResearchAgent(client)
        question = (
            "What are the latest treatments for Type 2 Diabetes with cardiovascular complications?"
        )

        result = agent.research_question(question)

        if result["status"] == "success":
            print(f"\n{result['answer']}\n")
        elif result["status"] == "blocked":
            print(f"\n⚠️  {result['reason']}\n")

        # Step 8: Test policy violation
        print("=== Policy Violation Example ===")
        blocked_question = "experimental_treatment for rare disease"
        agent.research_question(blocked_question)
        print()

        # Step 9: Flush
        print("=== Flush Telemetry ===")
        client.shutdown(timeout_millis=5000)
        registry_client.close()
        print("✓Telemetry sent to collector\n")

        # Step 10: Verification
        print("=== Verify in GCP ===")
        print("Cloud Logging: https://console.cloud.google.com/logs/query")
        print('  Filter: logName="projects/bhsf-ai-team-projects/logs/emms-model-telemetry"')
        print(f'          AND labels."service.name"="{model_config.service_name}"')
        print("\nCloud Trace: https://console.cloud.google.com/traces/list")
        print(f'  Search: service.name="{model_config.service_name}"')
        print("  Look for spans: agent.goal, agent.tool.*, agent.reasoning.*")
        print("\nMetrics: agent.goals_*, agent.tool_calls_*, agent.human_interventions_*")

    except Exception as e:
        print(f"\n✗ Error: {e}")
        client.shutdown()
        registry_client.close()
        raise


if __name__ == "__main__":
    main()
