nxtgauge-backend-rust/apps/users/src/handlers/ai.rs

use crate::AppState;
use axum::{
    extract::State,
    http::StatusCode,
    response::IntoResponse,
    routing::{get, post},
    Json, Router,
};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

pub fn ai_router() -> Router<AppState> {
    Router::new()
        .route("/chat/message", post(ai_chat_message))
        .route("/tickets/create", post(ai_create_ticket))
        .route("/tickets/{id}", get(ai_get_ticket))
        .route("/forms/extract", post(ai_extract_form))
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct OllamaChatRequest {
    pub model: Option<String>,
    pub message: String,
    pub conversation_id: Option<String>,
    pub user_id: Option<String>,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct OllamaChatResponse {
    pub message: String,
    pub conversation_id: String,
    pub intent: String,
    pub confidence: f32,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
struct OllamaGenerateRequest {
    model: String,
    prompt: String,
    stream: bool,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
struct OllamaGenerateResponse {
    response: String,
}

async fn call_ollama(_state: &AppState, model: &str, prompt: &str) -> Result<String, String> {
    let base_url = std::env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| "http://ollama.nxtgauge-ai.svc.cluster.local:11434".to_string());
    let url = format!("{}/api/generate", base_url);

    let req = OllamaGenerateRequest {
        model: model.to_string(),
        prompt: prompt.to_string(),
        stream: false,
    };

    let client = reqwest::Client::new();
    let response = client
        .post(&url)
        .json(&req)
        .send()
        .await
        .map_err(|e| format!("ollama request failed: {}", e))?;

    if !response.status().is_success() {
        return Err(format!("ollama returned status: {}", response.status()));
    }

    let result: OllamaGenerateResponse = response
        .json()
        .await
        .map_err(|e| format!("failed to parse ollama response: {}", e))?;

    Ok(result.response)
}

async fn classify_intent(message: &str, ollama_base: &str, model: &str) -> (String, f32) {
    let prompt = format!(
        "Classify this user message into one intent category. Categories: ticket_creation, form_filling, help_search, general. \
         Return ONLY the intent name, nothing else.\n\nMessage: {}",
        message
    );

    match call_ollama_inline(ollama_base, model, &prompt).await {
        Ok(response) => {
            let intent = response.trim().to_lowercase();
            let confidence = if intent.is_empty() { 0.5 } else { 0.85 };
            let intent = match intent.as_str() {
                "ticket_creation" => "ticket_creation",
                "form_filling" => "form_filling",
                "help_search" => "help_search",
                _ => "general",
            };
            (intent.to_string(), confidence)
        }
        Err(_) => ("general".to_string(), 0.5),
    }
}

async fn call_ollama_inline(base_url: &str, model: &str, prompt: &str) -> Result<String, String> {
    let url = format!("{}/api/generate", base_url);
    let req = OllamaGenerateRequest {
        model: model.to_string(),
        prompt: prompt.to_string(),
        stream: false,
    };

    let client = reqwest::Client::new();
    let response = client
        .post(&url)
        .json(&req)
        .send()
        .await
        .map_err(|e| format!("ollama request failed: {}", e))?;

    if !response.status().is_success() {
        return Err(format!("ollama returned status: {}", response.status()));
    }

    let result: OllamaGenerateResponse = response
        .json()
        .await
        .map_err(|e| format!("failed to parse ollama response: {}", e))?;

    Ok(result.response)
}

async fn ai_chat_message(
    State(state): State<AppState>,
    Json(body): Json<OllamaChatRequest>,
) -> impl IntoResponse {
    let ollama_base = std::env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| "http://ollama.nxtgauge-ai.svc.cluster.local:11434".to_string());
    let model = std::env::var("OLLAMA_CHAT_MODEL").unwrap_or_else(|_| "gemma3:270m".to_string());
    let default_conversation = Uuid::new_v4().to_string();

    let conversation_id = body.conversation_id.unwrap_or_else(|| default_conversation);

    let (intent, confidence) = classify_intent(&body.message, &ollama_base, &model).await;

    let system_prompt = match intent.as_str() {
        "ticket_creation" => {
            "You are a support ticket assistant. Help users create clear, actionable support tickets. \
             Ask for: subject, description of issue, category, priority if not provided. \
             Summarize the ticket in a structured way."
        }
        "form_filling" => {
            "You are a form filling assistant. Help users fill out forms by extracting relevant information \
             from their message. Extract key:value pairs when possible."
        }
        "help_search" => {
            "You are a help center assistant. Help users find relevant help articles based on their query. \
             Ask clarifying questions to narrow down the search."
        }
        _ => {
            "You are a helpful AI assistant for Nxtgauge platform. Provide clear, concise responses. \
             If the user needs support, guide them to create a ticket."
        }
    };

    let full_prompt = format!("{}\n\nUser: {}\nAssistant:", system_prompt, body.message);

    let response_text = match call_ollama(&state, &model, &full_prompt).await {
        Ok(r) => r,
        Err(e) => {
            tracing::error!("Ollama error: {}", e);
            "I'm having trouble processing your request right now. Please try again or contact support.".to_string()
        }
    };

    (
        StatusCode::OK,
        Json(OllamaChatResponse {
            message: response_text,
            conversation_id,
            intent,
            confidence,
        }),
    )
        .into_response()
}

async fn ai_create_ticket(
    State(state): State<AppState>,
    Json(body): Json<serde_json::Value>,
) -> impl IntoResponse {
    let subject = body.get("subject").and_then(|v| v.as_str()).unwrap_or("AI Assisted Request");
    let description = body.get("description").and_then(|v| v.as_str());
    let category = body.get("category").and_then(|v| v.as_str()).unwrap_or("ai_assisted");
    let priority = body.get("priority").and_then(|v| v.as_str()).unwrap_or("medium");
    let user_id = body.get("user_id").and_then(|v| v.as_str())
        .and_then(|s| Uuid::parse_str(s).ok())
        .unwrap_or_else(Uuid::nil);

    let result = sqlx::query_as::<_, TicketRow>(
        r#"
        INSERT INTO support_tickets (user_id, subject, description, category, priority, status)
        VALUES ($1, $2, $3, $4, $5, 'new')
        RETURNING id, subject, description, category, priority, status,
                  requester_name, requester_email, assigned_to, created_at, updated_at
        "#,
    )
    .bind(user_id)
    .bind(subject)
    .bind(description)
    .bind(category)
    .bind(priority)
    .fetch_one(&state.pool)
    .await;

    match result {
        Ok(r) => (
            StatusCode::CREATED,
            Json(serde_json::json!({
                "id": r.id,
                "subject": r.subject,
                "status": r.status,
                "ticket_id": r.id,
            })),
        )
            .into_response(),
        Err(e) => {
            tracing::error!("AI ticket creation failed: {}", e);
            (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Failed to create ticket" }))).into_response()
        }
    }
}

async fn ai_get_ticket(
    State(state): State<AppState>,
    axum::extract::Path(id): axum::extract::Path<Uuid>,
) -> impl IntoResponse {
    let result = sqlx::query_as::<_, TicketRow>(
        r#"
        SELECT id, subject, description, category, priority, status,
               requester_name, requester_email, assigned_to, created_at, updated_at
        FROM support_tickets WHERE id = $1
        "#,
    )
    .bind(id)
    .fetch_optional(&state.pool)
    .await;

    match result {
        Ok(Some(r)) => (
            StatusCode::OK,
            Json(serde_json::json!({
                "id": r.id,
                "subject": r.subject,
                "description": r.description,
                "category": r.category,
                "priority": r.priority,
                "status": r.status,
                "requester_name": r.requester_name,
                "requester_email": r.requester_email,
                "assigned_to": r.assigned_to,
                "created_at": r.created_at,
                "updated_at": r.updated_at,
            })),
        )
            .into_response(),
        Ok(None) => (StatusCode::NOT_FOUND, Json(serde_json::json!({ "error": "Ticket not found" }))).into_response(),
        Err(e) => {
            tracing::error!("Failed to fetch ticket {}: {}", id, e);
            (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Failed to fetch ticket" }))).into_response()
        }
    }
}

#[derive(Debug, Deserialize)]
struct FormExtractBody {
    message: String,
    form_type: Option<String>,
}

#[derive(Debug, Serialize)]
struct FormExtractResponse {
    fields: Vec<ExtractedField>,
    missing_fields: Vec<String>,
    confidence: f32,
}

#[derive(Debug, Serialize)]
struct ExtractedField {
    key: String,
    value: String,
    confidence: f32,
}

async fn ai_extract_form(
    State(_state): State<AppState>,
    Json(body): Json<FormExtractBody>,
) -> impl IntoResponse {
    let ollama_base = std::env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| "http://ollama.nxtgauge-ai.svc.cluster.local:11434".to_string());
    let model = std::env::var("OLLAMA_CHAT_MODEL").unwrap_or_else(|_| "gemma3:270m".to_string());

    let form_type = body.form_type.unwrap_or_else(|| "generic".to_string());

    let prompt = format!(
        "Extract key:value pairs from this message for a {} form. \
         Return ONLY a JSON object with the fields you can identify. \
         Use camelCase for field names.\n\nMessage: {}",
        form_type, body.message
    );

    let response_text = match call_ollama_inline(&ollama_base, &model, &prompt).await {
        Ok(r) => r,
        Err(e) => {
            tracing::error!("Ollama form extraction error: {}", e);
            return (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Form extraction failed" }))).into_response();
        }
    };

    let extracted: serde_json::Value = serde_json::from_str(&response_text)
        .unwrap_or_else(|_| serde_json::json!({}));

    let mut fields = Vec::new();
    let missing_fields = Vec::new();

    if let Some(obj) = extracted.as_object() {
        for (key, value) in obj {
            fields.push(ExtractedField {
                key: key.clone(),
                value: value.to_string(),
                confidence: 0.8,
            });
        }
    }

    let confidence = if fields.is_empty() { 0.3 } else { 0.75 };

    (StatusCode::OK, Json(FormExtractResponse {
        fields,
        missing_fields,
        confidence,
    })).into_response()
}

#[derive(sqlx::FromRow)]
struct TicketRow {
    id: Uuid,
    subject: String,
    description: Option<String>,
    category: String,
    priority: String,
    status: String,
    requester_name: Option<String>,
    requester_email: Option<String>,
    assigned_to: Option<Uuid>,
    created_at: chrono::DateTime<chrono::Utc>,
    updated_at: chrono::DateTime<chrono::Utc>,
}