mod parser;
use crate::parser::*;

use base64::prelude::*;
use yaml_rust2::parser::Parser;
use yaml_rust2::Event;
use std::collections::HashMap;
use std::ops::Range;
use std::sync;

fn help(cmd: &str) {
    println!(
        "{} [options] [--] <FILE>
Search FILE for docker images and suggest updates

Options:
--auth <REGISTRY> <USER> <PASS> Authenticate against REGISTRY (repeatable)",
        cmd
    );
}

#[derive(Debug)]
enum AuthMethod {
    Basic,
    Bearer{realm: String, service: String, scope: String},
}

impl AuthMethod {
    fn from_header(header: &ureq::http::HeaderValue) -> Option<Self> {
        let header_str = header.to_str().unwrap();
        let (parse, _) = parse_challenge(&header_str.chars().collect::<Vec<char>>(), 0).unwrap();

        if &header_str[parse.scheme.clone()] == "Basic" {
            return Some(AuthMethod::Basic);
        }

        if &header_str[parse.scheme.clone()] == "Bearer" {
            let mut realm_param = None;
            let mut scope_param = None;
            let mut service_param = None;
            for param in parse.params {
                match &header_str[param.key.clone()] {
                    "realm" => realm_param = Some(param),
                    "scope" => scope_param = Some(param),
                    "service" => service_param = Some(param),
                    _ => {},
                }
            }

            return Some(AuthMethod::Bearer{
                realm: header_str[realm_param.unwrap().value].to_string(),
                scope: header_str[scope_param.unwrap().value].to_string(),
                service: header_str[service_param.unwrap().value].to_string(),
            });
        }

        return None;
    }
}

struct AuthInfo {
    host: String,
    username: String,
    password: String,
}

enum AuthStage {
    Idle,
    Unauthorized,
    Authorized(String),
}

struct AuthState {
    info: AuthInfo,
    stage: AuthStage,
}

struct Auth {
    states: HashMap<String, AuthState>
}

impl Auth {
    fn new(infos: Vec<AuthInfo>) -> Self {
        let mut states = HashMap::new();
        for info in infos {
            states.insert(info.host.clone(), AuthState {
                info: info,
                stage: AuthStage::Idle,
            });
        }

        return Auth {
            states
        }
    }
}

impl AuthState{
    fn add_to_request<T>(&self, req: ureq::RequestBuilder<T>) -> ureq::RequestBuilder<T> {
        if let AuthStage::Authorized(x) = &self.stage {
            return req.header("Authorization", x);
        }

        return req
    }

    fn attempt_authorization(&mut self, response: &ureq::http::Response<ureq::Body>) -> Result<(), ()> {
        if let AuthStage::Idle = self.stage {
            if let Some(auth_header) = response.headers().get("www-authenticate") {
                let auth_header = AuthMethod::from_header(auth_header);
                match auth_header {
                    Some(AuthMethod::Basic) => {
                        let basic_auth = format!("Basic {}", BASE64_STANDARD.encode(format!("{}:{}", self.info.username, self.info.password)));
                        self.stage = AuthStage::Authorized(basic_auth);

                        return Ok(());
                    },
                    Some(AuthMethod::Bearer{realm, scope, service}) => {
                        let basic_auth = format!("Basic {}", BASE64_STANDARD.encode(format!("{}:{}", self.info.username, self.info.password)));

                        let url = format!("{}?service={}&scope={}", realm, service, scope);

                        let body = ureq::get(url)
                            .config().http_status_as_error(false).build()
                            .header("Authorization", basic_auth)
                            .call();

                        dbg!(&body);

                        let body = body.unwrap()
                            .body_mut().read_to_string().unwrap();

                        dbg!(&body);

                        let body: tinyjson::JsonValue = body
                            .parse().unwrap();

                        let token: &String = body["token"].get().unwrap();

                        self.stage = AuthStage::Authorized(format!("Bearer {}", token));
                        return Ok(());
                    },
                    None => todo!("Failed parsing the challenge header"),
                }
            } else {
                todo!("Server didn't ask us to authenticate");
            }
        } else {
            todo!("Authorized request somehow failed (expired token?)");
        }
    }
}

fn perform_registry_request(registry: &str, url: &str, auth: &mut Auth) -> Result<ureq::http::Response<ureq::Body>, ()> {
    let mut state = auth.states.get_mut(registry);

    let url = format!("https://{}{}", registry, &url);

    for _ in 0..2 {
        let mut request = ureq::get(&url)
            .config().http_status_as_error(false).build();

        if let Some(ref state) = state {
            request = state.add_to_request(request);
        }

        let response = request.call().unwrap();

        if response.status() == 401 {
            if let Some(ref mut state) = state {
                state.attempt_authorization(&response)?;
                continue;
            } else {
                todo!("Server returned 401 but we have no credentials");
            }
        }

        return Ok(response);
    }

    return Err(());
}

enum YContext {
    InDocument,
    InObject,
    InSequence,
    InValue(bool),
}

