use std::os::raw::c_void;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use swift_rs::SRString;
use tokio::sync::{oneshot, mpsc as tokio_mpsc};
use futures_util::Stream;

pub mod agent;
pub use agent::{Agent, AgentEvent};
pub mod mcp;
pub use mcp::McpClient;

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct GenerationOptions {
    pub temperature: f32, // -1.0 for default
    pub max_tokens: i32,   // -1 for default
}

impl Default for GenerationOptions {
    fn default() -> Self {
        Self {
            temperature: -1.0,
            max_tokens: -1,
        }
    }
}

/// Defines a tool (function) that the model can request to call.
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ToolDefinition {
    pub name: String,
    pub description: String,
    pub parameters: serde_json::Value,
    #[serde(skip)]
    pub tool_source: String, // "local" or the name of the MCP server
}

/// A request from the model to call a specific tool.
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ToolCall {
    pub tool_name: String,
    pub arguments: serde_json::Value,
}

/// The response from a model which can either be a plain string or a tool call.
#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(tag = "type", content = "data")]
pub enum ToolCallResponse {
    #[serde(rename = "text")]
    Plain(String),
    #[serde(rename = "tool_call")]
    Call(ToolCall),
}

unsafe extern "C" {
    fn call_apple_llm(
        prompt: *const SRString, 
        instructions: *const SRString,
        session_id: *const SRString,
        temperature: f32, 
        max_tokens: i32, 
        sender: *mut c_void, 
        callback: extern "C" fn(*const u8, usize, *mut c_void)
    );
    fn call_apple_llm_stream(
        prompt: *const SRString, 
        instructions: *const SRString,
        session_id: *const SRString,
        temperature: f32, 
        max_tokens: i32, 
        sender: *mut c_void, 
        callback: extern "C" fn(*const u8, usize, bool, *mut c_void)
    );
    fn apple_llm_prewarm();
    fn apple_llm_clear_session(session_id: *const SRString);
    fn apple_llm_count_tokens(text: *const SRString, sender: *mut c_void, callback: extern "C" fn(i32, *mut c_void));
    fn apple_llm_get_context_size(sender: *mut c_void, callback: extern "C" fn(i32, *mut c_void));
    fn apple_nl_get_embedding(text: *const SRString, sender: *mut c_void, callback: extern "C" fn(*const f32, usize, *mut c_void));
}

/// Prewarms the foundation model by loading it into memory.
pub fn prewarm() {
    unsafe { apple_llm_prewarm() };
}

/// Generates a vector embedding for the given text using Apple's Natural Language framework.
pub async fn get_embedding(text: &str) -> Result<Vec<f32>, String> {
    let (tx, rx) = oneshot::channel();
    
    {
        let text_sr = SRString::from(text);
        let tx_ptr = Box::into_raw(Box::new(tx)) as *mut c_void;

        extern "C" fn on_done(ptr: *const f32, len: usize, sender_ptr: *mut c_void) {
            unsafe {
                let tx = Box::from_raw(sender_ptr as *mut oneshot::Sender<Result<Vec<f32>, String>>);
                if ptr.is_null() || len == 0 {
                    let _ = tx.send(Err("Failed to generate embedding.".to_string()));
                } else {
                    let slice = std::slice::from_raw_parts(ptr, len);
                    let vec = slice.to_vec();
                    let _ = tx.send(Ok(vec));
                }
            }
        }

        unsafe {
            apple_nl_get_embedding(&text_sr, tx_ptr, on_done);
        }
    }

    rx.await.map_err(|e| format!("Oneshot error: {}", e))?
}

/// Counts the number of tokens in a string.
pub async fn count_tokens(text: &str) -> Result<i32, String> {
    let (tx, rx) = oneshot::channel();
    
    {
        let text_sr = SRString::from(text);
        let tx_ptr = Box::into_raw(Box::new(tx)) as *mut c_void;

        extern "C" fn on_done(count: i32, sender_ptr: *mut c_void) {
            unsafe {
                let tx = Box::from_raw(sender_ptr as *mut oneshot::Sender<i32>);
                let _ = tx.send(count);
            }
        }

        unsafe {
            apple_llm_count_tokens(&text_sr, tx_ptr, on_done);
        }
    }

    let res = rx.await.map_err(|e| format!("Oneshot error: {}", e))?;
    if res < 0 {
        Err("Failed to count tokens".to_string())
    } else {
        Ok(res)
    }
}

