use super::{common::PathField, Problem};
use crate::{belt_finding::brute_force::BruteforceBuilder, misc::Map};
use crate::{common::visualize::Visualize, prelude::*};
use std::{
    ops::RangeInclusive,
    time::{Duration, Instant},
};

#[derive(Default)]
struct Field {
    blocked: bool,
}

pub struct ConflictAvoidance {
    map: Map<Field>,
    belts: Vec<Vec<PathField>>,
    range: Option<(RangeInclusive<i32>, RangeInclusive<i32>)>,
}

#[derive(Debug, PartialEq, Eq)]
struct Candidate {
    min: Position,
    max: Position,
}

struct BruteForceEntry {
    path_id: usize,
    start: usize,
    end: usize,
    path: Vec<PathField>,
}

impl Candidate {
    fn new(min: Position, max: Position) -> Self {
        Self { min, max }
    }

    fn area(&self) -> PositionType {
        (self.max.x - self.min.x) * (self.max.y - self.min.y)
    }
    fn extend_range(&mut self, conflict_map: &Map<usize>) {
        for x in self.min.x..=self.max.x {
            if self.min.y > 0 && *conflict_map.get(x as usize, self.min.y as usize) > 1 {
                self.min.y -= 1;
                return self.extend_range(conflict_map);
            }
            if (self.max.y as usize) < conflict_map.height - 1
                && *conflict_map.get(x as usize, self.max.y as usize) > 1
            {
                self.max.y += 1;
                return self.extend_range(conflict_map);
            }
        }
        for y in self.min.y..=self.max.y {
            if self.min.x > 0 && *conflict_map.get(self.min.x as usize, y as usize) > 1 {
                self.min.x -= 1;
                return self.extend_range(conflict_map);
            }
            if (self.max.x as usize) < conflict_map.width - 1
                && *conflict_map.get(self.max.x as usize, y as usize) > 1
            {
                self.max.x += 1;
                return self.extend_range(conflict_map);
            }
        }
    }
}

impl ConflictAvoidance {
    pub fn new(problem: &Problem) -> Self {
        let mut map: Map<Field> = Map::new(problem.map.width, problem.map.height);
        for x in 0..problem.map.width {
            for y in 0..problem.map.height {
                map.get_mut(x, y).blocked = problem.map.get(x, y).blocked;
            }
        }
        let mut belts = Vec::new();
        for i in 0..problem.path.len() {
            // dbg!(&problem.path[i]);
            let mut p = Vec::new();

            p.push(PathField::Belt {
                pos: problem.start[i].0,
                dir: problem.start[i].1,
            });

            for (pos, dir) in &problem.path[i][1..] {
                let start = p.last().unwrap().end_pos();
                let start = start.0.in_direction(&start.1, 1);
                // dbg!(start, pos);
                if &start == pos {
                    p.push(PathField::Belt {
                        pos: *pos,
                        dir: *dir,
                    });
                } else {
                    p.push(PathField::Underground {
                        pos: start,
                        dir: *dir,
                        len: u8::max(
                            (start.x - pos.x).unsigned_abs() as u8,
                            (start.y - pos.y).unsigned_abs() as u8,
                        ),
                    })
                }
            }

            // p.push(PathField::Belt {
            //     pos: *problem.path[i].last().unwrap(),
            //     dir: problem.end[i].1,
            // });

            // p.push(PathField::Belt {
            //     pos: problem.end[i].0,
            //     dir: problem.end[i].1,
            // });

            belts.push(p);
        }

        Self {
            map,
            belts,
            range: None,
        }
    }

    pub fn get_paths(&self) -> &[Vec<PathField>] {
        &self.belts
    }

