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

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
    );
}

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
        }
    }
}

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

#[derive(Debug)]
struct Challenge {
    scheme: Range<usize>,
    params: Vec<AuthParam>,
}

fn parse_token(str: &Vec<char>, start: usize) -> Option<(Range<usize>, usize)> {
    let mut pos = start;
    while str.len() > pos {
        match str[pos] {
            'a'..'z' | 'A'..'Z' => pos += 1,
            _ => break,
        }
    }

    if pos == start {
        return None;
    }

    return Some((start..pos, pos));
}

fn parse_qdtext(str: &Vec<char>, start: usize) -> Option<usize> {
    match str[start] {
        '\u{0}'..'\u{25}' | '\u{127}' => return None,
        _ => return Some(start + 1),
    }
}

fn parse_quotedpair(str: &Vec<char>, start: usize) -> Option<usize> {
    if str[start] != '\\' {
        return None;
    }

    return Some(start + 2);
}

fn parse_quotedstring(str: &Vec<char>, start: usize) -> Option<(Range<usize>, usize)> {
    let mut pos = start;
    if str[pos] != '"' {
        return None;
    }
    pos += 1;

    loop {
        if let Some(npos) = parse_qdtext(str, pos) {
            pos = npos;
        } else if let Some(npos) = parse_quotedpair(str, pos) {
            pos = npos;
        } else {
            break;
        }
    }

    if str[pos] != '"' {
        return None;
    }
    pos += 1;

    return Some((start+1..pos-1, pos));
}

fn parse_sp(str: &Vec<char>, start: usize) -> Option<(Range<usize>, usize)> {
    if str.len() <= start {
        return None;
    }

    let mut pos = start;
    match str[pos] {
        ' ' => pos += 1,
        _ => return None,
    }

    if pos == start {
        return None;
    }

    return Some((start..pos, pos));
}

fn parse_crlf(str: &Vec<char>, start: usize) -> Option<usize> {
    if str.len() - 2 < start {
        return None;
    }

    if str[start] == '\r' && str[start+1] == '\n' {
        return Some(start + 2);
    }

    return None;
}

fn parse_lws(str: &Vec<char>, start: usize) -> Option<(Range<usize>, usize)> {
    let mut pos = start;

    match parse_crlf(str, pos) {
        Some(npos) => pos = npos,
        None => {},
    }

    match parse_sp(str, pos) {
        Some((_, npos)) => pos = npos,
        None => return None,
    }
    loop {
        match parse_sp(str, pos) {
            Some((_, npos)) => pos = npos,
            None => break,
        }
    }

    return Some((start..pos, pos));
}

#[derive(Debug)]
struct AuthParam {
    key: Range<usize>,
    value: Range<usize>,
}

fn parse_param(str: &Vec<char>, start: usize) -> Option<(AuthParam, usize)> {
    let pos = start;
    let (param_range, pos) = parse_token(str, pos)?;

    let mut pos = pos;
    if str[pos] == '=' {
        pos += 1;
    } else {
        todo!();
    }

    let value_range;
    if let Some((range, npos)) = parse_token(str, pos) {
        value_range = range;
        pos = npos;
    } else if let Some((range, npos)) = parse_quotedstring(str, pos) {
        value_range = range;
        pos = npos;
    } else {
        todo!();
    }

    return Some((AuthParam{key: param_range, value: value_range}, pos));
}

fn parse_challenge(str: &Vec<char>, start: usize) -> Option<(Challenge, usize)> {
    let pos = start;
    let (scheme_range, pos) = parse_token(str, pos).unwrap();

    let (_, pos) = parse_sp(str, pos).unwrap();
    let mut pos = pos;
    loop {
        match parse_sp(str, pos) {
            Some((_, npos)) => pos = npos,
            None => break,
        }
    }

    let mut params = vec!();

    let mut pos = pos;
    loop {
        match parse_lws(str, pos) {
            Some((_, npos)) => pos = npos,
            None => break,
        }
    }
    let (param, pos) = parse_param(str, pos).unwrap();
    params.push(param);
    let mut pos = pos;
    loop {
        loop {
            match parse_lws(str, pos) {
                Some((_, npos)) => pos = npos,
                None => break,
            }
        }

        if pos >= str.len() || str[pos] != ',' {
            break;
        }
        pos += 1;

        loop {
            match parse_lws(str, pos) {
                Some((_, npos)) => pos = npos,
                None => break,
            }
        }

        match parse_param(str, pos) {
            Some((param, npos)) => {
                pos = npos;
                params.push(param);
            },
            None => return None,
        }
    }

    if pos < str.len() {
        return None;
    }
    return Some((Challenge {scheme: scheme_range, params}, pos));
}

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 {
                dbg!(&header_str[param.key.clone()]);
                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;
    }
}

fn perform_registry_request(registry: &str, url: String, 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 {
            if let AuthStage::Authorized(x) = &state.stage {
                request = request.header("Authorization", x);
            }
        }

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

        if response.status() == 401 {
            if let Some(ref mut state) = state {
                if let AuthStage::Idle = state.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!("{}:{}", state.info.username, state.info.password)));
                                state.stage = AuthStage::Authorized(basic_auth);

                                continue; // try again
                            },
                            Some(AuthMethod::Bearer{realm, scope, service}) => {
                                let basic_auth = format!("Basic {}", BASE64_STANDARD.encode(format!("{}:{}", state.info.username, state.info.password)));

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

                                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();

                                state.stage = AuthStage::Authorized(format!("Bearer {}", token));
                                continue;
                            },
                            None => todo!(),
                        }
                    } else {
                        todo!("Server didn't ask us to authenticate");
                    }
                } else {
                    todo!("Authorized request somehow failed (expired token?)");
                }
            } 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,
}

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 tag = img.tag.map(|x| &file[x]).unwrap_or("latest");

    let digest = None;

    // Find the digest for the tag
    if let Some(digest) = img.digest {
        let url = format!("/v2/{}/manifests/{}", &file[img.image], tag);
        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");
    }

    let image_ref = {
        let prefix = &file[img.full_range.start..digest.start];
        let digest = &response.headers()["docker-content-digest"].to_str().unwrap();
        format!("{}{}", prefix, digest)
    };

    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);
}