# import requests
import msal
from datetime import datetime, timedelta
import streamlit as st
import os
from azure.ai.projects import AIProjectClient
from azure.ai.agents import AgentsClient
from azure.ai.agents.models import (
    MessageRole,
    ListSortOrder,
    AgentsNamedToolChoice
)
from azure.identity import OnBehalfOfCredential
# from azure.core.exceptions import ClientAuthenticationError

# ======================================================
# CONFIGURATION - TWO APP REGISTRATIONS
# ======================================================

# App Registration 1: Used for Easy Auth (user login)
AUTH_APP_CLIENT_ID = "<AuthAppClientID>"  # The one configured in App Service Authentication

# App Registration 2: Used for OBO credential (SDK access)
SDK_APP_TENANT_ID = "<AzureTenantID>"
SDK_APP_CLIENT_ID = "<SDK_App_Client_ID>"  # Different from AUTH_APP
SDK_APP_CLIENT_SECRET = "<SDK_App_client_secret>"

# CONFIG: AGENT 1 (TRP Job Agent)
JOB_PROJECT_ENDPOINT = "<Project1_Endpoint>"
JOB_AGENT_ID = "<Project1_Agent_ID>"

# CONFIG: AGENT 2 (TRP Guide Bot)
GUIDE_ENDPOINT = "<Project2_Endpoint>"
GUIDE_AGENT_ID = "<Project2_Agent_ID>"

def decode_token(token):
    """
    Decode JWT token to inspect claims and expiration.
    """
    import base64
    import json
    try:
        parts = token.split('.')
        if len(parts) == 3:
            payload = parts[1]
            payload += '=' * (4 - len(payload) % 4)
            decoded = base64.urlsafe_b64decode(payload)
            claims = json.loads(decoded)
            return claims
    except Exception as e:
        print(f"Could not decode token: {e}")
        return None


def is_token_expired(token, buffer_seconds=300):
    """
    Check if token is expired or will expire soon.
    
    Args:
        token: JWT token string
        buffer_seconds: Consider token expired if it expires within this many seconds (default: 5 minutes)
    
    Returns:
        True if token is expired or will expire soon, False otherwise
    """
    claims = decode_token(token)
    if not claims or 'exp' not in claims:
        return True
    
    # Token expiration timestamp
    exp_timestamp = claims['exp']
    exp_datetime = datetime.fromtimestamp(exp_timestamp)
    
    # Current time plus buffer
    now_with_buffer = datetime.now() + timedelta(seconds=buffer_seconds)
    
    is_expired = now_with_buffer >= exp_datetime
    
    if is_expired:
        print(f"Token expired or expiring soon. Expires at: {exp_datetime}, Current time: {datetime.now()}")
    else:
        time_remaining = exp_datetime - datetime.now()
        print(f"Token valid. Time remaining: {time_remaining}")
    
    return is_expired

# ======================================================
# ON-BEHALF-OF AUTHENTICATION (MINIMAL ADDITION)
# ======================================================

def get_user_access_token():
    """
    Retrieve user access token from Easy Auth headers.
    This token is issued for the AUTH app registration.
    Includes expiration checking and NO caching (always fresh from headers).
    """
    try:
        # Get headers from Streamlit context
        headers = st.context.headers
        
        # Easy Auth exposes the token in this header
        access_token = headers.get("X-Ms-Token-Aad-Access-Token")
        
        if not access_token:
            # Try alternative header names
            access_token = (
                headers.get("x-ms-token-aad-access-token") or
                headers.get("HTTP_X_MS_TOKEN_AAD_ACCESS_TOKEN")
            )
        
        if access_token:
            print("✓ Access token retrieved from Easy Auth headers")
            
            # Decode and validate token
            import base64
            import json
            try:
                parts = access_token.split('.')
                if len(parts) == 3:
                    payload = parts[1]
                    payload += '=' * (4 - len(payload) % 4)
                    decoded = base64.urlsafe_b64decode(payload)
                    claims = json.loads(decoded)
                    
                    print(f"DEBUG: Token audience: {claims.get('aud', 'unknown')}")
                    print(f"DEBUG: Token app ID: {claims.get('appid', 'unknown')}")
                    print(f"DEBUG: Token scopes: {claims.get('scp', 'unknown')}")
                    
                    # Check expiration
                    if is_token_expired(access_token):
                        raise Exception("User access token is expired. Please refresh the page to re-authenticate.")
                        
            except Exception as e:
                print(f"DEBUG: Could not decode/validate token: {e}")
                if "expired" in str(e).lower():
                    raise
            
            return access_token
        else:
            print("DEBUG: Available headers:")
            for key, value in headers.items():
                display_value = value[:40] + "..." if len(value) > 40 else value
                print(f"  {key}: {display_value}")
            
            raise Exception("X-Ms-Token-Aad-Access-Token header not found")
            
    except AttributeError:
        raise Exception(
            "Could not access st.context.headers. "
            "Ensure we're using Streamlit version that supports st.context"
        )
    except Exception as e:
        raise Exception(f"Error retrieving access token: {str(e)}")
    

