use png::{BitDepth, ColorType, Transformations};
use pyo3::exceptions::{PyIOError, PyKeyError};
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyTuple};
use touhou_formats::th06::pos::LoopPoints;
use touhou_formats::th06::pbg3;
use touhou_formats::th06::std as stage;
use touhou_formats::th06::msg;
use std::collections::{BTreeMap, HashMap};
use std::fs::File;
use std::io::BufReader;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};

#[cfg(feature = "glide")]
mod glide;

mod audio;

#[pyclass(module = "libtouhou")]
struct PyModel {
    inner: stage::Model,
}

#[pymethods]
impl PyModel {
    #[getter]
    fn quads(&self) -> Vec<(u16, f32, f32, f32, f32, f32)> {
        self.inner.quads.iter().map(|quad| (quad.anm_script, quad.pos.x, quad.pos.y, quad.pos.z, quad.size_override.width, quad.size_override.height)).collect()
    }

    #[getter]
    fn bounding_box(&self) -> [f32; 6] {
        self.inner.bounding_box
    }
}

#[pyclass(module = "libtouhou")]
struct PyStage {
    inner: stage::Stage,
}

#[pymethods]
impl PyStage {
    #[getter]
    fn models(&self) -> Vec<PyModel> {
        self.inner.models.clone().into_iter().map(|inner| PyModel { inner }).collect()
    }

    #[getter]
    fn object_instances(&self) -> Vec<(u16, f32, f32, f32)> {
        self.inner.instances.iter().map(|instance| (instance.id, instance.pos.x, instance.pos.y, instance.pos.z)).collect()
    }

    #[getter]
    fn name(&self) -> &str {
        &self.inner.name
    }

    #[getter]
    fn bgms(&self) -> Vec<Option<(String, String)>> {
        self.inner.musics.clone()
    }

    #[getter]
    fn script(&self, py: Python) -> Vec<(u32, u16, Py<PyTuple>)> {
        fn call_to_python(py: Python, call: &stage::Call) -> PyResult<(u32, u16, Py<PyTuple>)> {
            let (opcode, args) = match call.instr {
                stage::Instruction::SetViewpos(x, y, z) => (0, (x, y, z).into_pyobject(py)?.unbind()),
                stage::Instruction::SetFog(r, g, b, _a, near, far) => (1, (r, g, b, near, far).into_pyobject(py)?.unbind()),
                stage::Instruction::SetViewpos2(x, y, z) => (2, (x, y, z).into_pyobject(py)?.unbind()),
                stage::Instruction::StartInterpolatingViewpos2(frame, _unused1, _unused2) => (3, (frame,).into_pyobject(py)?.unbind()),
                stage::Instruction::StartInterpolatingFog(frame, _unused1, _unused2) => (4, (frame,).into_pyobject(py)?.unbind()),
                stage::Instruction::Unknown(unused1, unused2, unused3) => (5, (unused1, unused2, unused3).into_pyobject(py)?.unbind()),
            };
            Ok((call.time, opcode, args))
        }

        self.inner.script.iter().map(|call| call_to_python(py, call).unwrap()).collect()
    }
}

#[pyclass(module = "libtouhou")]
struct PyMsg {
    inner: msg::Msg,
}

#[pymethods]
impl PyMsg {
    #[getter]
    fn msgs(&self, py: Python) -> BTreeMap<u8, Vec<(u16, u8, Py<PyTuple>)>> {
        fn call_to_python(py: Python, call: &msg::Call) -> PyResult<(u16, u8, Py<PyTuple>)> {
            let (opcode, args) = match &call.instr {
                msg::Instruction::Unk1() => (0, ().into_pyobject(py)?),
                msg::Instruction::Enter(side, effect) => (1, (side, effect).into_pyobject(py)?),
                msg::Instruction::ChangeFace(side, index) => (2, (side, index).into_pyobject(py)?),
                msg::Instruction::DisplayText(side, index, text) => (3, (side, index, text).into_pyobject(py)?),
                msg::Instruction::Pause(duration) => (4, (duration,).into_pyobject(py)?),
                msg::Instruction::Animate(side, effect) => (5, (side, effect).into_pyobject(py)?),
                msg::Instruction::SpawnEnemySprite() => (6, ().into_pyobject(py)?),
                msg::Instruction::ChangeMusic(track) => (7, (track,).into_pyobject(py)?),
                msg::Instruction::DisplayDescription(side, index, text) => (8, (side, index, text).into_pyobject(py)?),
                msg::Instruction::ShowScores(unk1) => (9, (unk1,).into_pyobject(py)?),
                msg::Instruction::Freeze() => (10, ().into_pyobject(py)?),
                msg::Instruction::NextStage() => (11, ().into_pyobject(py)?),
                msg::Instruction::Unk2() => (12, ().into_pyobject(py)?),
                msg::Instruction::SetAllowSkip(boolean) => (13, (boolean,).into_pyobject(py)?),
                msg::Instruction::Unk3() => (14, ().into_pyobject(py)?),
            };
            Ok((call.time, opcode, args.unbind()))
        }
        self.inner.scripts.iter().map(|(index, script)| (
            *index,
            script.into_iter().map(|call| call_to_python(py, call).unwrap()).collect(),
        )).collect()
    }
}

