//! Model atomic ops for easier analysis.

use std::fmt::Display;

use rustc_middle::{mir::{BasicBlock, Body, Operand, Place}, ty::{Instance, TyCtxt}};

use crate::aa::facts::FullPlaceRef;
use log::debug;

// TODO: We need a good way to track the op ordering.
#[derive(Debug)]
pub enum AtomicOps<'tcx> {
    Load {
        ptr: FullPlaceRef<'tcx>,
        // ordering: Ordering
    },
    Store {
        ptr: FullPlaceRef<'tcx>,
        value: Operand<'tcx>,
        // ordering: Ordering
    },
    LoadStore {
        ptr: FullPlaceRef<'tcx>,
        value: FullPlaceRef<'tcx>,
        // success_ordering: Ordering,
        // failure_ordering: Ordering,
    },
}

impl<'tcx> AtomicOps<'tcx> {
    pub fn get_ptr(&self) -> FullPlaceRef<'tcx> {
        match self {
            &AtomicOps::Load{ptr} => ptr,
            &AtomicOps::Store{ptr, ..} => ptr,
            &AtomicOps::LoadStore { ptr, ..} => ptr,
        }
    }
}

impl<'tcx> Display for AtomicOps<'tcx> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            AtomicOps::Load{ ptr} => {
                write!(f, "Load({})", ptr)
            }
            AtomicOps::Store{ ptr, value} => {
                write!(f, "Store({}, {:?})", ptr, value)
            }
            _ => unimplemented!()
        }
    }
}

// pub fn try_parse_atomic_op<'tcx>(tcx: TyCtxt<'tcx>, args: &[Operand<'tcx>], callee_body: &Body<'tcx>) -> Option<AtomicOps<'tcx>> {
//     match tcx.def_path_str(callee_body.source.def_id()).as_str() {
//         // "std::sync::atomic::atomic_load" | "core::sync::atomic::atomic_load" => {
//         "std::sync::atomic::AtomicBool::load" | "core::sync::atomic::AtomicBool::load" => {
//             assert!(args.len() == 2);
//             Some(AtomicOps::Load{
//                 ptr: args[0].clone()
//             })
//         },
//         "std::sync::atomic::atomic_store" | "core::sync::atomic::atomic_store" => {
//             assert!(args.len() == 3);
//             Some(AtomicOps::Store {
//                 ptr: args[0].clone(),
//                 value: args[1].clone(),
//             })
//         },
//         "std::sync::atomic::atomic_add" | "core::sync::atomic::atomic_add" => {
//             assert!(args.len() == 3);
//             Some(AtomicOps::LoadStore {
//                 ptr: args[0].clone(),
//                 value: args[1].clone()
//             })
//         },
//         "std::sync::atomic::atomic_compare_exchange" | "core::sync::atomic::atomic_compare_exchange" => {
//             assert!(args.len() == 5);
//             Some(AtomicOps::LoadStore {
//                 ptr: args[0].clone(),
//                 value: args[2].clone()
//             })
//         },
//         _ => None,
//     }
// }

pub fn try_parse_atomic_op<'tcx>(
    tcx: TyCtxt<'tcx>,
    args: &[Operand<'tcx>],
    callee_body: &Body<'tcx>,
    caller_body: &Body<'tcx>,
    dest: Option<&(Place<'tcx>, BasicBlock)>,
) -> Option<AtomicOps<'tcx>> {
    match tcx.def_path_str(callee_body.source.def_id()).as_str() {
        "std::sync::atomic::AtomicBool::load" | "core::sync::atomic::AtomicBool::load" => {
            assert!(args.len() == 2);
            match args[0] {
                Operand::Copy(place) | Operand::Move(place) => {
                    let ptr = FullPlaceRef::new(caller_body.source.def_id(), place.as_ref());
                    Some(AtomicOps::Load { ptr })
                }
                _ => None,
            }
        }
        "std::sync::atomic::AtomicBool::store" | "core::sync::atomic::AtomicBool::store" => {
            assert!(args.len() == 3);
            match args[0] {
                Operand::Copy(place) | Operand::Move(place) => {
                    let ptr = FullPlaceRef::new(caller_body.source.def_id(), place.as_ref());
                    Some(AtomicOps::Store {
                        ptr,
                        value: args[1].clone(),
                    })
                }
                _ => None,
            }
        }
        "std::sync::atomic::AtomicBool::compare_exchange"
        | "core::sync::atomic::AtomicBool::compare_exchange" => {
            assert!(args.len() == 5);
            match args[0] {
                Operand::Copy(place) | Operand::Move(place) => {
                    let ptr = FullPlaceRef::new(caller_body.source.def_id(), place.as_ref());
                    Some(AtomicOps::Store {
                        ptr,
                        value: args[2].clone(),
                    })
                }
                _ => None,
            }
        }
        _ => None,
    }
}
