use std::collections::BTreeSet;

use dashmap::DashMap;

use tower_lsp::jsonrpc::Result as LspResult;
use tower_lsp::lsp_types::*;
use tower_lsp::{Client, LanguageServer, LspService, Server};

use dinoco_compiler::{compile, compile_only_ast};
use dinoco_formatter::{format_from_ast, FormatterConfig};

#[derive(Debug)]
struct LspServer {
    client: Client,

    documents: DashMap<Url, String>,
}

impl LspServer {
    pub fn new(client: Client) -> Self {
        Self { client, documents: DashMap::new() }
    }

    async fn validate_document(&self, uri: Url, text: String) {
        let mut diagnostics = vec![];

        if let Err(errors) = compile(&text) {
            for e in errors {
                let diagnostic = Diagnostic {
                    range: Range {
                        start: Position {
                            line: e.start_line.saturating_sub(1) as u32,
                            character: e.start_column.saturating_sub(1) as u32,
                        },
                        end: Position {
                            line: e.end_line.saturating_sub(1) as u32,
                            character: e.end_column.saturating_sub(1) as u32,
                        },
                    },
                    severity: Some(DiagnosticSeverity::ERROR),
                    code: Some(NumberOrString::String("E0001".to_string())),
                    source: Some("dinoco".to_string()),
                    message: e.message.clone(),
                    ..Default::default()
                };

                diagnostics.push(diagnostic);
            }
        }

        self.client.publish_diagnostics(uri, diagnostics, None).await;
    }
}