/// A loader for Touhou files.
#[pyclass(module = "libtouhou", subclass)]
#[derive(Default)]
struct Loader {
    /// The file names to the possible executable.
    #[pyo3(get)]
    exe_files: Vec<PathBuf>,

    /// The path to the game directory.
    #[pyo3(get)]
    game_dir: Option<PathBuf>,

    /// A map from inner filenames to the archive containing them.
    known_files: HashMap<String, Arc<Mutex<pbg3::PBG3<BufReader<File>>>>>,
}

#[pymethods]
impl Loader {
    /// Create a new Loader for the given game_dir.
    #[new]
    fn new(game_dir: Option<PathBuf>) -> Loader {
        Loader {
            exe_files: Vec::new(),
            game_dir,
            known_files: HashMap::new(),
        }
    }

    /// Scan the game_dir for archives.
    ///
    /// paths_lists is a list of ':'-separated glob patterns, the first matching file will be used
    /// and the other ones ignored.
    fn scan_archives(&mut self, paths_lists: Vec<String>) -> PyResult<()> {
        for paths in paths_lists.iter() {
            let found_paths: Vec<_> = paths.split(':').map(|path| {
                glob::glob(if let Some(game_dir) = self.game_dir.as_ref() {
                    game_dir.join(path)
                } else {
                    PathBuf::from(path)
                }.to_str().unwrap()).unwrap()
            }).flatten().map(Result::unwrap).map(PathBuf::from).collect();
            if found_paths.is_empty() {
                return Err(PyIOError::new_err(format!("No path found for {paths:?}")));
            }
            let path = &found_paths[0];
            if let Some(extension) = path.extension() && extension == "exe" {
                self.exe_files.extend(found_paths);
            } else {
                let pbg3 = pbg3::from_path_buffered(path)?;
                let filenames: Vec<_> = pbg3.list_files().cloned().collect();
                let pbg3 = Arc::new(Mutex::new(pbg3));
                for name in filenames {
                    self.known_files.insert(name.clone(), Arc::clone(&pbg3));
                }
            }
        }
        Ok(())
    }

    /// Return the given file as a Vec<u8>.
    fn get_file_internal(&self, name: &str) -> PyResult<Vec<u8>> {
        let name = if let Some((_, name)) = name.rsplit_once('/') {
            name
        } else {
            name
        };
        if let Some(archive) = self.known_files.get(name) {
            let mut archive = archive.lock().unwrap();
            let bytes = archive.get_file(name, true)?;
            Ok(bytes)
        } else {
            Err(PyKeyError::new_err(format!("Unknown file {name:?}")))
        }
    }

    /// Return the given file as an io.BytesIO object.
    fn get_file(&self, py: Python, name: &str) -> PyResult<Py<PyAny>> {
        let vec = self.get_file_internal(name)?;
        let bytes = PyBytes::new(py, &vec);
        let io = py.import("io")?;
        let bytesio_class = io.dict().get_item("BytesIO")?.unwrap();
        let file = bytesio_class.call1((bytes,))?;
        Ok(file.unbind())
    }

    fn get_stage(&self, py: Python, name: &str) -> PyResult<Py<PyStage>> {
        let vec = self.get_file_internal(name)?;
        let (_, inner) = stage::Stage::from_slice(&vec).unwrap();
        Ok(Py::new(py, PyStage { inner })?)
    }

    fn get_msg(&self, py: Python, name: &str) -> PyResult<Py<PyMsg>> {
        let vec = self.get_file_internal(name)?;
        let (_, inner) = msg::Msg::from_slice(&vec).unwrap();
        Ok(Py::new(py, PyMsg { inner })?)
    }