#[derive(Debug)]
struct Chunk {
    position: Range<usize>,
}

fn scan_yaml_for_images<T: Iterator<Item = char>>(mut yaml: Parser<T>) -> Vec<Chunk> {
    let mut images = vec!();

    let mut scope = vec!();

    loop {
        let (ev, mark) = yaml.next_token().unwrap();
        match ev {
            Event::StreamStart => {}
            Event::StreamEnd => { break; }
            Event::DocumentStart => { scope.push(YContext::InDocument); }
            Event::DocumentEnd => {
                assert!(matches!(scope.pop().unwrap(), YContext::InDocument));
                scope.pop_if(|x| matches!(x, YContext::InValue(_)));
            },
            Event::MappingStart(_, _) => { scope.push(YContext::InObject); },
            Event::MappingEnd => {
                assert!(matches!(scope.pop().unwrap(), YContext::InObject));
                scope.pop_if(|x| matches!(x, YContext::InValue(_)));
            },
            Event::SequenceStart(_, _) => { scope.push(YContext::InSequence); },
            Event::SequenceEnd => { 
                assert!(matches!(scope.pop().unwrap(), YContext::InSequence));
                scope.pop_if(|x| matches!(x, YContext::InValue(_)));
            },

            Event::Scalar(ref txt, _, _, _) => {
                let parent = scope.last().unwrap();
                match parent {
                    YContext::InObject => {
                        // We are the key of a mapping, which means the next even is the value
                        scope.push(YContext::InValue(txt == "image"));
                    },
                    YContext::InSequence => {},
                    YContext::InValue(img) => {
                        if *img {
                            let next_idx = images.len();
                            images.push(Chunk{
                                position: mark.index()..mark.index() + txt.len(),
                            });
                        }
                        scope.pop();
                    },

                    _ => panic!(),
                }
            },
            x => todo!("{:?}", x),
        }
    }

    return images;
}

pub trait SubsliceOffset<T> {
    fn subslice_range(&self, inner: &Self) -> Option<std::ops::Range<usize>>;
}

impl<T> SubsliceOffset<T> for [T] {
    fn subslice_range(&self, subslice: &[T]) -> Option<std::ops::Range<usize>> {
        if size_of::<T>() == 0 {
            panic!("elements are zero-sized");
        }

        let self_start = self.as_ptr().addr();
        let subslice_start = subslice.as_ptr().addr();

        let byte_start = subslice_start.wrapping_sub(self_start);

        if !byte_start.is_multiple_of(size_of::<T>()) {
            return None;
        }

        let start = byte_start / size_of::<T>();
        let end = start.wrapping_add(subslice.len());

        if start <= self.len() && end <= self.len() { Some(start..end) } else { None }
    }
}

#[derive(Debug, Clone)]
struct DockerRef {
    full_range: Range<usize>,

    registry: Option<Range<usize>>,
    image: Range<usize>,
    tag: Option<Range<usize>>,
    digest: Option<Range<usize>>,
}

impl DockerRef{
    fn parse(file: &str, chunk: &Chunk) -> DockerRef {
        let mut string_range = chunk.position.clone();

        let mut digest = None;
        if let Some(idx) = file[string_range.clone()].rfind("@") {
            digest = Some(string_range.start+idx+1..string_range.end);
            string_range.end = string_range.start+idx;
        }

        let mut tag = None;
        if let Some(idx) = file[string_range.clone()].rfind(":") {
            tag = Some(string_range.start+idx+1..string_range.end);
            string_range.end = string_range.start+idx;
        }

        let mut registry = None;
        let image;
        if let Some(idx) = file[string_range.clone()].find("/") {
            let head = &file[string_range.clone()][..idx];
            if head.contains(":") || head.contains(".") {
                registry = Some(string_range.start..string_range.start+idx);
                image = string_range.start+idx+1..string_range.end;
            } else {
                registry = None;
                image = string_range;
            }
        } else {
            image = string_range;
        }

        return DockerRef {
            full_range: chunk.position.clone(),

            registry,
            image,
            tag,
            digest,
        };
    }
}

#[derive(Debug)]
struct Update {
    position: Range<usize>,
    content: String,
}

#[derive(Debug)]
struct Version {
    // It always starts with a string part, so the first string can be empty
    string_parts: Vec<String>,
    number_parts: Vec<u64>,
}

#[derive(Debug)]
enum CompareOutcome {
    Higher,
    Lower,

    Incompatible,
    Identical,
}

impl Version {
    fn parse(tag: &str) -> Self {
        static RE: sync::LazyLock<regex::Regex> = sync::LazyLock::new(|| regex::Regex::new(r"(?<str>[^0-9]*)(?<num>[0-9]+)").unwrap());
        let mut string_parts = vec!();
        let mut number_parts = vec!();
        for it in RE.captures_iter(tag) {
            string_parts.push(it.name("str").unwrap().as_str().to_string());
            number_parts.push(it.name("num").unwrap().as_str().to_string().parse().unwrap());
        }

        return Version {
            string_parts,
            number_parts,
        }
    }