def get_obo_token_with_msal():
    """
    Use MSAL to acquire token on behalf of user.
    This is the recommended approach per Microsoft documentation.
    
    Returns:
        Access token string (not a credential object)
    """
    try:
        # Get the user's access token (from Easy Auth)
        user_assertion = get_user_access_token()
        
        # Define scopes - adjust based on our needs
        # For Azure AI Projects, Agents and TRP, we typically need:
        # scopes = ["https://cognitiveservices.azure.com/.default"]
        # scopes = [f"api://<sdkApp_Client_ID>/user_impersonation"]
        scopes = ["https://ai.azure.com/.default"]

        
        # Create MSAL confidential client application
        authority = f"https://login.microsoftonline.com/{SDK_APP_TENANT_ID}"
        app = msal.ConfidentialClientApplication(
            client_id=SDK_APP_CLIENT_ID,
            client_credential=SDK_APP_CLIENT_SECRET,
            authority=authority
        )
        
        # Acquire token on behalf of user
        result = app.acquire_token_on_behalf_of(
            user_assertion=user_assertion,
            scopes=scopes
        )
        
        if "access_token" in result:
            print("✓ OBO token acquired successfully via MSAL")
            
            # Decode new token for debugging
            claims = decode_token(result["access_token"])
            if claims:
                print(f"New token audience: {claims.get('aud', 'unknown')}")
                print(f"New token scopes: {claims.get('scp', 'unknown')}")
                print(f"New token issuer: {claims.get('iss', 'unknown')}")
                print(f"New token app ID: {claims.get('appid', 'unknown')}")
                
                # CRITICAL DEBUG: Check if audience matches expected
                expected_aud = "https://ai.azure.com"
                actual_aud = claims.get('aud', '')
                if expected_aud not in actual_aud and actual_aud not in expected_aud:
                    print(f"⚠️ WARNING: Token audience '{actual_aud}' may not match expected '{expected_aud}'")
            
            return result["access_token"]
        else:
            error_msg = result.get("error", "Unknown error")
            error_description = result.get("error_description", "No description")
            raise Exception(f"MSAL OBO failed: {error_msg} - {error_description}")
            
    except Exception as e:
        print(f"❌ Error in get_obo_token_with_msal: {e}")
        raise

def get_cached_obo_token():
    """
    Get OBO token with intelligent caching.
    Only use cached token if it exists and is not expired.
    
    Returns:
        Valid access token string
    """
    # Check if we have a cached token
    if "obo_token" in st.session_state:
        cached_token = st.session_state.obo_token
        
        # Check if cached token is still valid
        if not is_token_expired(cached_token):
            print("✓ Using cached OBO token")
            return cached_token
        else:
            print("⚠ Cached OBO token expired, acquiring new token")
            # Remove expired token from cache
            del st.session_state.obo_token
    
    # Acquire new token
    try:
        new_token = get_obo_token_with_msal()
        st.session_state.obo_token = new_token
        return new_token
    except Exception as e:
        print(f"❌ Failed to acquire OBO token: {e}")
        raise

def get_cached_obo_credential():
    """
    Create a credential object compatible with Azure SDK.
    This wraps the MSAL token for use with AIProjectClient and AgentsClient.
    """
    from azure.core.credentials import AccessToken
    
    class OBOTokenCredential:
        """Custom credential that uses MSAL OBO token."""
        
        def get_token(self, *scopes, **kwargs):
            """Get token for Azure SDK."""
            try:
                token = get_cached_obo_token()
                claims = decode_token(token)
                
                # Get expiration time
                expires_on = claims.get('exp', 0) if claims else 0
                
                return AccessToken(token=token, expires_on=expires_on)
            except Exception as e:
                raise Exception(f"Failed to get OBO token: {e}")
    
    return OBOTokenCredential()


