"""
MCA SDK Example: GenAI/LLM Clinical Documentation Assistant

Demonstrates GenAI-specific SDK features with Registry API integration.
"""

import os
import sys
import time
import random

# 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,
    DEFAULT_MODEL_ID
)

from mca_sdk.integrations import create_genai_client, track_llm_completion, estimate_openai_cost
from mca_sdk.registry import RegistryClient


def summarize_note(client, note_text: str, note_type: str) -> dict:
    """Simulate LLM-based clinical note summarization."""
    start = time.time()
    time.sleep(random.uniform(0.05, 0.15))

    # Simulate token counts and cost
    prompt_tokens = int(len(note_text.split()) * 1.3)
    completion_tokens = random.randint(60, 120)
    model = "gpt-4"
    cost = estimate_openai_cost(model, prompt_tokens, completion_tokens)
    latency = time.time() - start

    # Track LLM completion
    track_llm_completion(
        client=client,
        model=model,
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        latency=latency,
        status="success"
    )

    # Custom metric
    client.record_counter("genai.notes_processed", 1, note_type=note_type, model=model)

    return {
        "summary": f"[Generated summary for {note_type}]",
        "model": model,
        "tokens": prompt_tokens + completion_tokens,
        "cost": cost,
        "latency": latency
    }


def main():
    print("MCA SDK - GenAI Clinical Documentation Assistant\n")

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

    # Registry integration
    print("\n=== Registry Integration ===")
    model_id = os.environ.get("MODEL_ID", DEFAULT_MODEL_ID)

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

    # Fetch config
    model_config = registry_client.fetch_model_config(model_id)
    print(f"✓ Using: {model_config.service_name} ({model_config.team_name})")
    if model_config.thresholds:
        print(f"✓ Cost threshold: ${model_config.thresholds.get('token_cost_per_1k', 0.05):.4f} per 1k tokens")
    print()

    # Initialize GenAI client
    print("=== SDK Initialization ===")
    # NOTE: capture_input_data and capture_output_data default to False for HIPAA compliance.
    # This example explicitly enables them for demonstration purposes.
    # In production, only enable if sensitive data has been sanitized or you have appropriate data use agreements.
    client = create_genai_client(
        service_name=model_config.service_name,
        llm_provider="openai",
        llm_model="gpt-4",
        team_name=model_config.team_name,
        model_id=model_id,
        capture_input_data=True,   # Explicit opt-in (only if sensitive data is sanitized)
        capture_output_data=True,  # Explicit opt-in (only if sensitive data is sanitized)
        collector_endpoint=COLLECTOR_HTTP_ENDPOINT,
        allow_insecure_collector=True,
        registry_url=REGISTRY_URL,
        use_gcp_auth=True,
    )
    print("✓ GenAI client initialized\n")

    try:
        # Step 7: Using @predict decorator for LLM
        print("=== LLM with @predict Decorator ===")

        @client.predict(capture_input=True, capture_output=True)
        def summarize_with_decorator(note_text: str) -> str:
            time.sleep(0.05)
            # In production: return openai.ChatCompletion.create(...)
            return f"[Summary of note: {note_text[:30]}...]"

        summarize_with_decorator("Patient presents with acute symptoms...")
        print("✓ Summary generated with input/output capture\n")

        # Step 8: Clinical note summarization with cost tracking
        print("=== Clinical Note Summarization with Cost Tracking ===")

        notes = [
            ("Patient presents with fever and cough for 3 days...", "progress_note"),
            ("72-year-old male with hypertension...", "discharge_summary"),
            ("Type 2 Diabetes with suboptimal control...", "consultation_note")
        ]

        cost_threshold = client.thresholds.get('token_cost_per_1k', 0.05)
        total_cost = 0.0

        for note_text, note_type in notes:
            result = summarize_note(client, note_text, note_type)
            total_cost += result['cost']

            # Check cost threshold
            cost_per_1k = (result['cost'] / result['tokens']) * 1000
            if cost_per_1k > cost_threshold:
                print(f"  ⚠️  {note_type}: ${result['cost']:.6f} (${cost_per_1k:.4f}/1k tokens - over threshold)")
            else:
                print(f"  ✓ {note_type}: ${result['cost']:.6f}")

        print(f"\nTotal cost: ${total_cost:.6f}\n")

        # Step 9: Custom GenAI metrics
        print("=== Custom GenAI Metrics ===")
        client.record_counter("genai.cache_hits", 5, cache_type="semantic")
        client.record_gauge("genai.rate_limit_remaining", 4500, provider="openai")
        client.record_metric("genai.summary_quality", 4.2, metric_type="histogram", model="gpt-4")
        print("✓ Recorded custom GenAI metrics\n")

        # Step 10: LLM error tracking
        print("=== LLM Error Tracking ===")
        try:
            raise Exception("OpenAI API rate limit exceeded (429)")
        except Exception as e:
            track_llm_completion(client, model="gpt-4", prompt_tokens=0,
                               completion_tokens=0, latency=0.5, status="error")
            client.record_error(
                error=e,
                error_type="llm_api_error",
                severity="high",
                context={"provider": "openai", "model": "gpt-4", "error_code": "429"}
            )
            print(f"✓ Error tracked: {type(e).__name__}\n")

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

        # Step 12: 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("\nMetrics: genai.requests_total, genai.tokens_total, genai.cost_usd")
        print("  Filter by: model=gpt-4, note_type=progress_note")

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


if __name__ == "__main__":
    main()
