use std::fmt::Display;

use anyhow::Result;
use regex::Regex;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum Language {
    Python,
    Typescript,
}

impl Display for Language {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Language::Python => write!(f, "python"),
            Language::Typescript => write!(f, "typescript"),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct CodeBlock {
    pub language: Language,
    pub block: String,
}

impl Language {
    fn from_str(s: &str) -> Result<Self> {
        match s {
            "python" => Ok(Language::Python),
            "typescript" => Ok(Language::Typescript),
            _ => Err(anyhow::anyhow!("Unsupported language: {}", s)),
        }
    }
}

pub fn parse_code_block_from_response_content(content: String) -> Result<CodeBlock> {
    let block_regex = Regex::new(r"```(python|typescript)\n([\s\S]*?)```")?;

    if let Some(captures) = block_regex.captures(&content) {
        let language = Language::from_str(captures.get(1).unwrap().as_str())?;
        let code = captures.get(2).unwrap().as_str().trim().to_string();

        return Ok(CodeBlock {
            language,
            block: code,
        });
    }

    Err(anyhow::anyhow!("No code block found in response"))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_code_block() -> Result<()> {
        let response = include_str!("../../tests/test_data/llm_code_response.txt");
        let expected = include_str!("../../tests/test_data/llm_parsed_code.txt");

        let parsed = parse_code_block_from_response_content(response.to_string())?;
        assert_eq!(parsed.language, Language::Python);
        assert_eq!(parsed.block.trim(), expected.trim());
        Ok(())
    }
}
