use serde_json::{json, Value};
use std::collections::HashMap;
use std::io::{self, BufRead, Write};

fn main() {
    if let Err(error) = run() {
        eprintln!("iris-language-server: {error}");
        std::process::exit(1);
    }
}

fn run() -> Result<(), Box<dyn std::error::Error>> {
    let stdin = io::stdin();
    let stdout = io::stdout();
    let mut reader = io::BufReader::new(stdin.lock());
    let mut writer = io::BufWriter::new(stdout.lock());
    let mut documents: HashMap<String, String> = HashMap::new();
    let mut shutdown_requested = false;

    while let Some(message) = read_message(&mut reader)? {
        let id = message.get("id").cloned();
        let method = message.get("method").and_then(Value::as_str);

        match method {
            Some("initialize") => {
                let result = json!({
                    "capabilities": {
                        "textDocumentSync": 1,
                        "completionProvider": {
                            "resolveProvider": false,
                            "triggerCharacters": [".", ":"]
                        },
                        "hoverProvider": true,
                        "documentSymbolProvider": true
                    },
                    "serverInfo": {
                        "name": "Iris Language Server",
                        "version": "0.1.0"
                    }
                });
                send_response(&mut writer, id, result)?;
            }
            Some("initialized") => {}
            Some("shutdown") => {
                shutdown_requested = true;
                send_response(&mut writer, id, Value::Null)?;
            }
            Some("exit") => {
                std::process::exit(if shutdown_requested { 0 } else { 1 });
            }
            Some("textDocument/didOpen") => {
                if let Some(document) = message.get("params").and_then(|p| p.get("textDocument")) {
                    let uri = document.get("uri").and_then(Value::as_str).unwrap_or_default().to_string();
                    let text = document.get("text").and_then(Value::as_str).unwrap_or_default().to_string();
                    documents.insert(uri.clone(), text.clone());
                    publish_diagnostics(&mut writer, &uri, &text)?;
                }
            }
            Some("textDocument/didChange") => {
                if let Some(params) = message.get("params") {
                    let uri = params
                        .get("textDocument")
                        .and_then(|v| v.get("uri"))
                        .and_then(Value::as_str)
                        .unwrap_or_default()
                        .to_string();
                    let text = params
                        .get("contentChanges")
                        .and_then(Value::as_array)
                        .and_then(|changes| changes.last())
                        .and_then(|change| change.get("text"))
                        .and_then(Value::as_str)
                        .unwrap_or_default()
                        .to_string();
                    documents.insert(uri.clone(), text.clone());
                    publish_diagnostics(&mut writer, &uri, &text)?;
                }
            }
            Some("textDocument/didClose") => {
                if let Some(uri) = message
                    .get("params")
                    .and_then(|p| p.get("textDocument"))
                    .and_then(|d| d.get("uri"))
                    .and_then(Value::as_str)
                {
                    documents.remove(uri);
                    notify(
                        &mut writer,
                        "textDocument/publishDiagnostics",
                        json!({ "uri": uri, "diagnostics": [] }),
                    )?;
                }
            }
            Some("textDocument/completion") => {
                let items: Vec<Value> = completions()
                    .into_iter()
                    .map(|label| json!({"label": label, "kind": 14, "detail": "Iris keyword/builtin"}))
                    .collect();
                send_response(&mut writer, id, json!(items))?;
            }
            Some("textDocument/hover") => {
                let uri = message
                    .get("params")
                    .and_then(|p| p.get("textDocument"))
                    .and_then(|d| d.get("uri"))
                    .and_then(Value::as_str)
                    .unwrap_or_default();
                let position = message.get("params").and_then(|p| p.get("position")).cloned().unwrap_or(Value::Null);
                let line = position.get("line").and_then(Value::as_u64).unwrap_or(0) as usize;
                let character = position.get("character").and_then(Value::as_u64).unwrap_or(0) as usize;
                let text = documents.get(uri).cloned().unwrap_or_default();
                let word = word_at(&text, line, character);
                if let Some(doc) = keyword_doc(&word) {
                    send_response(
                        &mut writer,
                        id,
                        json!({
                            "contents": {
                                "kind": "markdown",
                                "value": format!("**{}**\n\n{}", word, doc)
                            }
                        }),
                    )?;
                } else {
                    send_response(&mut writer, id, Value::Null)?;
                }
            }
            Some("textDocument/documentSymbol") => {
                let uri = message
                    .get("params")
                    .and_then(|p| p.get("textDocument"))
                    .and_then(|d| d.get("uri"))
                    .and_then(Value::as_str)
                    .unwrap_or_default();
                let text = documents.get(uri).cloned().unwrap_or_default();
                let symbols = collect_symbols(&text);
                send_response(&mut writer, id, json!(symbols))?;
            }
            _ => {
                if id.is_some() {
                    send_response(&mut writer, id, Value::Null)?;
                }
            }
        }
    }

    Ok(())
}