    fn try_bruteforce(
        &self,
        candidate: &Candidate,
        timeout: Option<Instant>,
    ) -> Option<Vec<BruteForceEntry>> {
        let xrange = candidate.min.x as usize..=candidate.max.x as usize;
        let yrange = candidate.min.y as usize..=candidate.max.y as usize;

        // println!("x: {:?}, y: {:?}", xrange, yrange);

        let xsize = xrange.end() - xrange.start() + 1;
        let ysize = yrange.end() - yrange.start() + 1;

        // dbg!(xsize, ysize);

        let mut b = BruteforceBuilder::new(xsize + 2, ysize + 2);
        b.set_blocked_range(0, 0, xsize + 1, 0, true);
        b.set_blocked_range(0, ysize + 1, xsize + 1, ysize + 1, true);
        b.set_blocked_range(0, 0, 0, ysize + 1, true);
        b.set_blocked_range(xsize + 1, 0, xsize + 1, ysize + 1, true);

        for x in 0..xsize {
            for y in 0..ysize {
                b.set_blocked(
                    x + 1,
                    y + 1,
                    self.map.get(xrange.start() + x, yrange.start() + y).blocked,
                );
            }
        }

        let mut mapping = Vec::new();
        let offset = Position::new(*xrange.start() as i32 - 1, *yrange.start() as i32 - 1);

        // dbg!(&xrange, &yrange);

        for (i, path) in self.belts.iter().enumerate() {
            // index of first PathField where the next position is in the area
            let start_index = path
                .iter()
                .map(|p| {
                    let (pos, dir) = p.end_pos();
                    pos.in_direction(&dir, 1)
                })
                .position(|pos| {
                    xrange.contains(&(pos.x as usize)) && yrange.contains(&(pos.y as usize))
                });

            // index of last PathField that ends inside the area
            let end_index = path
                .iter()
                .rev()
                .map(|p| p.end_pos().0)
                .position(|pos| {
                    xrange.contains(&(pos.x as usize)) && yrange.contains(&(pos.y as usize))
                })
                .map(|rev_i| path.len() - rev_i - 1);

            if let Some((start_index, end_index)) = Option::zip(start_index, end_index) {
                // dbg!(start_index, end_index, path[start_index], path[end_index]);

                let (start_pos, start_dir) = path[start_index].end_pos();

                let (end_pos, end_dir) = path[end_index].end_pos();

                // edge cases with underground into and end

                if let Some(PathField::Underground { pos, dir, len }) = path.get(end_index + 1) {
                    if xrange.contains(&(pos.x as usize)) && yrange.contains(&(pos.y as usize)) {
                        let p = *pos - offset;
                        // println!("Blocked {:?}", p);
                        if b.get_blocked(p.x as usize, p.y as usize) {
                            return None;
                        }
                        b.set_blocked(p.x as usize, p.y as usize, true);
                    }

                    for l in 1..*len {
                        let p = pos.in_direction(dir, l as i32);

                        if xrange.contains(&(p.x as usize)) && yrange.contains(&(p.y as usize)) {
                            let p = p - offset;
                            b.set_underground(p.x as usize, p.y as usize, Some(*dir));
                        }
                    }
                }

                if xrange.contains(&(start_pos.x as usize))
                    && yrange.contains(&(start_pos.y as usize))
                {
                    let p = start_pos - offset;
                    // println!("Blocked {:?}", p);
                    // if b.get_blocked(p.x as usize, p.y as usize) {
                    //     return None;
                    // }
                    b.set_blocked(p.x as usize, p.y as usize, true);
                }

                if let PathField::Underground { pos, dir, len } = path[start_index] {
                    for l in 1..len {
                        let p = pos.in_direction(&dir, l as i32);

                        if xrange.contains(&(p.x as usize)) && yrange.contains(&(p.y as usize)) {
                            let p = p - offset;
                            b.set_underground(p.x as usize, p.y as usize, Some(dir));
                        }
                    }
                }

                let start_pos_offset = start_pos - offset;
                b.add_path((start_pos_offset, start_dir), (end_pos - offset, end_dir));

                mapping.push((i, start_index, end_index));
            } else if let Some(end_index) = end_index {
                if let Some(PathField::Underground { pos, dir, len }) = path.get(end_index) {
                    let p = pos.in_direction(dir, *len as i32);
                    let p = p - offset;
                    b.set_blocked(p.x as usize, p.y as usize, true);
                }
            }
        }

        // println!("{:?}", mapping);

        let mut b = b.build();

        // b.print();
        // self.print();

        let mut min_cost = f64::INFINITY;
        let mut solutions = Vec::new();

        while b.next_finish_state(timeout) {
            // println!("{}", b);
            // b.print();
            let c = b.cost();
            if c < min_cost {
                min_cost = c;
                solutions = b.get_paths();
            }
        }

        if b.solution_count() == 0 {
            return None;
        }

        Some(
            mapping
                .into_iter()
                .zip(solutions)
                .map(|((path_id, start, end), path)| BruteForceEntry {
                    path_id,
                    start,
                    end,
                    path,
                })
                .collect(),
        )
    }

