// Copyright 2018-2021 System76 <info@system76.com>
//
// SPDX-License-Identifier: GPL-3.0-only

use crate::pci::PciBus;
use serde::{Deserialize, Serialize};
use std::{fs, io, process::ExitStatus};
use sysfs_class::{PciDevice, SysClass};

const EXTERNAL_DISPLAY_REQUIRES_NVIDIA: &[&str] = &[
    "addw1",
    "addw2",
    "addw3",
    "addw4",
    "addw5",
    "bonw15",
    "bonw15-b",
    "bonw16",
    "gaze14",
    "gaze15",
    "gaze16-3050",
    "gaze16-3060",
    "gaze16-3060-b",
    "gaze17-3050",
    "gaze17-3060-b",
    "gaze20",
    "kudu6",
    "oryp4",
    "oryp4-b",
    "oryp5",
    "oryp6",
    "oryp7",
    "oryp8",
    "oryp9",
    "oryp10",
    "oryp11",
    "oryp12",
    "oryp13",
    "serw13",
    "serw14",
];

#[derive(Debug, thiserror::Error)]
pub enum GraphicsDeviceError {
    #[error("failed to execute {} command: {}", cmd, why)]
    Command { cmd: &'static str, why: io::Error },
    #[error("{} in use by {}", func, driver)]
    DeviceInUse { func: String, driver: String },
    #[error("failed to probe driver features: {}", _0)]
    Json(io::Error),
    #[error("failed to fetch list of active kernel modules: {}", _0)]
    ModulesFetch(io::Error),
    #[error("does not have switchable graphics")]
    NotSwitchable,
    #[error("PCI driver error on {}: {}", device, why)]
    PciDriver { device: String, why: io::Error },
    #[error("failed to remove PCI device {}: {}", device, why)]
    Remove { device: String, why: io::Error },
    #[error("failed to rescan PCI bus: {}", _0)]
    Rescan(io::Error),
    #[error("failed to access sysfs info: {}", _0)]
    SysFs(io::Error),
    #[error("failed to unbind {} on PCI driver {}: {}", func, driver, why)]
    Unbind { func: String, driver: String, why: io::Error },
    #[error("update-initramfs failed with {} status", _0)]
    UpdateInitramfs(ExitStatus),
    #[error("failed to access Xserver config: {}", _0)]
    XserverConf(io::Error),
}

pub struct GraphicsDevice {
    id:        String,
    devid:     u16,
    functions: Vec<PciDevice>,
}

impl GraphicsDevice {
    #[must_use]
    pub fn new(id: String, devid: u16, functions: Vec<PciDevice>) -> Self {
        Self { id, devid, functions }
    }

    #[must_use]
    pub fn exists(&self) -> bool { self.functions.iter().any(|func| func.path().exists()) }

    #[must_use]
    pub const fn device(&self) -> u16 { self.devid }

    pub unsafe fn unbind(&self) -> Result<(), GraphicsDeviceError> {
        for func in &self.functions {
            if func.path().exists() {
                match func.driver() {
                    Ok(driver) => {
                        log::info!("{}: Unbinding {}", driver.id(), func.id());
                        driver.unbind(func).map_err(|why| GraphicsDeviceError::Unbind {
                            driver: driver.id().to_owned(),
                            func: func.id().to_owned(),
                            why,
                        })?;
                    }
                    Err(why) => match why.kind() {
                        io::ErrorKind::NotFound => (),
                        _ => {
                            return Err(GraphicsDeviceError::PciDriver {
                                device: self.id.clone(),
                                why,
                            })
                        }
                    },
                }
            }
        }

        Ok(())
    }