fn read_message<R: BufRead>(reader: &mut R) -> Result<Option<Value>, Box<dyn std::error::Error>> {
    let mut content_length: Option<usize> = None;

    loop {
        let mut line = String::new();
        let bytes = reader.read_line(&mut line)?;
        if bytes == 0 {
            return Ok(None);
        }

        if line == "\r\n" || line == "\n" {
            break;
        }

        if let Some(rest) = line.strip_prefix("Content-Length:") {
            content_length = Some(rest.trim().parse()?);
        }
    }

    let length = match content_length {
        Some(value) => value,
        None => return Ok(None),
    };

    let mut body = vec![0_u8; length];
    reader.read_exact(&mut body)?;
    let message = serde_json::from_slice::<Value>(&body)?;
    Ok(Some(message))
}

fn send_response<W: Write>(
    writer: &mut W,
    id: Option<Value>,
    result: Value,
) -> Result<(), Box<dyn std::error::Error>> {
    let payload = json!({
        "jsonrpc": "2.0",
        "id": id.unwrap_or(Value::Null),
        "result": result
    });
    write_message(writer, &payload)
}

fn notify<W: Write>(
    writer: &mut W,
    method: &str,
    params: Value,
) -> Result<(), Box<dyn std::error::Error>> {
    let payload = json!({
        "jsonrpc": "2.0",
        "method": method,
        "params": params
    });
    write_message(writer, &payload)
}

fn write_message<W: Write>(writer: &mut W, payload: &Value) -> Result<(), Box<dyn std::error::Error>> {
    let body = serde_json::to_vec(payload)?;
    write!(writer, "Content-Length: {}\r\n\r\n", body.len())?;
    writer.write_all(&body)?;
    writer.flush()?;
    Ok(())
}

fn publish_diagnostics<W: Write>(
    writer: &mut W,
    uri: &str,
    text: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut diagnostics = Vec::new();
    diagnostics.extend(check_unbalanced_pairs(text));
    diagnostics.extend(check_trailing_comma_after_block(text));
    diagnostics.extend(check_missing_semicolons(text));

    notify(
        writer,
        "textDocument/publishDiagnostics",
        json!({
            "uri": uri,
            "diagnostics": diagnostics
        }),
    )
}

fn check_unbalanced_pairs(text: &str) -> Vec<Value> {
    let pairs = [('{', '}'), ('(', ')'), ('[', ']')];
    let mut diagnostics = Vec::new();

    for (open, close) in pairs {
        let mut balance = 0_i32;
        for ch in text.chars() {
            if ch == open {
                balance += 1;
            } else if ch == close {
                balance -= 1;
            }
        }
        if balance != 0 {
            diagnostics.push(diagnostic(
                0,
                0,
                0,
                1,
                &format!("Unbalanced '{}{}' pair detected.", open, close),
                1,
            ));
        }
    }

    diagnostics
}

fn check_trailing_comma_after_block(text: &str) -> Vec<Value> {
    text.lines()
        .enumerate()
        .filter_map(|(line_index, line)| {
            let trimmed = line.trim();
            if trimmed == "}," {
                Some(diagnostic(
                    line_index,
                    0,
                    line_index,
                    line.len(),
                    "Trailing comma after a block is not valid in the current Iris runtime.",
                    1,
                ))
            } else {
                None
            }
        })
        .collect()
}

fn check_missing_semicolons(text: &str) -> Vec<Value> {
    text.lines()
        .enumerate()
        .filter_map(|(line_index, raw)| {
            let line = raw.trim();
            if line.is_empty() || line.starts_with("//") {
                return None;
            }

            let exempt = line.ends_with(';')
                || line.ends_with('{')
                || line.ends_with('}')
                || line.starts_with("fn ")
                || line.starts_with("if ")
                || line.starts_with("else")
                || line.starts_with("while ")
                || line.starts_with("for ")
                || line.starts_with("match ")
                || line.starts_with("requires ")
                || line.starts_with("ensures ")
                || line == "where";

            if exempt {
                return None;
            }

            let suspicious = ["let", "return", "println", "print", "lookup", "input"]
                .iter()
                .any(|prefix| line.starts_with(prefix));

            if suspicious {
                Some(diagnostic(
                    line_index,
                    0,
                    line_index,
                    raw.len(),
                    "Statement may be missing a trailing semicolon.",
                    2,
                ))
            } else {
                None
            }
        })
        .collect()
}

