use orch::lm::{
    AnthropicBuilder, LanguageModel, LanguageModelBuilder, LanguageModelProvider, OllamaBuilder,
    OpenAiBuilder,
};

pub fn get_lm() -> (Box<dyn LanguageModel>, LanguageModelProvider) {
    let args = std::env::args().collect::<Vec<_>>();
    let provider_name = args.get(1).unwrap_or_else(|| {
        eprintln!("ERROR: Please provide a provider name");
        std::process::exit(1);
    });
    let provider = LanguageModelProvider::try_from(provider_name.as_str())
        .expect("Invalid provider name. Supported values: 'ollama', 'openai', 'anthropic'");

    let open_ai_api_key = {
        if provider == LanguageModelProvider::OpenAi {
            std::env::var("OPENAI_API_KEY")
                .unwrap_or_else(|_| panic!("OPENAI_API_KEY environment variable not set"))
        } else {
            String::new()
        }
    };
    let anthropic_api_key = {
        if provider == LanguageModelProvider::Anthropic {
            std::env::var("ANTHROPIC_API_KEY")
                .unwrap_or_else(|_| panic!("ANTHROPIC_API_KEY environment variable not set"))
        } else {
            String::new()
        }
    };
    let lm: Box<dyn LanguageModel> = match provider {
        LanguageModelProvider::Ollama => Box::new(OllamaBuilder::new().try_build().unwrap()),
        LanguageModelProvider::OpenAi => Box::new(
            OpenAiBuilder::new()
                .with_api_key(open_ai_api_key)
                .try_build()
                .unwrap(),
        ),
        LanguageModelProvider::Anthropic => Box::new(
            AnthropicBuilder::new()
                .with_api_key(anthropic_api_key)
                .try_build()
                .unwrap(),
        ),
    };

    (lm, provider)
}

#[allow(dead_code)]
fn main() {}