/// Gets the total context size (max tokens) of the on-device model.
pub async fn get_context_size() -> Result<i32, String> {
    let (tx, rx) = oneshot::channel();
    let tx_ptr = Box::into_raw(Box::new(tx)) as *mut c_void;

    extern "C" fn on_done(size: i32, sender_ptr: *mut c_void) {
        unsafe {
            let tx = Box::from_raw(sender_ptr as *mut oneshot::Sender<i32>);
            let _ = tx.send(size);
        }
    }

    unsafe {
        apple_llm_get_context_size(tx_ptr, on_done);
    }

    rx.await.map_err(|e| format!("Oneshot error: {}", e))
}

/// A stateful conversation session with the foundation model.
#[derive(Clone)]
pub struct Session {
    id: String,
    instructions: String,
}

impl Session {
    /// Creates a new session with optional system instructions (persona).
    pub fn new(id: &str, instructions: &str) -> Self {
        Self {
            id: id.to_string(),
            instructions: instructions.to_string(),
        }
    }

    /// Asynchronously responds to a prompt within this session context.
    pub async fn respond(&self, prompt: &str) -> Result<String, String> {
        self.respond_with_options(prompt, GenerationOptions::default()).await
    }

    pub async fn respond_with_options(&self, prompt: &str, options: GenerationOptions) -> Result<String, String> {
        let (tx, rx) = oneshot::channel();
        
        {
            let prompt_sr = SRString::from(prompt);
            let inst_sr = SRString::from(self.instructions.as_str());
            let id_sr = SRString::from(self.id.as_str());
            
            let tx_ptr = Box::into_raw(Box::new(tx)) as *mut c_void;

            extern "C" fn on_done(ptr: *const u8, len: usize, sender_ptr: *mut c_void) {
                unsafe {
                    let tx = Box::from_raw(sender_ptr as *mut oneshot::Sender<Result<String, String>>);
                    if !ptr.is_null() && len > 0 {
                        let slice = std::slice::from_raw_parts(ptr, len);
                        let response = String::from_utf8_lossy(slice).to_string();
                        let _ = tx.send(Ok(response));
                    } else {
                        let _ = tx.send(Ok(String::new()));
                    }
                }
            }

            unsafe {
                call_apple_llm(&prompt_sr, &inst_sr, &id_sr, options.temperature, options.max_tokens, tx_ptr, on_done);
            }
        }

        rx.await.map_err(|e| format!("Oneshot channel error: {}", e))?
    }