    pub fn remove_conflict(&mut self, timeout: Option<Instant>) -> bool {
        let mut conflicts: Map<usize> = Map::new(self.map.width, self.map.height);

        for x in 0..self.map.width {
            for y in 0..self.map.height {
                if self.map.get(x, y).blocked {
                    *conflicts.get_mut(x, y) += 1;
                }
            }
        }

        for path in &self.belts {
            for p in &path[1..] {
                match p {
                    PathField::Belt { pos, dir: _ } => {
                        *conflicts.get_mut(pos.x as usize, pos.y as usize) += 1
                    }
                    PathField::Underground { pos, dir, len } => {
                        *conflicts.get_mut(pos.x as usize, pos.y as usize) += 1;
                        let end = pos.in_direction(dir, *len as PositionType);
                        *conflicts.get_mut(end.x as usize, end.y as usize) += 1;
                    }
                }
            }
        }

        // for y in 0..self.map.height {
        //     for x in 0..self.map.width {
        //         if *conflicts.get(x, y) > 1 {
        //             print!("#");
        //         } else {
        //             print!(" ");
        //         }
        //     }
        //     println!();
        // }

        // self.print();

        let mut candidates = Vec::new();

        for y in 0..self.map.height {
            for x in 0..self.map.width {
                if *conflicts.get(x, y) > 1 {
                    let mut candidate = Candidate::new(
                        Position::new(x as PositionType, y as PositionType),
                        Position::new(x as PositionType, y as PositionType),
                    );
                    candidate.extend_range(&conflicts);
                    if !candidates.iter().any(|c| c == &candidate) {
                        candidates.push(candidate);
                    }
                }
            }
        }

        if candidates.is_empty() {
            return false;
        }
        // dbg!(&candidates);

        while timeout.is_none_or(|t| t > Instant::now()) {
            candidates.sort_by_key(|c| -c.area());
            // dbg!(&candidates);
            let c = match candidates.pop() {
                Some(c) => c,
                None => {
                    return false;
                }
            };

            self.range = Some((c.min.x..=c.max.x, c.min.y..=c.max.y));

            let result = self.try_bruteforce(&c, timeout);

            // dbg!(&solutions);

            // println!("{}", b.solution_count());

            let xrange = c.min.x as usize..=c.max.x as usize;
            let yrange = c.min.y as usize..=c.max.y as usize;
            let offset = Position::new(*xrange.start() as i32 - 1, *yrange.start() as i32 - 1);
            if let Some(r) = result {
                // println!("Here");
                for BruteForceEntry {
                    path_id,
                    start,
                    end,
                    path,
                } in r
                {
                    let mut t = Vec::new();
                    // println!("{:?}", p);

                    t.extend_from_slice(&self.belts[path_id][0..=start]);
                    t.extend(path[1..].iter().map(|p| p.offset(&offset)));
                    t.extend_from_slice(&self.belts[path_id][end + 1..]);
                    // println!("{:?}", &t);
                    // println!();
                    // println!("{:?}", &self.belts[index]);

                    self.belts[path_id] = t;
                }
                return true;
            }
            let mut candidate = Candidate::new(
                Position::new(
                    xrange.start().saturating_sub(1) as PositionType,
                    *yrange.start() as PositionType,
                ),
                Position::new(*xrange.end() as PositionType, *yrange.end() as PositionType),
            );
            candidate.extend_range(&conflicts);
            if candidate != c && !candidates.iter().any(|c| c == &candidate) {
                candidates.push(candidate);
            }
            let mut candidate = Candidate::new(
                Position::new(
                    *xrange.start() as PositionType,
                    yrange.start().saturating_sub(1) as PositionType,
                ),
                Position::new(*xrange.end() as PositionType, *yrange.end() as PositionType),
            );
            candidate.extend_range(&conflicts);
            if candidate != c && !candidates.iter().any(|c| c == &candidate) {
                candidates.push(candidate);
            }
            let mut candidate = Candidate::new(
                Position::new(
                    *xrange.start() as PositionType,
                    *yrange.start() as PositionType,
                ),
                Position::new(
                    usize::min(xrange.end() + 1, self.map.width - 1) as PositionType,
                    *yrange.end() as PositionType,
                ),
            );
            candidate.extend_range(&conflicts);
            if candidate != c && !candidates.iter().any(|c| c == &candidate) {
                candidates.push(candidate);
            }
            let mut candidate = Candidate::new(
                Position::new(
                    *xrange.start() as PositionType,
                    *yrange.start() as PositionType,
                ),
                Position::new(
                    *xrange.end() as PositionType,
                    usize::min(yrange.end() + 1, self.map.height - 1) as PositionType,
                ),
            );
            candidate.extend_range(&conflicts);
            if candidate != c && !candidates.iter().any(|c| c == &candidate) {
                candidates.push(candidate);
            }
        }

        false
    }