def get_job_client():
    cred = get_cached_obo_credential()
    return AIProjectClient(
        credential=cred,
        endpoint=JOB_PROJECT_ENDPOINT
    )

def get_guide_client():
    cred = get_cached_obo_credential()
    return AgentsClient(
        endpoint=GUIDE_ENDPOINT,
        credential=cred
    )

def debug_token_details():
    """
    Detailed token analysis to identify the issue
    """
    try:
        headers = st.context.headers
        access_token = headers.get("X-Ms-Token-Aad-Access-Token") or headers.get("x-ms-token-aad-access-token")
        
        if not access_token:
            st.error("No token found in headers")
            return
        
        # Decode token
        import base64
        import json
        
        parts = access_token.split('.')
        if len(parts) != 3:
            st.error(f"Invalid token format - has {len(parts)} parts, expected 3")
            return
        
        # Decode header
        header = parts[0]
        header += '=' * (4 - len(header) % 4)
        header_decoded = base64.urlsafe_b64decode(header)
        header_claims = json.loads(header_decoded)
        
        # Decode payload
        payload = parts[1]
        payload += '=' * (4 - len(payload) % 4)
        payload_decoded = base64.urlsafe_b64decode(payload)
        payload_claims = json.loads(payload_decoded)
        
        st.write("### Token Header")
        st.json(header_claims)
        
        st.write("### Token Payload (Claims)")
        st.json(payload_claims)
        
        st.write("### Key Findings:")
        st.write(f"**Token Type (typ):** {header_claims.get('typ', 'Not specified')}")
        st.write(f"**Algorithm (alg):** {header_claims.get('alg', 'Not specified')}")
        st.write(f"**Audience (aud):** {payload_claims.get('aud', 'Not specified')}")
        st.write(f"**Issuer (iss):** {payload_claims.get('iss', 'Not specified')}")
        st.write(f"**App ID (appid):** {payload_claims.get('appid', 'Not specified')}")
        st.write(f"**Token Version (ver):** {payload_claims.get('ver', 'Not specified')}")
        st.write(f"**Scopes (scp):** {payload_claims.get('scp', 'Not specified')}")
        st.write(f"**Roles (roles):** {payload_claims.get('roles', 'Not specified')}")
        
        # Check for ID token indicators
        if payload_claims.get('idtyp') or 'oid' in payload_claims:
            st.warning("⚠️ This might be an ID token, not an access token!")
        
        # Check audience
        expected_audiences = [
            AUTH_APP_CLIENT_ID,
            f"api://{AUTH_APP_CLIENT_ID}",
            SDK_APP_CLIENT_ID,
            f"api://{SDK_APP_CLIENT_ID}"
        ]
        
        actual_aud = payload_claims.get('aud', '')
        if actual_aud not in expected_audiences:
            st.error(f"❌ Token audience '{actual_aud}' doesn't match expected values")
            st.write("Expected one of:")
            for aud in expected_audiences:
                st.write(f"  - {aud}")
        else:
            st.success(f"✓ Token audience matches: {actual_aud}")
            
    except Exception as e:
        st.error(f"Error decoding token: {str(e)}")
        import traceback
        st.code(traceback.format_exc())


# ==============================================
# TRP JOB AGENT LOGIC
# ==============================================

def run_job_agent(thread,user_input: str):
    try:
        cred = get_cached_obo_credential()

        job_client = AIProjectClient(
            credential=cred,
            endpoint=JOB_PROJECT_ENDPOINT
        )

        job_agent = job_client.agents.get_agent(JOB_AGENT_ID)
        
        # Send message
        job_client.agents.messages.create(
            # thread_id=job_thread.id,
            thread_id=thread.id,
            # role="user",
            role=MessageRole.USER,
            content=user_input
        )

        # Run agent
        job_client.agents.runs.create_and_process(
            # thread_id=job_thread.id,
            thread_id=thread.id,
            agent_id=job_agent.id
        )

        # Fetch response
        messages = job_client.agents.messages.list(
            # thread_id=job_thread.id,
            thread_id=thread.id,
            order=ListSortOrder.ASCENDING
        )

        last = None
        for m in messages:
            if m.role == "assistant" and m.text_messages:
                last = m.text_messages[-1].text.value

        return last or "No response."

    except Exception as e:
        return f"⚠️ Error: {str(e)}"