    /// Asynchronously streams the response within this session context.
    pub fn stream(&self, prompt: &str) -> impl Stream<Item = Result<String, String>> + 'static {
        self.stream_with_options(prompt, GenerationOptions::default())
    }

    pub fn stream_with_options(&self, prompt: &str, options: GenerationOptions) -> impl Stream<Item = Result<String, String>> + 'static {
        let (tx, rx) = tokio_mpsc::unbounded_channel();
        
        {
            let prompt_sr = SRString::from(prompt);
            let inst_sr = SRString::from(self.instructions.as_str());
            let id_sr = SRString::from(self.id.as_str());
            
            struct StreamState {
                tx: tokio_mpsc::UnboundedSender<Result<String, String>>,
                previous_len: usize,
            }
            
            let state = StreamState { tx, previous_len: 0 };
            let state_ptr = Box::into_raw(Box::new(state)) as *mut c_void;

            extern "C" fn on_stream(ptr: *const u8, len: usize, is_done: bool, sender_ptr: *mut c_void) {
                unsafe {
                    if is_done {
                        let state = Box::from_raw(sender_ptr as *mut StreamState);
                        if !ptr.is_null() && len > 0 {
                            let slice = std::slice::from_raw_parts(ptr, len);
                            let response = String::from_utf8_lossy(slice).to_string();
                            if response.len() > state.previous_len {
                                let delta = response[state.previous_len..].to_string();
                                let _ = state.tx.send(Ok(delta));
                            }
                        }
                    } else {
                        let state = &mut *(sender_ptr as *mut StreamState);
                        if !ptr.is_null() && len > 0 {
                            let slice = std::slice::from_raw_parts(ptr, len);
                            let response = String::from_utf8_lossy(slice).to_string();
                            
                            if response.len() > state.previous_len {
                                let delta = response[state.previous_len..].to_string();
                                state.previous_len = response.len();
                                let _ = state.tx.send(Ok(delta));
                            }
                        }
                    }
                }
            }

            unsafe {
                call_apple_llm_stream(&prompt_sr, &inst_sr, &id_sr, options.temperature, options.max_tokens, state_ptr, on_stream);
            }
        }

        tokio_stream::wrappers::UnboundedReceiverStream::new(rx)
    }

    pub fn clear(&self) {
        let id_sr = SRString::from(self.id.as_str());
        unsafe { apple_llm_clear_session(&id_sr) };
    }

    pub async fn respond_json<T: DeserializeOwned>(&self, prompt: &str) -> Result<T, String> {
        self.respond_json_with_options(prompt, GenerationOptions::default()).await
    }

    pub async fn respond_json_with_options<T: DeserializeOwned>(&self, prompt: &str, options: GenerationOptions) -> Result<T, String> {
        let json_prompt = format!(
            "{}\n\nReturn ONLY a JSON object. Respond strictly in valid JSON format.",
            prompt
        );
        let response = self.respond_with_options(&json_prompt, options).await?;
        let cleaned = response.trim().trim_start_matches("```json").trim_start_matches("```").trim_end_matches("```").trim();
        serde_json::from_str(cleaned).map_err(|e| format!("Failed to parse JSON: {}\nResponse: {}", e, response))
    }
}

/// Stateless wrapper for the Apple Foundation Models API
pub struct FoundationModel;

impl FoundationModel {
    pub async fn respond(prompt: &str) -> Result<String, String> {
        let session = Session::new("default", "");
        let res = session.respond(prompt).await;
        session.clear();
        res
    }

    pub async fn respond_with_options(prompt: &str, options: GenerationOptions) -> Result<String, String> {
        let session = Session::new("default", "");
        let res = session.respond_with_options(prompt, options).await;
        session.clear();
        res
    }

    pub async fn respond_json<T: DeserializeOwned>(prompt: &str) -> Result<T, String> {
        Self::respond_json_with_options(prompt, GenerationOptions::default()).await
    }

    pub async fn respond_json_with_options<T: DeserializeOwned>(prompt: &str, options: GenerationOptions) -> Result<T, String> {
        let json_prompt = format!(
            "{}\n\nReturn ONLY a JSON object. Respond strictly in valid JSON format.",
            prompt
        );
        let response = Self::respond_with_options(&json_prompt, options).await?;
        let cleaned = response.trim().trim_start_matches("```json").trim_start_matches("```").trim_end_matches("```").trim();
        serde_json::from_str(cleaned).map_err(|e| format!("Failed to parse JSON: {}\nResponse: {}", e, response))
    }

    pub async fn respond_with_tools(prompt: &str, tools: &[ToolDefinition]) -> Result<ToolCallResponse, String> {
        let tools_json = serde_json::to_string(tools).map_err(|e| e.to_string())?;
        let tool_prompt = format!(
            "Tools: {}\nUser: {}\nInstructions:\n1. Use tool_call JSON if needed.\n2. Use text JSON for plain response.\nRespond ONLY with JSON.",
            tools_json, prompt
        );
        Self::respond_json::<ToolCallResponse>(&tool_prompt).await
    }

    pub fn stream(prompt: &str) -> impl Stream<Item = Result<String, String>> + '_ {
        static SESSION: std::sync::LazyLock<Session> = std::sync::LazyLock::new(|| Session::new("default_stream", ""));
        SESSION.stream(prompt)
    }
}
