use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

use serde::{Deserialize, Serialize};

use crate::cache::Cache;
use crate::env::EnvLoader;
use crate::fingerprint::Fingerprinter;
use crate::graph::{Action, Graph, shell_quote};
use crate::{EnvMap, Error, Reporter, Result};

#[derive(Clone, Default)]
pub struct CancelToken(Arc<AtomicBool>);

impl CancelToken {
    pub fn new() -> Self {
        CancelToken(Arc::new(AtomicBool::new(false)))
    }
    pub fn cancel(&self) {
        self.0.store(true, Ordering::SeqCst);
    }
    pub fn is_cancelled(&self) -> bool {
        self.0.load(Ordering::SeqCst)
    }
}

pub struct RunOptions {
    pub root: PathBuf,
    pub cache: Option<Cache>,
    pub force: bool,
    pub extra_args: Vec<String>,
    pub env: EnvMap,
    pub reporter: Box<dyn Reporter>,
    pub cancelled: CancelToken,
}

impl RunOptions {
    pub fn new(root: PathBuf) -> Self {
        RunOptions {
            root,
            cache: None,
            force: false,
            extra_args: Vec::new(),
            env: EnvMap::new(),
            reporter: Box::new(crate::NoopReporter),
            cancelled: CancelToken::new(),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ReadyTask {
    pub name: String,
    pub action: Action,
    /// Extra arguments; non-empty only for the run's target.
    pub args: Vec<String>,
    pub env: EnvMap,
}

impl ReadyTask {
    pub fn display(&self) -> String {
        let mut text = self.action.display();
        for arg in &self.args {
            text.push(' ');
            text.push_str(&shell_quote(arg));
        }
        text
    }
}

#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
pub struct FailedTask {
    pub name: String,
    pub code: u8,
}

#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
pub struct Report {
    pub code: u8,
    pub completed: Vec<String>,
    pub failed: Vec<FailedTask>,
    pub cached: Vec<String>,
}

impl Report {
    pub fn ok(&self) -> bool {
        self.code == 0
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub enum Step {
    Ready(Vec<ReadyTask>),
    Done(Report),
}

struct Frame {
    name: String,
    claim: usize,
    parent: Option<usize>,
    dep_index: usize,
    env: EnvMap,
    ancestors: Vec<String>,
    dep_outcomes: Vec<Option<u8>>,
    active_children: HashSet<usize>,
    deps_started: bool,
    waiting_claim: Option<usize>,
    dispatched: bool,
    surfaced: bool,
    done: bool,
}

struct Claim {
    owner: Option<usize>,
    outcome: Option<u8>,
    fingerprint: Option<String>,
    /// Whether the task's own action was dispatched (vs. completing
    /// because a dependency failed). Only dispatched tasks are failures.
    dispatched: bool,
}

pub struct Execution {
    graph: Graph,
    root: PathBuf,
    cache: Option<Cache>,
    force: bool,
    extra_args: Vec<String>,
    base_env: EnvMap,
    reporter: Box<dyn Reporter>,
    cancelled: CancelToken,
    env: EnvLoader,
    fingerprinter: Fingerprinter,
    frames: Vec<Frame>,
    claims: Vec<Claim>,
    claim_by_name: HashMap<String, usize>,
    root_frame: usize,
    completed: Vec<String>,
    cached: Vec<String>,
}

impl Execution {
    pub fn new(graph: Graph, target: &str, opts: RunOptions) -> Result<Self> {
        graph.get(target)?;
        let root = opts.root.clone();
        let mut execution = Execution {
            fingerprinter: Fingerprinter::new(graph.clone(), root.clone()),
            env: EnvLoader::new(root.clone()),
            graph,
            root,
            cache: opts.cache,
            force: opts.force,
            extra_args: opts.extra_args,
            base_env: opts.env,
            reporter: opts.reporter,
            cancelled: opts.cancelled,
            frames: Vec::new(),
            claims: Vec::new(),
            claim_by_name: HashMap::new(),
            root_frame: 0,
            completed: Vec::new(),
            cached: Vec::new(),
        };
        let base = execution.base_env.clone();
        let root_frame = execution.create_frame(None, 0, target, &base, Vec::new())?;
        execution.root_frame = root_frame;
        Ok(execution)
    }

    fn create_frame(
        &mut self,
        parent: Option<usize>,
        dep_index: usize,
        name: &str,
        inherited_env: &EnvMap,
        ancestors: Vec<String>,
    ) -> Result<usize> {
        if let Some(start) = ancestors.iter().position(|a| a == name) {
            let mut cycle = ancestors[start..].to_vec();
            cycle.push(name.to_owned());
            return Err(Error::Cycle(cycle.join(" -> ")));
        }
        let task = self.graph.get(name)?.clone();
        let env = self.env.effective(&task, inherited_env)?;
        let id = self.frames.len();
        let (claim, waiting_claim) = match self.claim_by_name.get(name) {
            Some(&claim) => (claim, Some(claim)),
            None => {
                let claim = self.claims.len();
                self.claims.push(Claim {
                    owner: Some(id),
                    outcome: None,
                    fingerprint: None,
                    dispatched: false,
                });
                self.claim_by_name.insert(name.to_owned(), claim);
                (claim, None)
            }
        };
        self.frames.push(Frame {
            name: name.to_owned(),
            claim,
            parent,
            dep_index,
            env,
            ancestors,
            dep_outcomes: vec![None; task.dependencies.len()],
            active_children: HashSet::new(),
            deps_started: false,
            waiting_claim,
            dispatched: false,
            surfaced: false,
            done: false,
        });
        Ok(id)
    }

    fn complete(&mut self, frame_id: usize, code: u8) {
        let (claim, parent, dep_index) = {
            let frame = &mut self.frames[frame_id];
            if frame.done {
                return;
            }
            frame.done = true;
            (frame.claim, frame.parent, frame.dep_index)
        };
        if self.claims[claim].owner == Some(frame_id) && self.claims[claim].outcome.is_none() {
            self.claims[claim].outcome = Some(code);
            let name = self.frames[frame_id].name.clone();
            self.completed.push(name);
        }
        if let Some(parent) = parent {
            self.frames[parent].dep_outcomes[dep_index] = Some(code);
            self.frames[parent].active_children.remove(&frame_id);
        }
    }

    /// Advances every frame whose dependencies have resolved, marking
    /// cached tasks complete and dispatchable tasks ready. Returns whether
    /// anything changed.
    fn advance(&mut self) -> Result<bool> {
        let mut progressed = false;
        let mut index = 0;
        while index < self.frames.len() {
            let frame = &self.frames[index];
            if frame.done || frame.dispatched {
                index += 1;
                continue;
            }
            // A frame borrowing another frame's in-flight execution
            // completes with that execution's outcome.
            if let Some(claim) = frame.waiting_claim {
                if let Some(code) = self.claims[claim].outcome {
                    self.complete(index, code);
                    progressed = true;
                }
                index += 1;
                continue;
            }
            let task = self.graph.get(&frame.name)?.clone();

            // Start dependencies.
            if !frame.deps_started || !task.parallel && frame.active_children.is_empty() {
                let next = if task.parallel {
                    if frame.deps_started {
                        None
                    } else {
                        Some(0..task.dependencies.len())
                    }
                } else {
                    // Sequential: a failed dependency ends the task now.
                    if let Some(code) = frame.dep_outcomes.iter().flatten().find(|code| **code != 0)
                    {
                        let code = *code;
                        self.complete(index, code);
                        progressed = true;
                        index += 1;
                        continue;
                    }
                    let started = frame.dep_outcomes.iter().filter(|o| o.is_some()).count()
                        + frame.active_children.len();
                    if started < task.dependencies.len() {
                        Some(started..started + 1)
                    } else {
                        None
                    }
                };
                if let Some(range) = next {
                    self.frames[index].deps_started = true;
                    let env = self.frames[index].env.clone();
                    let mut child_ancestors = self.frames[index].ancestors.clone();
                    child_ancestors.push(self.frames[index].name.clone());
                    for dep_index in range {
                        let dependency = task.dependencies[dep_index].clone();
                        let child = self.create_frame(
                            Some(index),
                            dep_index,
                            &dependency,
                            &env,
                            child_ancestors.clone(),
                        )?;
                        self.frames[index].active_children.insert(child);
                    }
                    progressed = true;
                    index += 1;
                    continue;
                }
            }
            if !self.frames[index].active_children.is_empty() {
                index += 1;
                continue;
            }
            if task.parallel
                && let Some(_failed) = frame.dep_outcomes.iter().flatten().find(|c| **c != 0)
            {
                self.complete(index, 1);
                progressed = true;
                index += 1;
                continue;
            }
            if !task.parallel
                && let Some(code) = frame.dep_outcomes.iter().flatten().find(|c| **c != 0)
            {
                let code = *code;
                self.complete(index, code);
                progressed = true;
                index += 1;
                continue;
            }

            // All dependencies succeeded: resolve this task's own action.
            match task.action {
                Action::None => {
                    self.complete(index, 0);
                    progressed = true;
                }
                _ => {
                    if task.files.is_some() {
                        let args: Vec<String> = if index == self.root_frame {
                            self.extra_args.clone()
                        } else {
                            Vec::new()
                        };
                        let env = self.frames[index].env.clone();
                        let name = self.frames[index].name.clone();
                        let fingerprint = self.fingerprinter.fingerprint(&name, &args, &env)?;
                        self.claims[self.frames[index].claim].fingerprint =
                            Some(fingerprint.clone());
                        let hit = !self.force
                            && self.cache.as_ref().is_some_and(|cache| {
                                cache
                                    .manifest_matches(&self.root, &task, &fingerprint)
                                    .unwrap_or(false)
                            });
                        if hit {
                            self.reporter.cached(&name);
                            self.cached.push(name);
                            self.complete(index, 0);
                            progressed = true;
                            index += 1;
                            continue;
                        }
                    }
                    self.frames[index].dispatched = true;
                    let claim = self.frames[index].claim;
                    self.claims[claim].dispatched = true;
                    progressed = true;
                }
            }
            index += 1;
        }
        Ok(progressed)
    }

    pub fn step(&mut self) -> Result<Step> {
        // Advance until nothing else resolves without driver input.
        loop {
            if !self.advance()? {
                break;
            }
        }
        if self.frames[self.root_frame].done {
            return Ok(Step::Done(self.report(self.claim_code())));
        }
        if self.cancelled.is_cancelled() {
            return Ok(Step::Done(self.report(1)));
        }
        let mut ready = Vec::new();
        let root_name = self.frames[self.root_frame].name.clone();
        for frame in self.frames.iter_mut() {
            if frame.dispatched && !frame.surfaced && !frame.done {
                frame.surfaced = true;
                ready.push(ReadyTask {
                    name: frame.name.clone(),
                    action: self.graph.get(&frame.name).unwrap().action.clone(),
                    args: if frame.name == root_name {
                        self.extra_args.clone()
                    } else {
                        Vec::new()
                    },
                    env: frame.env.clone(),
                });
            }
        }
        Ok(Step::Ready(ready))
    }

    /// Records the outcome of a task handed out by `step`. On success the
    /// task's outputs are committed to the cache.
    pub fn finish(&mut self, name: &str, code: u8) {
        let Some(frame_id) = self
            .frames
            .iter()
            .position(|f| f.name == name && f.dispatched && !f.done)
        else {
            return;
        };
        if code == 0 && !self.cancelled.is_cancelled() {
            let fingerprint = self.claims[self.frames[frame_id].claim].fingerprint.clone();
            if let (Some(cache), Some(fingerprint)) = (&self.cache, fingerprint) {
                let task = self.graph.get(name).unwrap().clone();
                if let Err(error) =
                    cache.record_at(&self.root, &task, &fingerprint, &*self.reporter, name)
                {
                    self.reporter.cache_error(name, &error.to_string());
                }
            }
        }
        self.complete(frame_id, code);
    }

    fn claim_code(&self) -> u8 {
        let claim = self.frames[self.root_frame].claim;
        self.claims[claim].outcome.unwrap_or(1)
    }

    fn report(&self, code: u8) -> Report {
        let failed = self
            .claims
            .iter()
            .enumerate()
            .filter_map(|(index, claim)| {
                let code = claim.outcome?;
                if code == 0 || !claim.dispatched {
                    return None;
                }
                let name = self
                    .claim_by_name
                    .iter()
                    .find(|&(_, &claim_index)| claim_index == index)
                    .map(|(name, _)| name.clone())?;
                Some(FailedTask { name, code })
            })
            .collect();
        Report {
            code,
            completed: self.completed.clone(),
            failed,
            cached: self.cached.clone(),
        }
    }
}