#[tower_lsp::async_trait]
impl LanguageServer for LspServer {
    async fn initialize(&self, _: InitializeParams) -> LspResult<InitializeResult> {
        Ok(InitializeResult {
            capabilities: ServerCapabilities {
                text_document_sync: Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::FULL)),
                completion_provider: Some(CompletionOptions {
                    resolve_provider: Some(false),

                    trigger_characters: Some(vec![
                        "@".to_string(),
                        " ".to_string(),
                        "(".to_string(),
                        ",".to_string(),
                        ":".to_string(),
                        "[".to_string(),
                    ]),
                    ..Default::default()
                }),
                document_formatting_provider: Some(OneOf::Left(true)),
                ..Default::default()
            },

            ..Default::default()
        })
    }

    async fn initialized(&self, _: InitializedParams) {
        self.client.log_message(MessageType::INFO, "Dinoco LSP Server is running!").await;
    }

    async fn shutdown(&self) -> LspResult<()> {
        Ok(())
    }

    async fn did_open(&self, params: DidOpenTextDocumentParams) {
        let uri = params.text_document.uri.clone();
        let text = params.text_document.text.clone();

        self.documents.insert(uri.clone(), text.clone());
        self.validate_document(uri, text).await;
    }

    async fn did_change(&self, mut params: DidChangeTextDocumentParams) {
        let uri = params.text_document.uri.clone();
        let text = params.content_changes.remove(0).text;

        self.documents.insert(uri.clone(), text.clone());
        self.validate_document(uri, text).await;
    }

    async fn did_close(&self, params: DidCloseTextDocumentParams) {
        self.documents.remove(&params.text_document.uri);
    }

    async fn formatting(&self, params: DocumentFormattingParams) -> LspResult<Option<Vec<TextEdit>>> {
        let uri = params.text_document.uri;

        let Some(doc) = self.documents.get(&uri) else {
            return Ok(None);
        };

        let text = doc.value().clone();
        let formatted = match compile_only_ast(&text) {
            Ok(schema) => {
                let config = FormatterConfig::default();

                format_from_ast(&schema, &config)
            }
            Err(_) => {
                return Ok(None);
            }
        };

        let line_count = text.lines().count() as u32;
        let full_range =
            Range { start: Position { line: 0, character: 0 }, end: Position { line: line_count, character: 0 } };

        Ok(Some(vec![TextEdit { range: full_range, new_text: formatted }]))
    }

    async fn completion(&self, params: CompletionParams) -> LspResult<Option<CompletionResponse>> {
        let uri = params.text_document_position.text_document.uri;
        let position = params.text_document_position.position;

        let doc_text = self.documents.get(&uri).map(|v| v.clone()).unwrap_or_default();
        let (model_names, enum_names) = collect_schema_identifiers(&doc_text);
        let lines: Vec<&str> = doc_text.lines().collect();

        let current_line = lines.get(position.line as usize).copied().unwrap_or("");
        let prefix_up_to_cursor = prefix_up_to_character(current_line, position.character);
        let last_word = prefix_up_to_cursor.split_whitespace().last().unwrap_or("");
        let mut completions = vec![];

        if is_inside_default_arguments(prefix_up_to_cursor) {
            completions.extend(default_value_completions());
            return Ok(Some(CompletionResponse::Array(completions)));
        }

        if let Some(relation_fragment) = relation_argument_fragment(prefix_up_to_cursor) {
            if is_relation_action_context(relation_fragment) {
                completions.extend(relation_action_completions());
            } else {
                completions.extend(relation_argument_completions());
            }
            return Ok(Some(CompletionResponse::Array(completions)));
        }

        if last_word.starts_with("@@") || prefix_up_to_cursor.ends_with("@@") {
            completions.extend(model_decorator_completions());
            return Ok(Some(CompletionResponse::Array(completions)));
        }

        if last_word.starts_with('@') || prefix_up_to_cursor.ends_with('@') {
            completions.extend(field_decorator_completions());
            completions.extend(model_decorator_completions());
            return Ok(Some(CompletionResponse::Array(completions)));
        }

        if is_field_type_context(prefix_up_to_cursor) {
            completions.extend(base_type_completions());
            completions.extend(custom_type_completions(&model_names, &enum_names));
            return Ok(Some(CompletionResponse::Array(completions)));
        }

        let is_root_level = !prefix_up_to_cursor.starts_with(' ') && !prefix_up_to_cursor.starts_with('\t');

        if is_root_level {
            completions.extend(vec![
                create_snippet_completion(
                    "model",
                    CompletionItemKind::KEYWORD,
                    "Define a new model",
                    "model ${1:ModelName} {\n\t$0\n}",
                ),
                create_snippet_completion(
                    "enum",
                    CompletionItemKind::KEYWORD,
                    "Define an enum",
                    "enum ${1:EnumName} {\n\t$0\n}",
                ),
                create_snippet_completion(
                    "config",
                    CompletionItemKind::KEYWORD,
                    "Database configuration",
                    "config {\n\tdatabase = \"${1:postgresql}\"\n\tdatabase_url = ${2:env(\"DATABASE_URL\")}\n\tredis = {\n\t\turl = ${3:env(\"REDIS_URL\")}\n\t}\n}",
                ),
            ]);
        } else {
            completions.extend(base_type_completions());
            completions.extend(custom_type_completions(&model_names, &enum_names));
        }

        Ok(Some(CompletionResponse::Array(completions)))
    }
}

fn prefix_up_to_character(line: &str, character: u32) -> &str {
    let requested = character as usize;
    let mut current = 0usize;

    for (index, _) in line.char_indices() {
        if current == requested {
            return &line[..index];
        }

        current += 1;
    }

    line
}

fn is_inside_default_arguments(prefix: &str) -> bool {
    prefix.rfind("@default(").is_some_and(|index| !prefix[index..].contains(')'))
}

fn relation_argument_fragment(prefix: &str) -> Option<&str> {
    let start = prefix.rfind("@relation(")?;
    let fragment = &prefix[(start + "@relation(".len())..];

    if fragment.contains(')') {
        None
    } else {
        Some(fragment)
    }
}

fn is_relation_action_context(fragment: &str) -> bool {
    let segment = fragment.rsplit(',').next().unwrap_or(fragment).trim_start();
    let is_on_delete = segment.strip_prefix("onDelete:").is_some_and(|value| value.trim().is_empty());
    let is_on_update = segment.strip_prefix("onUpdate:").is_some_and(|value| value.trim().is_empty());

    is_on_delete || is_on_update
}