    fn get_image(&self, py: Python, name: &str) -> PyResult<Py<Image>> {
        let vec = self.get_file_internal(name)?;
        let image = Image::from_rgb_png(&vec);
        Py::new(py, image)
    }

    fn get_image_with_alpha(&self, py: Python, name: &str, alpha: &str) -> PyResult<Py<Image>> {
        let rgb = self.get_file_internal(name)?;
        let alpha = self.get_file_internal(alpha)?;
        let image = Image::from_rgb_and_alpha_png(&rgb, &alpha);
        Py::new(py, image)
    }
}

impl Loader {
    fn get_loop_points(&self, name: &str) -> Result<LoopPoints, ()> {
        let vec = self.get_file_internal(name).unwrap();
        let (_, inner) = LoopPoints::from_slice(&vec).unwrap();
        Ok(inner)
    }
}

#[pyclass(module = "libtouhou")]
#[derive(Debug)]
struct Image {
    width: u32,
    height: u32,
    data: Vec<u8>,
}

#[pymethods]
impl Image {
    #[getter]
    fn dimensions(&self) -> (u32, u32) {
        (self.width, self.height)
    }

    #[getter]
    fn pixels(&self) -> &[u8] {
        &self.data
    }
}

impl Image {
    fn load_png(data: &[u8], add_alpha: bool) -> (u32, u32, ColorType, Vec<u8>) {
        let cursor = std::io::Cursor::new(data);
        let mut decoder = png::Decoder::new(cursor);
        // Request either rgb8 or rgba8 data.
        decoder.set_transformations(if add_alpha {
            Transformations::ALPHA
        } else {
            Transformations::EXPAND
        });
        let mut reader = decoder.read_info().unwrap();
        let mut buf = vec![0; reader.output_buffer_size().unwrap()];
        let info = reader.next_frame(&mut buf).unwrap();
        assert_eq!(buf.capacity(), info.buffer_size());
        assert_eq!(info.bit_depth, BitDepth::Eight);
        (info.width, info.height, info.color_type, buf)
    }

    fn from_rgb_png(rgb: &[u8]) -> Self {
        let (width, height, color_type, data) = Self::load_png(rgb, true);
        assert_eq!(color_type, ColorType::Rgba);
        Self { width, height, data }
    }

    fn from_rgb_and_alpha_png(rgb: &[u8], alpha: &[u8]) -> Self {
        let (width, height, color_type, data) = Self::load_png(rgb, false);
        if color_type == ColorType::Rgba {
            // TODO: Check which should be used, the alpha channel in the primary PNG, or the alpha
            // mask in the secondary PNG, in case both are present (such as for Patchouli’s face).
            return Self { width, height, data };
        }

        let (alpha_width, alpha_height, color_type, alpha) = Self::load_png(alpha, false);
        assert_eq!(color_type, ColorType::Rgb);
        assert_eq!(width, alpha_width);
        assert_eq!(height, alpha_height);

        let data = data
            .as_chunks::<3>().0
            .into_iter()
            .zip(alpha.as_chunks::<3>().0)
            .flat_map(|([r, g, b], [luma, _, _])| [*r, *g, *b, *luma])
            .collect();
        Self { width, height, data }
    }
}

/// A loader for Touhou files.
#[pyclass(module = "libtouhou")]
struct Prng {
    inner: touhou_utils::prng::Prng,
}

#[pymethods]
impl Prng {
    #[pyo3(signature = (seed=None))]
    #[new]
    fn new(seed: Option<u16>) -> Prng {
        let inner = touhou_utils::prng::Prng::new(seed);
        Prng { inner }
    }

    #[getter]
    fn seed(&self) -> u16 {
        self.inner.get_seed()
    }

    fn rand_uint16(&mut self) -> u16 {
        self.inner.get_u16()
    }

    fn rand_uint32(&mut self) -> u32 {
        self.inner.get_u32()
    }

    fn rand_double(&mut self) -> f64 {
        self.inner.get_f64()
    }
}

#[pymodule]
mod libtouhou {
    #[pymodule_export]
    use super::Loader;

    #[pymodule_export]
    use super::Image;

    #[pymodule_export]
    use super::Prng;

    #[pymodule_export]
    use crate::audio::Audio;

    #[cfg(feature = "glide")]
    #[pymodule_export]
    use super::glide::module;
}