use axum::{
    extract::{Path, State},
    http::{StatusCode, HeaderMap},
    response::{IntoResponse, Json},
    Extension,
};
use std::collections::HashMap;
use sqlx::PgPool;
use base64::prelude::*;
use crate::{FunctionsState, models::{DeployRequest, InvokeRequest, InvokeResponse, Function}};

pub async fn invoke_function(
    State(state): State<FunctionsState>,
    db: Option<Extension<PgPool>>,
    Path(name): Path<String>,
    headers: HeaderMap,
    Json(payload): Json<InvokeRequest>,
) -> impl IntoResponse {
    tracing::info!("Invoking function: {}", name);
    let db = db.map(|Extension(p)| p).unwrap_or_else(|| state.db.clone());

    // Convert headers
    let mut header_map = HashMap::new();
    for (k, v) in headers.iter() {
        if let Ok(val) = v.to_str() {
            header_map.insert(k.as_str().to_string(), val.to_string());
        }
    }

    // 1. Fetch function
    let func = sqlx::query_as::<_, Function>("SELECT * FROM functions.functions WHERE name = $1")
        .bind(&name)
        .fetch_optional(&db)
        .await;

    let func = match func {
        Ok(Some(f)) => f,
        Ok(None) => {
            tracing::warn!("Function not found: {}", name);
            return (StatusCode::NOT_FOUND, "Function not found").into_response();
        },
        Err(e) => {
            tracing::error!("DB error fetching function: {}", e);
            return (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response();
        }
    };

    // 2. Execute
    let result = if func.runtime == "deno" || func.runtime == "typescript" || func.runtime == "javascript" {
        let code = match String::from_utf8(func.code) {
             Ok(c) => c,
             Err(e) => {
                 tracing::error!("Invalid UTF-8 in Deno function code: {}", e);
                 return (StatusCode::INTERNAL_SERVER_ERROR, "Invalid function code".to_string()).into_response();
             }
        };
        state.deno_runtime.execute(code, payload.payload, header_map).await
    } else {
        // Assume WASM
        let payload_str = payload.payload.as_ref().map(|v| v.to_string());
        state.runtime.execute(&func.code, payload_str).await.map(|(out, err)| (out, err, 200, HashMap::new()))
    };

    match result {
        Ok((stdout, stderr, status, headers)) => {
            tracing::info!("Function executed successfully. Stdout len: {}, Stderr len: {}", stdout.len(), stderr.len());
            let resp = InvokeResponse {
                result: Some(stdout),
                error: if stderr.is_empty() { None } else { Some(stderr) },
                logs: vec![],
                status,
                headers: Some(headers),
            };
            Json(resp).into_response()
        },
        Err(e) => {
            tracing::error!("Runtime execution error: {:?}", e);
            (StatusCode::INTERNAL_SERVER_ERROR, format!("Runtime error: {:?}", e)).into_response()
        },
    }
}

pub async fn deploy_function(
    State(state): State<FunctionsState>,
    db: Option<Extension<PgPool>>,
    Json(payload): Json<DeployRequest>,
) -> impl IntoResponse {
    tracing::info!("Deploying function: {}", payload.name);
    let db = db.map(|Extension(p)| p).unwrap_or_else(|| state.db.clone());

    // Decode base64
    let code = match BASE64_STANDARD.decode(&payload.code_base64) {
        Ok(c) => c,
        Err(e) => {
             tracing::error!("Invalid base64: {}", e);
             return (StatusCode::BAD_REQUEST, format!("Invalid base64: {}", e)).into_response();
        }
    };

    // Store in DB
    let runtime = payload.runtime.unwrap_or("wasm".to_string());
    
    let res = sqlx::query(
        "INSERT INTO functions.functions (name, code, runtime) VALUES ($1, $2, $3) ON CONFLICT (name) DO UPDATE SET code = $2, runtime = $3, updated_at = NOW() RETURNING id"
    )
    .bind(&payload.name)
    .bind(&code)
    .bind(&runtime)
    .fetch_one(&db)
    .await;

    match res {
        Ok(_) => {
            tracing::info!("Function deployed successfully");
            (StatusCode::OK, "Function deployed").into_response()
        },
        Err(e) => {
            tracing::error!("DB error deploying function: {}", e);
            (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response()
        },
    }
}