//! Main LSP server implementation with full feature support

use std::{
    collections::HashMap,
    sync::Arc,
};

use tokio::sync::RwLock;
use tower_lsp::{
    jsonrpc::Result,
    lsp_types::{
        CodeActionParams,
        CodeActionResponse,
        CompletionOptions,
        CompletionParams,
        CompletionResponse,
        Diagnostic,
        DiagnosticSeverity,
        DidChangeTextDocumentParams,
        DidCloseTextDocumentParams,
        DidOpenTextDocumentParams,
        DocumentFormattingParams,
        DocumentSymbolParams,
        DocumentSymbolResponse,
        GotoDefinitionParams,
        GotoDefinitionResponse,
        Hover,
        HoverParams,
        HoverProviderCapability,
        InitializeParams,
        InitializeResult,
        InitializedParams,
        InlayHint,
        InlayHintParams,
        Location,
        MessageType,
        OneOf,
        Position,
        Range,
        ReferenceParams,
        RenameParams,
        SemanticTokensFullOptions,
        SemanticTokensLegend,
        SemanticTokensOptions,
        SemanticTokensParams,
        SemanticTokensResult,
        SemanticTokensServerCapabilities,
        ServerCapabilities,
        ServerInfo,
        TextDocumentSyncCapability,
        TextDocumentSyncKind,
        TextEdit,
        Url,
        WorkDoneProgressOptions,
        WorkspaceEdit,
    },
    Client,
    LanguageServer,
};

use super::{
    code_actions,
    completion,
    definition,
    document::{
        Document,
        ErrorSeverity,
    },
    formatting,
    hover,
    inlay_hints,
    references,
    rename,
    semantic_tokens,
    symbols,
};
use crate::resolve::names::NameTable;

/// Workspace-level state tracking all documents and cross-file references
#[derive(Debug)]
struct WorkspaceState {
    /// Combined name table from all open documents
    name_table: NameTable,
    /// Mapping from file index to URL
    file_urls: HashMap<usize, Url>,
}

impl WorkspaceState {
    fn new() -> Self {
        Self {
            name_table: NameTable::new(),
            file_urls: HashMap::new(),
        }
    }

    /// Rebuild the workspace name table from all documents
    fn rebuild(&mut self, documents: &HashMap<Url, Document>) {
        self.name_table = NameTable::new();
        self.file_urls.clear();

        // Build name table from all parsed documents
        for (file_index, (url, doc)) in documents.iter().enumerate() {
            self.file_urls.insert(file_index, url.clone());

            if let Some(ref ast) = doc.ast {
                // Build name table for this file
                // TODO: Properly merge file name tables with file_index
                // For now, we'll use the simple single-file approach
                let _ = (ast, file_index); // Use variables to avoid warnings
            }
        }
    }
}

/// The main language server instance
pub struct StorybookLanguageServer {
    client: Client,
    documents: Arc<RwLock<HashMap<Url, Document>>>,
    workspace: Arc<RwLock<WorkspaceState>>,
}

impl StorybookLanguageServer {
    pub fn new(client: Client) -> Self {
        Self {
            client,
            documents: Arc::new(RwLock::new(HashMap::new())),
            workspace: Arc::new(RwLock::new(WorkspaceState::new())),
        }
    }

    /// Rebuild workspace state after document changes
    async fn rebuild_workspace(&self) {
        let documents = self.documents.read().await;
        let mut workspace = self.workspace.write().await;
        workspace.rebuild(&documents);
    }