    pub unsafe fn remove(&self) -> Result<(), GraphicsDeviceError> {
        for func in &self.functions {
            if func.path().exists() {
                match func.driver() {
                    Ok(driver) => {
                        log::error!("{}: in use by {}", func.id(), driver.id());
                        return Err(GraphicsDeviceError::DeviceInUse {
                            func:   func.id().to_owned(),
                            driver: driver.id().to_owned(),
                        });
                    }
                    Err(why) => match why.kind() {
                        io::ErrorKind::NotFound => {
                            log::info!("{}: Removing", func.id());
                            func.remove().map_err(|why| GraphicsDeviceError::Remove {
                                device: self.id.clone(),
                                why,
                            })?;
                        }
                        _ => {
                            return Err(GraphicsDeviceError::PciDriver {
                                device: self.id.clone(),
                                why,
                            })
                        }
                    },
                }
            } else {
                log::warn!("{}: Already removed", func.id());
            }
        }

        Ok(())
    }
}

// supported-gpus.json
#[derive(Serialize, Deserialize, Debug)]
struct NvidiaDevice {
    devid:        String,
    subdeviceid:  Option<String>,
    subvendorid:  Option<String>,
    name:         String,
    legacybranch: Option<String>,
    features:     Vec<String>,
}

#[derive(Serialize, Deserialize, Debug)]
struct SupportedGpus {
    chips: Vec<NvidiaDevice>,
}

pub struct Graphics {
    pub bus:    PciBus,
    pub amd:    Vec<GraphicsDevice>,
    pub intel:  Vec<GraphicsDevice>,
    pub nvidia: Vec<GraphicsDevice>,
    pub other:  Vec<GraphicsDevice>,
}

impl Graphics {
    pub fn new() -> io::Result<Self> {
        let bus = PciBus::new()?;

        log::info!("Rescanning PCI bus");
        bus.rescan()?;

        let devs = PciDevice::all()?;

        let functions = |parent: &PciDevice| -> Vec<PciDevice> {
            let mut functions = Vec::new();
            if let Some(parent_slot) = parent.id().split('.').next() {
                for func in &devs {
                    if let Some(func_slot) = func.id().split('.').next() {
                        if func_slot == parent_slot {
                            log::info!("{}: Function for {}", func.id(), parent.id());
                            functions.push(func.clone());
                        }
                    }
                }
            }
            functions
        };

        let mut amd = Vec::new();
        let mut intel = Vec::new();
        let mut nvidia = Vec::new();
        let mut other = Vec::new();
        for dev in &devs {
            let c = dev.class()?;
            if (c >> 16) & 0xFF == 0x03 {
                match dev.vendor()? {
                    0x1002 => {
                        log::info!("{}: AMD graphics", dev.id());
                        amd.push(GraphicsDevice::new(
                            dev.id().to_owned(),
                            dev.device()?,
                            functions(dev),
                        ));
                    }
                    0x10DE => {
                        log::info!("{}: NVIDIA graphics", dev.id());
                        nvidia.push(GraphicsDevice::new(
                            dev.id().to_owned(),
                            dev.device()?,
                            functions(dev),
                        ));
                    }
                    0x8086 => {
                        log::info!("{}: Intel graphics", dev.id());
                        intel.push(GraphicsDevice::new(
                            dev.id().to_owned(),
                            dev.device()?,
                            functions(dev),
                        ));
                    }
                    vendor => {
                        log::info!("{}: Other({:X}) graphics", dev.id(), vendor);
                        other.push(GraphicsDevice::new(
                            dev.id().to_owned(),
                            dev.device()?,
                            functions(dev),
                        ));
                    }
                }
            }
        }

        Ok(Self { bus, amd, intel, nvidia, other })
    }

    pub fn is_desktop(&self) -> bool {
        let chassis = fs::read_to_string("/sys/class/dmi/id/chassis_type")
            .map_err(GraphicsDeviceError::SysFs)
            .unwrap_or_default();

        chassis.trim() == "3"
    }

    #[must_use]
    pub fn can_switch(&self) -> bool {
        !self.is_desktop()
            && (!self.nvidia.is_empty() && (!self.intel.is_empty() || !self.amd.is_empty()))
    }

    pub fn get_external_displays_require_dgpu(&self) -> Result<bool, GraphicsDeviceError> {
        self.switchable_or_fail()?;

        let model = fs::read_to_string("/sys/class/dmi/id/product_version")
            .map_err(GraphicsDeviceError::SysFs)?;

        Ok(EXTERNAL_DISPLAY_REQUIRES_NVIDIA.contains(&model.trim()))
    }

    fn switchable_or_fail(&self) -> Result<(), GraphicsDeviceError> {
        if self.can_switch() {
            Ok(())
        } else {
            Err(GraphicsDeviceError::NotSwitchable)
        }
    }
}