# ==============================================
# TRP GUIDE BOT LOGIC
# ==============================================
def run_guide_agent(thread, user_input: str):
    try:
        cred = get_cached_obo_credential()

        guide_client = AgentsClient(
            endpoint=GUIDE_ENDPOINT,
            credential=cred
        )

        guide_client.messages.create(
            thread_id=thread.id,
            role=MessageRole.USER,
            content=user_input
        )

        guide_client.runs.create_and_process(
            thread_id=thread.id,
            agent_id=GUIDE_AGENT_ID,
            tool_choice=AgentsNamedToolChoice(type="file_search")
        )

        messages = guide_client.messages.list(
            thread_id=thread.id,
            order=ListSortOrder.ASCENDING
        )

        last = None
        for m in messages:
            if m.role == "assistant" and m.text_messages:
                last = m.text_messages[-1].text.value

        return last or "No response."

    except Exception as e:
        return f"⚠️ Error: {str(e)}"


# ==============================================
# STREAMLIT UI
# ==============================================
st.set_page_config(page_title="TRP Bots", page_icon="🤖", layout="wide")

# DEBUG Section
# ==============================================
# Add this button to our UI temporarily
if st.sidebar.button("🔍 Debug Token"):
    debug_token_details()

if st.sidebar.button("🔍 Debug: Show Headers & Environment"):
    st.write("### Headers from st.context")
    try:
        headers_dict = dict(st.context.headers)
        st.json(headers_dict)
    except Exception as e:
        st.error(f"Could not access headers: {e}")
    
    st.write("### Environment Variables")
    env_dict = {}
    jwt_vars = []
    
    for key in sorted(os.environ.keys()):
        value = os.environ[key]
        
        # Check if it's a JWT token
        if isinstance(value, str) and value.startswith("eyJ"):
            jwt_vars.append(key)
            env_dict[key] = f"JWT_TOKEN (length: {len(value)}, starts: {value[:30]}...)"
        elif len(value) > 50:
            env_dict[key] = value[:50] + f"... (length: {len(value)})"
        else:
            env_dict[key] = value
    
    st.json(env_dict)
    
    if jwt_vars:
        st.warning(f"🎯 Found potential JWT tokens in: {', '.join(jwt_vars)}")

# ---------------------------
# GLOBAL THEME + LAweT STYLE
# ---------------------------

custom_css = """
<style>

    /* ------------------------------
       PAGE BACKGROUND (CENTER AREA)
       ------------------------------ */
    [data-testid="stAppViewContainer"] {
        background-color: #dedede !important;
        display: flex;
        justify-content: center;
    }

    /* Constrain main content to center */
    .main-centered {
        max-width: 820px;
        margin-left: auto;
        margin-right: auto;
        padding-top: 20px;
        text-align: center;
    }

    /* Center all headers inside main container */
    .main-centered h1, 
    .main-centered h2, 
    .main-centered h3, 
    .main-centered .subheader {
        text-align: center !important;
    }

    /* ------------------------------
       SIDEBAR STYLE
       ------------------------------ */
    [data-testid="stSidebar"] {
        background-color: #44c9d0 !important;
        padding: 20px;
    }

    /* Bigger sidebar title */
    .sidebar-title {
        font-size: 1.3rem;
        font-weight: bold;
        padding-bottom: 10px;
    }

    /* ------------------------------
       CHAT MESSAGE STYLE
       ------------------------------ */
    .stChatMessage {
        display: flex;
        justify-content: center !important;
    }

    /* ------------------------------
       CHAT INPUT FIXED AT BOTTOM
       ------------------------------ */
    .stChatInputContainer {
        position: fixed !important;
        bottom: 0 !important;
        left: 0;
        right: 0;
        padding: 10px 25%;
        background-color: #f4f7fb !important;
        z-index: 9999;
    }

    /* Rounded input */
    .stChatInputContainer textarea {
        border-radius: 20px !important;
        padding: 10px 15px !important;
    }

</style>
"""
st.markdown(custom_css, unsafe_allow_html=True)

# ---------------------------
# SIDEBAR AGENT SELECTOR
# ---------------------------
st.sidebar.markdown("<div class='sidebar-title'>Choose our Assistant</div>", unsafe_allow_html=True)