fn is_field_type_context(prefix: &str) -> bool {
    if prefix.contains('=') {
        return false;
    }

    let without_comment = prefix.split('#').next().unwrap_or(prefix);
    let trimmed_start = without_comment.trim_start();

    if trimmed_start.is_empty()
        || trimmed_start.starts_with('@')
        || trimmed_start.starts_with("model ")
        || trimmed_start.starts_with("enum ")
        || trimmed_start.starts_with("config ")
        || trimmed_start.starts_with("}")
    {
        return false;
    }

    let tokens = trimmed_start.split_whitespace().collect::<Vec<_>>();

    match tokens.len() {
        1 => without_comment.ends_with(' ') || without_comment.ends_with('\t'),
        2 => !without_comment.contains('@') && !without_comment.ends_with(' ') && !without_comment.ends_with('\t'),
        _ => false,
    }
}

fn collect_schema_identifiers(doc_text: &str) -> (Vec<String>, Vec<String>) {
    let mut model_names = BTreeSet::new();
    let mut enum_names = BTreeSet::new();

    if let Ok(schema) = compile_only_ast(doc_text) {
        for table in schema.tables {
            model_names.insert(table.name);
        }

        for item in schema.enums {
            enum_names.insert(item.name);
        }
    }

    collect_block_names_from_source(doc_text, "model ", &mut model_names);
    collect_block_names_from_source(doc_text, "enum ", &mut enum_names);

    (model_names.into_iter().collect(), enum_names.into_iter().collect())
}

fn collect_block_names_from_source(doc_text: &str, keyword: &str, output: &mut BTreeSet<String>) {
    for line in doc_text.lines() {
        let trimmed = line.trim_start();
        let Some(rest) = trimmed.strip_prefix(keyword) else {
            continue;
        };

        let rest = rest.trim_start();
        let Some(identifier) = parse_identifier(rest) else {
            continue;
        };

        output.insert(identifier);
    }
}

fn parse_identifier(input: &str) -> Option<String> {
    let identifier: String =
        input.chars().take_while(|character| character.is_ascii_alphanumeric() || *character == '_').collect();

    if identifier.is_empty() {
        None
    } else {
        Some(identifier)
    }
}

fn default_value_completions() -> Vec<CompletionItem> {
    vec![
        create_completion(
            "autoincrement()",
            CompletionItemKind::FUNCTION,
            "Auto increment integer ID",
            "autoincrement()",
        ),
        create_completion("uuid()", CompletionItemKind::FUNCTION, "UUID v7 string ID", "uuid()"),
        create_completion("snowflake()", CompletionItemKind::FUNCTION, "Snowflake integer ID", "snowflake()"),
        create_completion("now()", CompletionItemKind::FUNCTION, "Current UTC timestamp", "now()"),
    ]
}

fn relation_argument_completions() -> Vec<CompletionItem> {
    vec![
        create_snippet_completion("name", CompletionItemKind::PROPERTY, "Relation name", "name: \"${1:RelationName}\""),
        create_snippet_completion("fields", CompletionItemKind::PROPERTY, "Local fields", "fields: [${1:fieldId}]"),
        create_snippet_completion(
            "references",
            CompletionItemKind::PROPERTY,
            "Referenced fields",
            "references: [${1:id}]",
        ),
        create_snippet_completion(
            "onDelete",
            CompletionItemKind::PROPERTY,
            "Action on delete",
            "onDelete: ${1:Cascade}",
        ),
        create_snippet_completion(
            "onUpdate",
            CompletionItemKind::PROPERTY,
            "Action on update",
            "onUpdate: ${1:Cascade}",
        ),
    ]
}

fn relation_action_completions() -> Vec<CompletionItem> {
    vec![
        create_completion("Cascade", CompletionItemKind::ENUM_MEMBER, "Cascade action", "Cascade"),
        create_completion("Restrict", CompletionItemKind::ENUM_MEMBER, "Restrict action", "Restrict"),
        create_completion("SetNull", CompletionItemKind::ENUM_MEMBER, "Set field to null", "SetNull"),
        create_completion("NoAction", CompletionItemKind::ENUM_MEMBER, "No action taken", "NoAction"),
    ]
}