fn collect_symbols(text: &str) -> Vec<Value> {
    text.lines()
        .enumerate()
        .filter_map(|(line_index, line)| {
            let trimmed = line.trim_start();
            if let Some(name) = trimmed.strip_prefix("fn ").and_then(extract_ident) {
                return Some(symbol(name, 12, line_index, line.len()));
            }

            for prefix in ["schema ", "state ", "trait ", "class ", "enum "] {
                if let Some(name) = trimmed.strip_prefix(prefix).and_then(extract_ident) {
                    return Some(symbol(name, 5, line_index, line.len()));
                }
            }

            None
        })
        .collect()
}

fn extract_ident(input: &str) -> Option<&str> {
    let end = input
        .char_indices()
        .take_while(|(_, ch)| ch.is_ascii_alphanumeric() || *ch == '_')
        .last()
        .map(|(index, ch)| index + ch.len_utf8())?;
    Some(&input[..end])
}

fn symbol(name: &str, kind: i32, line: usize, end_character: usize) -> Value {
    json!({
        "name": name,
        "kind": kind,
        "range": range(line, 0, line, end_character),
        "selectionRange": range(line, 0, line, end_character)
    })
}

fn diagnostic(
    start_line: usize,
    start_character: usize,
    end_line: usize,
    end_character: usize,
    message: &str,
    severity: i32,
) -> Value {
    json!({
        "range": range(start_line, start_character, end_line, end_character),
        "severity": severity,
        "source": "iris-lsp",
        "message": message
    })
}

fn range(start_line: usize, start_character: usize, end_line: usize, end_character: usize) -> Value {
    json!({
        "start": { "line": start_line, "character": start_character },
        "end": { "line": end_line, "character": end_character }
    })
}

fn completions() -> &'static [&'static str] {
    &[
        "fn", "let", "mut", "return", "if", "else", "while", "for", "in", "match", "schema", "state",
        "spec", "axiom", "trait", "class", "impl", "enum", "println", "print", "input", "lookup",
        "true", "false",
    ]
}

fn keyword_doc(word: &str) -> Option<&'static str> {
    match word {
        "fn" => Some("Defines a function in Iris."),
        "let" => Some("Creates a local binding."),
        "mut" => Some("Marks a value as mutable in Iris syntax."),
        "if" => Some("Conditional branch. Iris runtime currently supports block-style if/else."),
        "else" => Some("Fallback branch for if."),
        "while" => Some("Repeats a block while a condition remains true."),
        "for" => Some("Iterates over an iterable value."),
        "match" => Some("Pattern matching expression."),
        "schema" => Some("Declares a structural data type in full Iris design."),
        "state" => Some("Declares mutable state in full Iris design."),
        "spec" => Some("Declares a formal specification block."),
        "axiom" => Some("Declares a formal axiom block."),
        "trait" => Some("Declares shared behavior."),
        "class" => Some("Declares an Iris class."),
        "impl" => Some("Implements trait behavior for a target type."),
        "println" => Some("Built-in runtime output function."),
        "print" => Some("Built-in runtime output function without newline."),
        "input" => Some("Built-in runtime input function. Reads one line from stdin."),
        "lookup" => Some("Built-in runtime lookup function for simple key/value retrieval."),
        _ => None,
    }
}

fn word_at(text: &str, target_line: usize, target_character: usize) -> String {
    let line = text.lines().nth(target_line).unwrap_or_default();
    let character = target_character.min(line.len());
    let bytes = line.as_bytes();

    let mut start = character;
    while start > 0 && is_ident_char(bytes[start - 1] as char) {
        start -= 1;
    }

    let mut end = character;
    while end < bytes.len() && is_ident_char(bytes[end] as char) {
        end += 1;
    }

    line[start..end].to_string()
}

fn is_ident_char(ch: char) -> bool {
    ch.is_ascii_alphanumeric() || ch == '_'
}