    /// Publish diagnostics for a document
    async fn publish_diagnostics(&self, uri: &Url, doc: &Document) {
        let mut positions = doc.positions.clone();
        let mut diagnostics = Vec::new();

        // Add parse errors
        for error in &doc.parse_errors {
            let (start_line, start_col) = positions.offset_to_position(error.start);
            let (end_line, end_col) = positions.offset_to_position(error.end);

            diagnostics.push(Diagnostic {
                range: Range {
                    start: Position {
                        line: start_line as u32,
                        character: start_col as u32,
                    },
                    end: Position {
                        line: end_line as u32,
                        character: end_col as u32,
                    },
                },
                severity: Some(match error.severity {
                    | ErrorSeverity::Error => DiagnosticSeverity::ERROR,
                    | ErrorSeverity::Warning => DiagnosticSeverity::WARNING,
                }),
                code: None,
                source: Some("storybook-parser".to_string()),
                message: error.message.clone(),
                related_information: None,
                tags: None,
                code_description: None,
                data: None,
            });
        }

        // Add semantic validation errors
        for error in &doc.resolve_errors {
            // For now, show resolve errors at the start of the document
            // TODO: Extract proper position information from ResolveError spans
            diagnostics.push(Diagnostic {
                range: Range {
                    start: Position {
                        line: 0,
                        character: 0,
                    },
                    end: Position {
                        line: 0,
                        character: 1,
                    },
                },
                severity: Some(DiagnosticSeverity::ERROR),
                code: None,
                source: Some("storybook-validator".to_string()),
                message: format!("{}", error),
                related_information: None,
                tags: None,
                code_description: None,
                data: None,
            });
        }

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

#[tower_lsp::async_trait]
impl LanguageServer for StorybookLanguageServer {
    async fn initialize(&self, _params: InitializeParams) -> Result<InitializeResult> {
        Ok(InitializeResult {
            capabilities: ServerCapabilities {
                text_document_sync: Some(TextDocumentSyncCapability::Kind(
                    TextDocumentSyncKind::FULL,
                )),
                hover_provider: Some(HoverProviderCapability::Simple(true)),
                completion_provider: Some(CompletionOptions {
                    trigger_characters: Some(vec![
                        ".".to_string(),
                        ":".to_string(),
                        "@".to_string(),
                    ]),
                    ..Default::default()
                }),
                document_symbol_provider: Some(OneOf::Left(true)),
                definition_provider: Some(OneOf::Left(true)),
                references_provider: Some(OneOf::Left(true)),
                document_formatting_provider: Some(OneOf::Left(true)),
                rename_provider: Some(OneOf::Right(tower_lsp::lsp_types::RenameOptions {
                    prepare_provider: Some(true),
                    work_done_progress_options: Default::default(),
                })),
                code_action_provider: Some(
                    tower_lsp::lsp_types::CodeActionProviderCapability::Simple(true),
                ),
                semantic_tokens_provider: Some(
                    SemanticTokensServerCapabilities::SemanticTokensOptions(
                        SemanticTokensOptions {
                            work_done_progress_options: WorkDoneProgressOptions::default(),
                            legend: SemanticTokensLegend {
                                token_types: semantic_tokens::LEGEND_TYPES.to_vec(),
                                token_modifiers: semantic_tokens::LEGEND_MODIFIERS
                                    .iter()
                                    .map(|s| tower_lsp::lsp_types::SemanticTokenModifier::new(s))
                                    .collect(),
                            },
                            range: Some(false),
                            full: Some(SemanticTokensFullOptions::Bool(true)),
                        },
                    ),
                ),
                inlay_hint_provider: Some(OneOf::Left(true)),
                ..Default::default()
            },
            server_info: Some(ServerInfo {
                name: "storybook-lsp".to_string(),
                version: Some(env!("CARGO_PKG_VERSION").to_string()),
            }),
        })
    }

    async fn initialized(&self, _params: InitializedParams) {
        self.client
            .log_message(MessageType::INFO, "Storybook LSP server initialized! 🎉")
            .await;
    }

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

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

        // Parse and store document
        let doc = Document::new(text);

        // Publish diagnostics
        self.publish_diagnostics(&uri, &doc).await;

        // Store document
        self.documents.write().await.insert(uri, doc);

        // Rebuild workspace for cross-file features
        self.rebuild_workspace().await;
    }

    async fn did_change(&self, params: DidChangeTextDocumentParams) {
        let uri = params.text_document.uri;
        if let Some(change) = params.content_changes.into_iter().next() {
            let text = change.text;

            // Update document
            let mut documents = self.documents.write().await;
            if let Some(doc) = documents.get_mut(&uri) {
                doc.update(text);
                self.publish_diagnostics(&uri, doc).await;
            }
            drop(documents); // Release lock before rebuilding
        }

        // Rebuild workspace for cross-file features
        self.rebuild_workspace().await;
    }

    async fn did_close(&self, params: DidCloseTextDocumentParams) {
        // Remove document from cache
        self.documents
            .write()
            .await
            .remove(&params.text_document.uri);

        // Rebuild workspace for cross-file features
        self.rebuild_workspace().await;
    }

    async fn hover(&self, params: HoverParams) -> Result<Option<Hover>> {
        let uri = params.text_document_position_params.text_document.uri;
        let position = params.text_document_position_params.position;

        let documents = self.documents.read().await;
        if let Some(doc) = documents.get(&uri) {
            // Try keyword hover first
            if let Some(hover) = hover::get_hover_info(
                &doc.text,
                position.line as usize,
                position.character as usize,
            ) {
                return Ok(Some(hover));
            }

            // Try semantic hover (symbols)
            if let Some(hover) = hover::get_semantic_hover_info(
                doc,
                position.line as usize,
                position.character as usize,
            ) {
                return Ok(Some(hover));
            }
        }

        Ok(None)
    }

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

        let documents = self.documents.read().await;
        if let Some(doc) = documents.get(&uri) {
            return Ok(completion::get_completions(doc, &params));
        }

        Ok(None)
    }