fn field_decorator_completions() -> Vec<CompletionItem> {
    vec![
        create_completion("@id", CompletionItemKind::PROPERTY, "Primary key", "id"),
        create_completion("@unique", CompletionItemKind::PROPERTY, "Unique constraint", "unique"),
        create_snippet_completion("@default", CompletionItemKind::FUNCTION, "Default value", "default(${1:uuid()})"),
        create_snippet_completion(
            "@relation",
            CompletionItemKind::FUNCTION,
            "Relation definition",
            "relation(name: \"${1}\", fields: [${2}], references: [${3}], onDelete: ${4:Cascade}, onUpdate: ${5:Cascade})",
        ),
    ]
}

fn model_decorator_completions() -> Vec<CompletionItem> {
    vec![
        create_snippet_completion(
            "@@ids",
            CompletionItemKind::PROPERTY,
            "Composite primary key",
            "ids([${1:fieldA}, ${2:fieldB}])",
        ),
        create_snippet_completion(
            "@@uniques",
            CompletionItemKind::PROPERTY,
            "Composite unique key",
            "uniques([${1:fieldA}, ${2:fieldB}])",
        ),
        create_snippet_completion(
            "@@indexes",
            CompletionItemKind::PROPERTY,
            "Composite index",
            "indexes([${1:fieldA}, ${2:fieldB}])",
        ),
        create_snippet_completion(
            "@@table_name",
            CompletionItemKind::PROPERTY,
            "Custom table name",
            "table_name(\"${1:table_name}\")",
        ),
    ]
}

fn base_type_completions() -> Vec<CompletionItem> {
    vec![
        create_completion("String", CompletionItemKind::TYPE_PARAMETER, "String text type", "String"),
        create_completion("Integer", CompletionItemKind::TYPE_PARAMETER, "Integer number type", "Integer"),
        create_completion("Float", CompletionItemKind::TYPE_PARAMETER, "Floating point number", "Float"),
        create_completion("Boolean", CompletionItemKind::TYPE_PARAMETER, "True or false type", "Boolean"),
        create_completion("DateTime", CompletionItemKind::TYPE_PARAMETER, "Date and time in UTC", "DateTime"),
        create_completion("Date", CompletionItemKind::TYPE_PARAMETER, "Date without time", "Date"),
        create_completion("Json", CompletionItemKind::TYPE_PARAMETER, "JSON object type", "Json"),
    ]
}

fn custom_type_completions(model_names: &[String], enum_names: &[String]) -> Vec<CompletionItem> {
    let mut completions = Vec::new();

    for model_name in model_names {
        completions.push(create_completion(model_name, CompletionItemKind::CLASS, "Model type", model_name));
    }

    for enum_name in enum_names {
        completions.push(create_completion(enum_name, CompletionItemKind::ENUM, "Enum type", enum_name));
    }

    completions
}

fn create_completion(label: &str, kind: CompletionItemKind, detail: &str, insert_text: &str) -> CompletionItem {
    CompletionItem {
        label: label.to_string(),
        kind: Some(kind),
        detail: Some(detail.to_string()),
        insert_text: Some(insert_text.to_string()),
        ..Default::default()
    }
}

fn create_snippet_completion(label: &str, kind: CompletionItemKind, detail: &str, snippet: &str) -> CompletionItem {
    CompletionItem {
        label: label.to_string(),
        kind: Some(kind),
        detail: Some(detail.to_string()),
        insert_text: Some(snippet.to_string()),
        insert_text_format: Some(InsertTextFormat::SNIPPET),
        ..Default::default()
    }
}

#[tokio::main]
async fn main() {
    let stdin = tokio::io::stdin();
    let stdout = tokio::io::stdout();

    let (service, socket) = LspService::new(|client| LspServer::new(client));
    Server::new(stdin, stdout, socket).serve(service).await;
}