agent_choice = st.sidebar.radio(
    "Select Bot",
    ["📘 TRP Guide Bot", "🔁 TRP Job Agent"],
    label_visibility="collapsed"
)

# Wrapper div for center-aligned content
st.markdown("<div class='main-centered'>", unsafe_allow_html=True)

# st.title("🤖 Unified TRP AI Assistant")
# st.markdown("<h1 class='subheader'>🤖 Unified TRP AI Assistant</h1>", unsafe_allow_html=True)
st.markdown("<h1 style='text-align:center; width:100%;'>🤖 Unified TRP AI Assistant</h1>", unsafe_allow_html=True)

if agent_choice == "📘 TRP Guide Bot":
    # show Guide Bot UI
    # ==============================================
    # TAB 1 → TRP GUIDE BOT
    # ==============================================
    # with tab1:
        # st.subheader("🧭 TRP Guide Bot")
        # st.markdown("<h2 class='subheader'>📘 TRP Guide Bot</h2>", unsafe_allow_html=True)
        st.markdown("<h2 style='text-align:center; width:100%;'>📘 TRP Guide Bot</h2>", unsafe_allow_html=True)

        # st.write("Ask questions about TRP metadata, configuration and pipeline guidance.")
        # st.markdown("<h3 style='text-align:center;'>Ask questions about TRP metadata, configuration and pipeline guidance.</h3>", unsafe_allow_html=True)
        st.markdown(
            "<p style='text-align:center; width:100%; font-size:1.2rem;'>"
            "Ask questions about TRP metadata, configuration and modules."
            "</p>",
            unsafe_allow_html=True
        )

        # Create thread if not exists
        if "guide_thread" not in st.session_state:
            # st.session_state.guide_thread = guide_client.threads.create()
            st.session_state.guide_thread = get_guide_client().threads.create()

        # Init chat history
        if "guide_msgs" not in st.session_state:
            st.session_state.guide_msgs = []

        # Display messages
        for msg in st.session_state.guide_msgs:
            with st.chat_message(msg["role"]):
                st.markdown(msg["text"])

        user_input = st.chat_input("Ask Guide Bot...")
        if user_input:
            st.session_state.guide_msgs.append({"role": "user", "text": user_input})
            with st.chat_message("user"):
                st.markdown(user_input)

            result = run_guide_agent(st.session_state.guide_thread, user_input)

            st.session_state.guide_msgs.append({"role": "assistant", "text": result})
            with st.chat_message("assistant"):
                st.markdown(result)

else:
    # ==============================================
    # TAB 2 → TRP JOB AGENT
    # ==============================================

        # st.subheader("🔧 TRP Job Agent")
        # st.markdown("<h2 class='subheader'>🔁 TRP Job Agent</h2>", unsafe_allow_html=True)
        st.markdown("<h2 style='text-align:center; width:100%;'>🔁 TRP Job Agent</h2>", unsafe_allow_html=True)

        # st.write("Ask questions about TRP TRP Pipelines execution, logs and job metadata.")
        # st.markdown("<h3 style='text-align:center;'>Ask questions about TRP TRP Pipelines execution, logs and job metadata.</h3>", unsafe_allow_html=True)
        st.markdown(
            "<p style='text-align:center; width:100%; font-size:1.2rem;'>"
            "Ask questions about TRP TRP Pipelines execution, logs and job metadata."
            "</p>",
            unsafe_allow_html=True
        )

        # Create thread if not exists
        if "job_thread" not in st.session_state:
            # st.session_state.job_thread = job_client.agents.threads.create()
            st.session_state.job_thread = get_job_client().agents.threads.create()


        # Init chat history
        if "job_msgs" not in st.session_state:
            st.session_state.job_msgs = []

        # Display messages
        for msg in st.session_state.job_msgs:
            with st.chat_message(msg["role"]):
                st.markdown(msg["text"])

        user_input = st.chat_input("Ask Job Agent...")
        if user_input:
            st.session_state.job_msgs.append({"role": "user", "text": user_input})
            with st.chat_message("user"):
                st.markdown(user_input)

            result = run_job_agent(st.session_state.job_thread, user_input)

            st.session_state.job_msgs.append({"role": "assistant", "text": result})
            with st.chat_message("assistant"):
                st.markdown(result)

# Close center wrapper
st.markdown("</div>", unsafe_allow_html=True)