    async fn document_symbol(
        &self,
        params: DocumentSymbolParams,
    ) -> Result<Option<DocumentSymbolResponse>> {
        let uri = params.text_document.uri;

        let documents = self.documents.read().await;
        if let Some(doc) = documents.get(&uri) {
            if let Some(ref ast) = doc.ast {
                let mut positions = doc.positions.clone();
                let symbols = symbols::extract_symbols_from_ast(ast, &mut positions);
                return Ok(Some(DocumentSymbolResponse::Nested(symbols)));
            }
        }

        Ok(None)
    }

    async fn goto_definition(
        &self,
        params: GotoDefinitionParams,
    ) -> Result<Option<GotoDefinitionResponse>> {
        let uri = params
            .text_document_position_params
            .text_document
            .uri
            .clone();

        let documents = self.documents.read().await;
        if let Some(doc) = documents.get(&uri) {
            return Ok(definition::get_definition(doc, &params, &uri));
        }

        Ok(None)
    }

    async fn references(&self, params: ReferenceParams) -> Result<Option<Vec<Location>>> {
        let uri = params.text_document_position.text_document.uri.clone();

        let documents = self.documents.read().await;
        if let Some(doc) = documents.get(&uri) {
            return Ok(references::find_references(doc, &params, &uri));
        }

        Ok(None)
    }

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

        let documents = self.documents.read().await;
        if let Some(doc) = documents.get(&uri) {
            return Ok(formatting::format_document(doc, &params.options));
        }

        Ok(None)
    }

    async fn prepare_rename(
        &self,
        params: tower_lsp::lsp_types::TextDocumentPositionParams,
    ) -> Result<Option<tower_lsp::lsp_types::PrepareRenameResponse>> {
        let uri = params.text_document.uri;
        let position = params.position;

        let documents = self.documents.read().await;
        if let Some(doc) = documents.get(&uri) {
            if let Some(range) = rename::prepare_rename(doc, position) {
                return Ok(Some(tower_lsp::lsp_types::PrepareRenameResponse::Range(
                    range,
                )));
            }
        }

        Ok(None)
    }

    async fn rename(&self, params: RenameParams) -> Result<Option<WorkspaceEdit>> {
        let uri = params.text_document_position.text_document.uri.clone();

        let documents = self.documents.read().await;

        // Pass all documents for workspace-wide rename
        Ok(rename::get_rename_edits(&documents, &params, &uri))
    }

    async fn code_action(&self, params: CodeActionParams) -> Result<Option<CodeActionResponse>> {
        let documents = self.documents.read().await;

        // Generate code actions for the given position
        Ok(code_actions::get_code_actions(&documents, &params))
    }

    async fn semantic_tokens_full(
        &self,
        params: SemanticTokensParams,
    ) -> Result<Option<SemanticTokensResult>> {
        let uri = params.text_document.uri;
        let documents = self.documents.read().await;

        if let Some(doc) = documents.get(&uri) {
            Ok(semantic_tokens::get_semantic_tokens(doc))
        } else {
            Ok(None)
        }
    }

    async fn inlay_hint(&self, params: InlayHintParams) -> Result<Option<Vec<InlayHint>>> {
        let uri = params.text_document.uri;
        let range = params.range;
        let documents = self.documents.read().await;

        if let Some(doc) = documents.get(&uri) {
            Ok(inlay_hints::get_inlay_hints(doc, range.start, range.end))
        } else {
            Ok(None)
        }
    }
}