    fn compare(&self, other: &Self) -> CompareOutcome {
        assert!(self.number_parts.len() == self.string_parts.len());
        assert!(other.number_parts.len() == other.string_parts.len());

        if self.number_parts.len() != other.number_parts.len() {
            return CompareOutcome::Incompatible;
        }

        for (self_string, other_string) in self.string_parts.iter().zip(other.string_parts.iter()) {
            if self_string != other_string {
                return CompareOutcome::Incompatible;
            }
        }

        for (self_number, other_number) in self.number_parts.iter().zip(other.number_parts.iter()) {
            if self_number > other_number {
                return CompareOutcome::Lower;
            } else if self_number < other_number {
                return CompareOutcome::Higher;
            }
        }

        return CompareOutcome::Identical;
    }
}

fn fetch_new_image(file: &str, auth: &mut Auth, img: DockerRef, edits: &mut Vec<Update>) {
    let registry = img.registry.map(|x| &file[x]).unwrap_or("registry.jnsn.dev/");
    let mut tag = img.tag.map(|x| file[x].to_string());

    let mut digest = None;


    if let Some(ref tag_str) = tag {
        let mut current = Version::parse(&tag_str);

        let mut url = format!("/v2/{}/tags/list", &file[img.image.clone()]);
        loop {
            let mut response = perform_registry_request(registry, &url, auth).unwrap();

            let link_str = response.headers()["link"].to_str().unwrap();
            let (link, _) = parse_link(&link_str.chars().collect(), 0).unwrap();
            dbg!(&link_str[link.get(0).unwrap().params.get(0).unwrap().value().clone()]);

            let body: tinyjson::JsonValue = response.body_mut().read_to_string().unwrap().parse().unwrap();
            for it in body["tags"].get::<Vec<tinyjson::JsonValue>>().unwrap().iter() {
                let candidate_str = it.get::<String>().unwrap();
                let candidate = Version::parse(candidate_str);
                match current.compare(&candidate) {
                    CompareOutcome::Higher => {
                        tag = Some(candidate_str.clone());
                        current = candidate;
                    },
                    CompareOutcome::Lower => {},
                    CompareOutcome::Incompatible => {},
                    CompareOutcome::Identical => {},
                }
            }
            dbg!(&tag);
        }
    }

    if img.digest.is_some() {
        // Find the digest for the selected tag
        let url = format!("/v2/{}/manifests/{}", &file[img.image], tag.unwrap_or("latest".to_string()));
        let mut response = perform_registry_request(registry, &url, auth).unwrap();

        let body: tinyjson::JsonValue = response.body_mut().read_to_string().unwrap().parse().unwrap();
        dbg!(&body);

        // let media_type : &String = body["mediaType"].get().unwrap();
        // assert!(media_type == "application/vnd.docker.distribution.manifest.v2+json");

        digest = Some(response.headers()["docker-content-digest"].to_str().unwrap().to_string());
    }

    let image_ref = {
        let prefix = &file[img.full_range.start..img.digest.unwrap().start];
        format!("{}{}", prefix, digest.unwrap_or("".to_string()))
    };

    edits.push(Update{
        position: img.full_range,
        content: image_ref,
    });
}

fn main() {
    let argv: Vec<String> = std::env::args().collect();
    let mut it = argv.iter();
    let cmd = &it.next().unwrap();

    let mut auths = vec![];
    let mut positional: Vec<&str> = vec!();

    loop {
        match it.next().map(|x| x.as_str()) {
            None => break,
            Some("--auth") => {
                if let Some(registry) = it.next() && let Some(username) = it.next() && let Some(password) = it.next() {
                    auths.push(AuthInfo {
                        host: registry.clone(),
                        username: username.clone(),
                        password: password.clone(),
                    });
                } else {
                    println!("Error: --auth requires three parameters");
                    help(cmd);
                    std::process::exit(1);
                }
            },
            Some("-h") | Some("--help") => {
                help(cmd);
                std::process::exit(0);
            },
            Some(arg) => positional.push(arg),
        };
    }


    if positional.len() != 1 {
        panic!("Bad arguments");
    }
    let file = positional[0];

    let mut auth = Auth::new(auths);

    let file_content = &std::fs::read_to_string(file).unwrap();
    let yaml = Parser::new_from_str(&file_content);
    let images : Vec<_> = scan_yaml_for_images(yaml);

    let mut edits = vec![];

    let images : Vec<_> = images.iter()
        .map(|x| DockerRef::parse(&file_content, x))
        .map(|x| fetch_new_image(&file_content, &mut auth, x, &mut edits))
        .collect();

    dbg!(&images);
    dbg!(&edits);
    // dbg!(&file_content[images[0].digest.as_ref().unwrap().clone()]);

    // let body: String = auth.apply(ureq::get("https://registry.jnsn.dev/v2/autobrr/tags/list"))
    //     .call().unwrap()
    //     .body_mut()
    //     .read_to_string().unwrap();
    // dbg!(body);
}