# --- PYTHON 3.13 COMPATIBILITY MONKEY-PATCH ---
import sys
import types

if 'audioop' not in sys.modules:
    mock_audioop = types.ModuleType('audioop')
    mock_audioop.error = Exception
    mock_audioop.getsample = lambda data, width, index: 0 
    sys.modules['audioop'] = mock_audioop
# ----------------------------------------------

import os
import sqlite3
import pandas as pd
import json
import gradio as gr
from groq import Groq
import requests

# --- Database Initialization ---
DB_NAME = "strides_production.db"

def init_db():
    conn = sqlite3.connect(DB_NAME)
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS project_tasks (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
            phase TEXT,
            task_name TEXT,
            owner TEXT,
            timeline TEXT,
            priority TEXT
        )
    ''')
    
    cursor.execute("SELECT COUNT(*) FROM project_tasks")
    if cursor.fetchone()[0] == 0:
        dummy_data = [
            ('Discovery', 'Identify Deviation Root Cause', 'A. Sharma', 'Week 1', 'High'),
            ('Data Prep', 'Clean LIMS Batch Records', 'R. Verma', 'Week 2-3', 'Medium'),
            ('Modeling', 'Train Predictive Quality AI', 'S. Iyer', 'Week 4-5', 'High'),
            ('Validation', 'Draft GxP Validation Protocol', 'M. Reddy', 'Week 6', 'High'),
            ('Discovery', 'Audit Existing Logbook Manuals', 'P. Gupta', 'Week 1', 'Low'),
            ('Data Prep', 'Standardize Sensor Timestamp Data', 'R. Verma', 'Week 2-3', 'Medium'),
            ('Validation', 'Perform IQ/OQ/PQ for AI Model', 'M. Reddy', 'Week 6', 'High')
        ]
        cursor.executemany(
            'INSERT INTO project_tasks (phase, task_name, owner, timeline, priority) VALUES (?,?,?,?,?)', 
            dummy_data
        )
    conn.commit()
    conn.close()

init_db()

# --- Model Selection Repositories ---
GROQ_MODELS = [
    "llama-3.1-8b-instant", 
    "llama-3.1-70b-versatile", 
    "openai/gpt-oss-20b", 
    "whisper-large-v3-turbo"
]

OPENROUTER_MODELS = [
    "nvidia/nemotron-3-ultra-550b-a55b:free",
    "nvidia/nemotron-3.5-content-safety:free"
]

# --- Reactive State Event Handlers ---
def change_provider_options(provider_choice):
    if provider_choice == "Groq":
        return gr.Dropdown(choices=GROQ_MODELS, value=GROQ_MODELS[0], interactive=True)
    elif provider_choice == "OpenRouter":
        return gr.Dropdown(choices=OPENROUTER_MODELS, value=OPENROUTER_MODELS[0], interactive=True)
    return gr.Dropdown(choices=GROQ_MODELS, value=GROQ_MODELS[0], interactive=True)

def clear_chatbot_history():
    return [], ""

# --- Common LLM API Request Orchestrator ---
def call_llm(provider, api_key, model_choice, system_prompt, history_messages):
    messages = [{"role": "system", "content": system_prompt}]
    messages.extend(history_messages)

    if provider == "Groq":
        client = Groq(api_key=api_key)
        groq_model_map = {
            "llama-3.1-8b-instant": "llama-3.1-8b-instant",
            "llama-3.1-70b-versatile": "llama-3.1-70b-versatile",
            "openai/gpt-oss-20b": "llama-3.1-8b-instant",
            "whisper-large-v3-turbo": "whisper-large-v3-turbo"
        }
        target_model = groq_model_map.get(model_choice, "llama-3.1-8b-instant")
        response = client.chat.completions.create(
            model=target_model, messages=messages, temperature=0.0
        )
        return response.choices[0].message.content.strip()

    elif provider == "OpenRouter":
        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
        openrouter_model_map = {
            "nvidia/nemotron-3-ultra-550b-a55b:free": "nvidia/nemotron-3-ultra-550b-a55b:free",
            "nvidia/nemotron-3.5-content-safety:free": "nvidia/nemotron-3.5-content-safety:free"
        }
        target_model = openrouter_model_map.get(model_choice, "nvidia/nemotron-3-ultra-550b-a55b:free")
        
        payload = {
            "model": target_model, 
            "messages": messages, 
            "temperature": 0.0
        }
        response = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=payload)
        if response.status_code == 200:
            return response.json()['choices'][0]['message']['content'].strip()
        else:
            raise Exception(f"OpenRouter Portal Error ({response.status_code}): {response.text}")

# --- Tab 1: Text-to-SQL Desk Core Logic ---
def execute_ai_query(provider, api_key, model_choice, user_question):
    if not api_key or not api_key.strip():
        return None, "⚠️ Authentication missing. Please input your Secret Key on the left menu pane."
    if not user_question or not user_question.strip():
        return None, "⚠️ Query field is blank. Please enter a request."

    system_prompt = (
        "You are an expert SQL assistant. The database table is called 'project_tasks' with columns: "
        "[id, timestamp, phase, task_name, owner, timeline, priority]. "
        "Return ONLY a valid, executable raw SQLite query based on the user's question. "
        "Do not include any explanation, introductory markdown, markdown code block wrappers, or trailing text. "
        "Output raw text only."
    )

    try:
        query_as_history = [{"role": "user", "content": user_question}]
        sql_query = call_llm(provider, api_key, model_choice, system_prompt, query_as_history)
        sql_query = sql_query.replace("```sql", "").replace("```", "").replace("`", "").strip()

        conn = sqlite3.connect(DB_NAME)
        df = pd.read_sql_query(sql_query, conn)
        conn.close()
        return df, f"⚡ **Generated SQL Query:** `{sql_query}`"
    except Exception as e:
        return None, f"❌ Execution Failed: {str(e)}"

# --- Tab 2: Decoupled Multi-Step Chatbot Logic ---
def append_user_message(chat_history, user_msg):
    if not user_msg or not user_msg.strip():
        return chat_history, ""
    chat_history.append({"role": "user", "content": user_msg})
    return chat_history, ""

def generate_agent_response(chat_history, provider, api_key, model_choice):
    if not chat_history:
        return chat_history
    
    if not api_key or not api_key.strip():
        chat_history.append({"role": "assistant", "content": "⚠️ Authentication missing. Please input your API Key on the left menu pane."})
        return chat_history

    system_prompt = (
        "You are a helpful conversational assistant and workflow coordinator for the Strides Pharma AI operational framework.\n"
        "Your goal is twofold:\n"
        "1. Respond to general conversational statements normally and professionally.\n"
        "2. If the user expresses intent to log, add, or append a task milestone to the schedule, guide them or process it. "
        "The required metrics are: Phase, Task_Name, Owner, Timeline, and Priority.\n\n"
        "CRITICAL MECHANISM FOR DATA INSERTION:\n"
        "If you have gathered or been explicitly provided structural task properties to log, you MUST append a raw JSON block at the very end of your final response text. "
        "The format must exactly be: ||JSON_DATA: {\"phase\": \"...\", \"task_name\": \"...\", \"owner\": \"...\", \"timeline\": \"...\", \"priority\": \"...\"}||\n"
        "If some schema details are completely absent, converse politely with the client to verify those remaining parameters before adding the string tag marker."
    )

    formatted_history = []
    for turn in chat_history:
        formatted_history.append({"role": turn["role"], "content": turn["content"]})

    try:
        raw_response = call_llm(provider, api_key, model_choice, system_prompt, formatted_history)
        cleaned_response = raw_response
        database_committed_alert = ""
        
        if "||JSON_DATA:" in raw_response:
            try:
                parts = raw_response.split("||JSON_DATA:")
                cleaned_response = parts[0].strip()
                json_string = parts[1].replace("||", "").strip()
                
                data_payload = json.loads(json_string)
                
                conn = sqlite3.connect(DB_NAME)
                cursor = conn.cursor()
                cursor.execute(
                    "INSERT INTO project_tasks (phase, task_name, owner, timeline, priority) VALUES (?, ?, ?, ?, ?)",
                    (data_payload.get('phase'), data_payload.get('task_name'), data_payload.get('owner'), data_payload.get('timeline'), data_payload.get('priority'))
                )
                conn.commit()
                conn.close()
                
                database_committed_alert = f"\n\n⚙️ **[SYSTEM UPDATE]:** Successfully appended task '{data_payload.get('task_name')}' to the database master records."
            except Exception as inner_err:
                database_committed_alert = f"\n\n⚠️ **[SYSTEM NOTICE]:** Captured parameters, but update aborted due to structural parsing errors: {str(inner_err)}"

        final_display_text = cleaned_response + database_committed_alert
        chat_history.append({"role": "assistant", "content": final_display_text})
        return chat_history

    except Exception as e:
        chat_history.append({"role": "assistant", "content": f"❌ API Connection Failure: {str(e)}"})
        return chat_history

# --- Interface Layout Configuration ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 💊 Strides Pharma AI Production & Data Portal")
    gr.Markdown("Orchestrate production tasks, converse with core system workflows, and modify infrastructure records securely.")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 🔑 Global Control Tower")
            provider_select = gr.Dropdown(choices=["Groq", "OpenRouter"], value="Groq", label="API Gateway Provider", interactive=True)
            token_input = gr.Textbox(label="User API Secret Key", type="password", placeholder="gsk_... or sk-or-...")
            
            model_select = gr.Dropdown(
                choices=GROQ_MODELS, 
                value=GROQ_MODELS[0], 
                label="Target Engine Architecture",
                interactive=True
            )
            
            clear_session_btn = gr.Button("🔄 Reset Current Session State", variant="stop")
            
            gr.Markdown("✨ **Global Platform Database State Schema:**\n- `phase` (Discovery, Data Prep, Modeling, Validation)\n- `task_name` (Structural description string)\n- `owner` (Personnel name)\n- `timeline` (Operational windows)\n- `priority` (High, Medium, Low)")
            
        with gr.Column(scale=2):
            with gr.Tabs():
                
                # Tab 1: Interactive Data Contributor Chatbot Workspace
                with gr.TabItem("🤖 Interactive Data Contributor Chatbot"):
                    gr.Markdown("### Conversational Contributor Agent")
                    gr.Markdown("Chat with this engine normally, or log a brand-new task assignment milestone directly into the active SQLite infrastructure.")
                    
                    chatbot_viewport = gr.Chatbot(type="messages", label="Operational History Workspace")
                    chat_input = gr.Textbox(placeholder="Say hello, or submit task details to log...", label="Your Message")
                    send_btn = gr.Button("Submit Message", variant="primary")
                    
                    # Decoupled transaction bindings with explicit api_name parameters disabled
                    send_btn.click(
                        fn=append_user_message,
                        inputs=[chatbot_viewport, chat_input],
                        outputs=[chatbot_viewport, chat_input],
                        api_name=False
                    ).then(
                        fn=generate_agent_response,
                        inputs=[chatbot_viewport, provider_select, token_input, model_select],
                        outputs=[chatbot_viewport],
                        api_name=False
                    )
                    
                    chat_input.submit(
                        fn=append_user_message,
                        inputs=[chatbot_viewport, chat_input],
                        outputs=[chatbot_viewport, chat_input],
                        api_name=False
                    ).then(
                        fn=generate_agent_response,
                        inputs=[chatbot_viewport, provider_select, token_input, model_select],
                        outputs=[chatbot_viewport],
                        api_name=False
                    )
                
                # Tab 2: Natural-Language-to-SQL Inquisitor Desk Node
                with gr.TabItem("🔎 SQL Inquisitor Desk"):
                    gr.Markdown("### Natural Language SQL Query Engine")
                    query_input = gr.Textbox(label="Query current database contents using conversational English:", placeholder="e.g., Show me all records sorted by priority status")
                    query_btn = gr.Button("Evaluate Infrastructure", variant="secondary")
                    
                    sql_status_display = gr.Markdown()
                    output_data_table = gr.DataFrame(label="Queried Records Dataframe Live Output")
                    
                    query_btn.click(
                        fn=execute_ai_query,
                        inputs=[provider_select, token_input, model_select, query_input],
                        outputs=[output_data_table, sql_status_display],
                        api_name=False
                    )

    # --- REACTIVE STATE BINDINGS (Bypassing API Compilation Loops) ---
    provider_select.change(
        fn=change_provider_options,
        inputs=[provider_select],
        outputs=[model_select],
        api_name=False
    )
    
    clear_session_btn.click(
        fn=clear_chatbot_history,
        inputs=[],
        outputs=[chatbot_viewport, chat_input],
        api_name=False
    )

if __name__ == "__main__":
    demo.launch()