    pub fn remove_all_conflicts(&mut self, timeout: Option<Duration>) -> bool {
        let end = timeout.map(|t| std::time::Instant::now() + t);
        while self.remove_conflict(end) {}

        let mut conflicts: Map<bool> = Map::new(self.map.width, self.map.height);

        for x in 0..self.map.width {
            for y in 0..self.map.height {
                if self.map.get(x, y).blocked {
                    if conflicts.get(x, y) == &true {
                        return false;
                    } else {
                        conflicts.set(x, y, true);
                    }
                }
            }
        }

        for path in &self.belts {
            for p in &path[1..] {
                match p {
                    PathField::Belt { pos, dir: _ } => {
                        if conflicts.get(pos.x as usize, pos.y as usize) == &true {
                            return false;
                        } else {
                            conflicts.set(pos.x as usize, pos.y as usize, true);
                        }
                    }
                    PathField::Underground { pos, dir, len } => {
                        if conflicts.get(pos.x as usize, pos.y as usize) == &true {
                            return false;
                        } else {
                            conflicts.set(pos.x as usize, pos.y as usize, true);
                        }
                        let end = pos.in_direction(dir, *len as PositionType);
                        if conflicts.get(end.x as usize, end.y as usize) == &true {
                            return false;
                        } else {
                            conflicts.set(end.x as usize, end.y as usize, true);
                        }
                    }
                }
            }
        }
        true
    }
}

impl Visualize for ConflictAvoidance {
    fn visualize(&self) -> crate::common::visualize::Visualization {
        let mut v = crate::common::visualize::Visualization::new(Position::new(
            self.map.width as i32,
            self.map.height as i32,
        ));

        // create conflict map
        let mut conflicts: Map<usize> = Map::new(self.map.width, self.map.height);

        for x in 0..self.map.width {
            for y in 0..self.map.height {
                if self.map.get(x, y).blocked {
                    *conflicts.get_mut(x, y) += 1;

                    v.add_symbol(
                        Position::new(x as i32, y as i32),
                        crate::common::visualize::Symbol::Block,
                        Some(crate::common::visualize::Color::white()),
                        None,
                    );
                }
            }
        }

        for path in &self.belts {
            for p in &path[1..] {
                match p {
                    PathField::Belt { pos, dir: _ } => {
                        *conflicts.get_mut(pos.x as usize, pos.y as usize) += 1
                    }
                    PathField::Underground { pos, dir, len } => {
                        *conflicts.get_mut(pos.x as usize, pos.y as usize) += 1;
                        let end = pos.in_direction(dir, *len as PositionType);
                        *conflicts.get_mut(end.x as usize, end.y as usize) += 1;
                    }
                }
            }
        }

        for (i, path) in self.belts.iter().enumerate() {
            for p in path {
                match p {
                    PathField::Belt { pos, dir } => {
                        v.add_symbol(
                            *pos,
                            crate::common::visualize::Symbol::Arrow(*dir),
                            Some(crate::common::visualize::Color::index(i)),
                            None,
                        );
                    }
                    PathField::Underground { pos, dir, len } => {
                        v.add_symbol(
                            *pos,
                            crate::common::visualize::Symbol::ArrowEnter(*dir),
                            Some(crate::common::visualize::Color::index(i)),
                            None,
                        );
                        v.add_symbol(
                            pos.in_direction(dir, *len as i32),
                            crate::common::visualize::Symbol::ArrowExit(*dir),
                            Some(crate::common::visualize::Color::index(i)),
                            None,
                        );
                    }
                }
            }
        }

        for x in 0..self.map.width {
            for y in 0..self.map.height {
                if conflicts.get(x, y) > &1 {
                    v.overwrite_background(
                        Position::new(x as i32, y as i32),
                        Some(crate::common::visualize::Color::new(100, 80, 80)),
                    );
                }
            }
        }

        for path in &self.belts {
            for p in &path[1..] {
                match p {
                    PathField::Belt { pos, dir: _ } => {
                        *conflicts.get_mut(pos.x as usize, pos.y as usize) += 1
                    }
                    PathField::Underground { pos, dir, len } => {
                        *conflicts.get_mut(pos.x as usize, pos.y as usize) += 1;
                        let end = pos.in_direction(dir, *len as PositionType);
                        *conflicts.get_mut(end.x as usize, end.y as usize) += 1;
                    }
                }
            }
        }

        v
    }
}