use own_assist_common::config_loader::ConfigLoadingError;
use std::collections::HashMap;
use ollama_rs::headers::{HeaderMap, HeaderValue};
use ollama_rs::Ollama;
use own_mcp::audio::{AudioClient, AudioClientTrait};
use rmcp::model::Implementation;
use serde::Deserialize;
use url::Url;
use own_mcp::mcp;
use own_mcp::mcp::{guaranteed_mcp_server_name, MCPClient};
use own_assist_common::config_from_file;

#[derive(Debug, Deserialize)]
pub struct Config {
    ollama: OllamaConfig,
    #[serde(rename = "audio-server")]
    audio_server: AudioServerConfig,
    #[serde(rename = "mcp-servers")]
    mcp_servers: Vec<MCPServerConfig>
}

#[derive(Debug, Deserialize)]
pub struct AudioServerConfig {
    url: Url,
    authorization: Option<String>,
}

impl Config {
    pub fn from_file() -> Result<Self, ConfigLoadingError> {
        config_from_file("assist.toml")
    }
    
    pub fn ollama_config(&self) -> &OllamaConfig {
        &self.ollama
    }

    pub fn ollama_instance(&self) -> Ollama {
        let mut ollama_headers = HeaderMap::new();

        if let Some(authorization_header) = &self.ollama.authorization {
            ollama_headers.append("Authorization", HeaderValue::from_str(&authorization_header).unwrap());
        }

        let ollama_default_url = Url::parse("http://127.0.0.1:11434").unwrap();
        let mut ollama = Ollama::from_url(self.ollama.url.clone().unwrap_or(ollama_default_url));
        ollama.set_headers(Some(ollama_headers));

        ollama
    }

    pub fn audio_client(&self) -> AudioClient {
        let mut audio_client = AudioClient::new(self.audio_server.url.clone());
        if let Some(authorization) = &self.audio_server.authorization {
            audio_client = audio_client.with_authorization(authorization.clone());
        }

        audio_client
    }

    pub async fn mcp_clients(&self) -> HashMap<String, MCPClient> {
        let mut mcp_clients: HashMap<String, MCPClient> = HashMap::new();

        for mcp_server in &self.mcp_servers {
            let client = mcp::get_client(
                mcp_server.url.as_str(),
                mcp_server.authorization.clone(),
                Implementation::new("own_assist", env!("CARGO_PKG_VERSION")),
            ).await.unwrap();

            let server_name = guaranteed_mcp_server_name(mcp_server.name.clone(), &client);
            let server_data = client;

            mcp_clients.insert(server_name.clone(), server_data);
        }

        mcp_clients
    }
}

#[derive(Debug, Deserialize)]
pub struct OllamaConfig {
    pub url: Option<Url>,
    pub model: OllamaModelConfig,
    pub authorization: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct OllamaModelConfig {
    pub from_model: String,
    pub name: String,
    pub context_size: Option<u64>,
    pub temperature: Option<f32>,
}

#[derive(Debug, Deserialize)]
pub struct MCPServerConfig {
    pub name: Option<String>,
    pub url: Url,
    pub authorization: Option<String>, // should probably implement oauth some time // actually, fuck oauth
}