use pyo3::prelude::*;
use std::collections::HashMap;
use std::path::PathBuf;

use code2prompt_core::configuration::Code2PromptConfigBuilder;
use code2prompt_core::session::Code2PromptSession;
use code2prompt_core::sort::FileSortMethod;
use code2prompt_core::template::OutputFormat;
use code2prompt_core::tokenizer::{TokenFormat, TokenizerType};

#[pyclass(from_py_object)]
#[derive(Clone)]
struct PyCode2PromptSession {
    inner: Code2PromptSession,
}

#[pymethods]
impl PyCode2PromptSession {
    #[new]
    fn new(path: &str) -> PyResult<Self> {
        let config = Code2PromptConfigBuilder::default()
            .path(PathBuf::from(path))
            .build()
            .map_err(|e| {
                PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
                    "Failed to create config: {}",
                    e
                ))
            })?;

        Ok(Self {
            inner: Code2PromptSession::new(config),
        })
    }

    // Configure methods that modify the config
    fn include(&mut self, patterns: Vec<String>) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        config.include_patterns = patterns;
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    fn exclude(&mut self, patterns: Vec<String>) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        config.exclude_patterns = patterns;
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    fn with_line_numbers(&mut self, value: bool) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        config.line_numbers = value;
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    fn with_absolute_paths(&mut self, value: bool) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        config.absolute_path = value;
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    fn with_full_directory_tree(&mut self, value: bool) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        config.full_directory_tree = value;
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    fn with_code_blocks(&mut self, value: bool) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        config.no_codeblock = !value; // Invert because API is different
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    fn follow_symlinks(&mut self, value: bool) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        config.follow_symlinks = value;
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    fn include_hidden(&mut self, value: bool) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        config.hidden = value;
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    fn no_ignore(&mut self, value: bool) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        config.no_ignore = value;
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    fn sort_by(&mut self, method: &str) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        match method.to_lowercase().as_str() {
            "name" | "name_asc" => config.sort_method = Some(FileSortMethod::NameAsc),
            "name_desc" => config.sort_method = Some(FileSortMethod::NameDesc),
            "date" | "date_asc" => config.sort_method = Some(FileSortMethod::DateAsc),
            "date_desc" => config.sort_method = Some(FileSortMethod::DateDesc),
            _ => {
                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
                    "Invalid sort method: {}. Valid values: name_asc, name_desc, date_asc, date_desc",
                    method
                )));
            }
        }
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    fn output_format(&mut self, format: &str) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        match format.to_lowercase().as_str() {
            "markdown" => config.output_format = OutputFormat::Markdown,
            // Assuming from the error that there's a Plain variant - please replace if needed
            "xml" | "text" => config.output_format = OutputFormat::Xml,
            "json" => config.output_format = OutputFormat::Json,
            _ => {
                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
                    "Invalid output format: {}",
                    format
                )));
            }
        }
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    fn with_token_encoding(&mut self, encoding: &str) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        match encoding.to_lowercase().as_str() {
            "cl100k" => config.encoding = TokenizerType::Cl100kBase,
            "o200k" => config.encoding = TokenizerType::O200kBase,
            "p50k" => config.encoding = TokenizerType::P50kBase,
            "p50k_edit" => config.encoding = TokenizerType::P50kEdit,
            "r50k" => config.encoding = TokenizerType::R50kBase,
            _ => {
                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
                    "Invalid token encoding: {}",
                    encoding
                )));
            }
        }
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    fn with_token_format(&mut self, format: &str) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        match format.to_lowercase().as_str() {
            "raw" => config.token_format = TokenFormat::Raw,
            "format" => config.token_format = TokenFormat::Format,
            _ => {
                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
                    "Invalid token format: {}. Use 'raw' or 'format'.",
                    format
                )));
            }
        }
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    #[pyo3(signature = (template, name=None))]
    fn with_template(&mut self, template: String, name: Option<String>) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        config.template_str = template;
        if let Some(name_val) = name {
            config.template_name = name_val;
        } else {
            config.template_name = "custom".to_string();
        }
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    #[pyo3(signature = (key, value))]
    fn with_variable(&mut self, key: String, value: String) -> PyResult<Py<Self>> {
        let mut config = self.inner.config.clone();
        config.user_variables.insert(key, value);
        self.inner = Code2PromptSession::new(config);

        Python::attach(|py| {
            Ok(Py::new(
                py,
                Self {
                    inner: self.inner.clone(),
                },
            )?)
        })
    }

    fn generate(&mut self) -> PyResult<String> {
        match self.inner.generate_prompt() {
            Ok(rendered) => Ok(rendered.prompt),
            Err(e) => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
                "Failed to generate prompt: {}",
                e
            ))),
        }
    }

    fn info(&self) -> PyResult<HashMap<String, String>> {
        // Since there's no direct info() method, we'll create a simple info map
        let mut info = HashMap::new();
        info.insert(
            "path".to_string(),
            self.inner.config.path.to_string_lossy().to_string(),
        );
        info.insert(
            "include_patterns".to_string(),
            format!("{:?}", self.inner.config.include_patterns),
        );
        info.insert(
            "exclude_patterns".to_string(),
            format!("{:?}", self.inner.config.exclude_patterns),
        );

        Ok(info)
    }

    fn token_count(&self) -> PyResult<usize> {
        // Generate the prompt and count tokens
        match self.inner.clone().generate_prompt() {
            Ok(rendered) => Ok(rendered.token_count),
            Err(e) => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
                "Failed to count tokens: {}",
                e
            ))),
        }
    }
}

// Module definition - Updated PyO3 syntax
#[pymodule(name = "code2prompt_rs")]
fn code2prompt_rs(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyCode2PromptSession>()?;
    Ok(())
}
