use clap::Parser as ClapParser;
use colored::Colorize;
use miette::{IntoDiagnostic, Result, bail};
use nexus_lang::{
    FileId, InterpretWarning, Network, NexusError, PortDirection, Session, WalkEvent,
};
use rustyline::{DefaultEditor, error::ReadlineError};
use std::path::{Path, PathBuf};

#[derive(clap::ValueEnum, Clone, Copy, Debug, Default)]
enum OutputFormat {
    #[default]
    Table,
    Json,
    Dot,
    Mermaid,
    Plantuml,
}
use tabled::{
    Table, Tabled,
    settings::{Color, Modify, Style, object::Rows},
};

/// Nexus command-line interpreter tool.
#[derive(ClapParser, Debug)]
#[command(author, version, about)]
#[command(group = clap::ArgGroup::new("input").required(false).multiple(false).args(&["source_code", "file"]))]
#[command(after_help = "Omit any source code input arguments to run the interactive REPL.")]
struct Args {
    /// Scan input source code from a string.
    #[arg(short, long, group = "input")]
    source_code: Option<String>,

    /// Scan input source code from an input file.
    #[arg(short, long, group = "input")]
    file: Option<PathBuf>,

    /// Additional library search paths for 'use' imports (repeatable).
    #[arg(short = 'I', long = "search-path", value_name = "DIR")]
    search_paths: Vec<PathBuf>,

    /// Output format.
    #[arg(long, value_enum, default_value_t = OutputFormat::Table)]
    output: OutputFormat,

    /// Pretty-print output (JSON only).
    #[arg(long)]
    pretty: bool,
}

fn main() -> Result<()> {
    let Args {
        source_code,
        file,
        search_paths,
        output,
        pretty,
    } = Args::parse();

    if let Some(source_code) = source_code {
        execute_code(source_code.as_str(), search_paths, output, pretty)?;
    } else if let Some(file) = file {
        execute_file(file, search_paths, output, pretty)?;
    } else {
        run_repl()?;
    }

    Ok(())
}

/// Run an in-memory source string through a fresh session.
fn execute_code(
    source_code: &str,
    search_paths: Vec<PathBuf>,
    output: OutputFormat,
    pretty: bool,
) -> Result<()> {
    let mut session = Session::new().with_search_paths(search_paths);
    let id = session.load_string(source_code, Some("source"));
    run_session(&mut session, id, output, pretty)
}

/// Drive a loaded file through the session pipeline.
fn run_session(
    session: &mut Session,
    id: FileId,
    output: OutputFormat,
    pretty: bool,
) -> Result<()> {
    match output {
        OutputFormat::Table => run_session_table(session, id),
        OutputFormat::Json => run_session_json(session, id, pretty),
        OutputFormat::Dot => run_session_dot(session, id),
        OutputFormat::Mermaid => run_session_mermaid(session, id),
        OutputFormat::Plantuml => run_session_plantuml(session, id),
    }
}

/// Debug table output: prints source, tokens, AST, then the network as tables.
fn run_session_table(session: &mut Session, id: FileId) -> Result<()> {
    // Source input.
    {
        let source = session.source(id).expect("freshly loaded source");
        print!("{}:\n{source}\n", "Source input".white().bold());
        if !source.code().ends_with('\n') {
            println!();
        }
    }

    // Scanned tokens.
    let tokens = {
        let tokens = session.tokens(id).map_err(NexusError::from)?;
        tokens
            .iter()
            .map(|token| format!("{token}"))
            .collect::<Vec<_>>()
            .join(", ")
    };
    println!("{}:\n[{tokens}]\n", "Scanned tokens".white().bold());

    // Parsed AST.
    let ast = {
        let ast = session.ast(id)?;
        ast.iter()
            .map(|n| format!("  {n}"))
            .collect::<Vec<_>>()
            .join("\n")
    };
    println!("{}:\n[\n{ast}\n]\n", "Parsed AST".white().bold());

    // Interpreted result.
    println!("{}:\n", "Interpreted result".white().bold());
    let network = session.evaluate(id)?;
    println!();

    print_warnings(session.warnings());
    print_network(&network);

    Ok(())
}

/// JSON output: evaluates the program and serializes the network to stdout.
fn run_session_json(session: &mut Session, id: FileId, pretty: bool) -> Result<()> {
    let network = session.evaluate(id)?;

    print_warnings(session.warnings());

    let json = if pretty {
        serde_json::to_string_pretty(&network)
    } else {
        serde_json::to_string(&network)
    }
    .into_diagnostic()?;

    println!("{json}");

    Ok(())
}

/// DOT output: evaluates the program and prints the network as a Graphviz digraph.
fn run_session_dot(session: &mut Session, id: FileId) -> Result<()> {
    let network = session.evaluate(id)?;

    print_warnings(session.warnings());
    print!("{}", network.to_dot());

    Ok(())
}

/// Mermaid output: evaluates the program and prints the network as a Mermaid flowchart.
fn run_session_mermaid(session: &mut Session, id: FileId) -> Result<()> {
    let network = session.evaluate(id)?;

    print_warnings(session.warnings());
    print!("{}", network.to_mermaid());

    Ok(())
}

/// `PlantUML` output: evaluates the program and prints the network as a `PlantUML` component diagram.
fn run_session_plantuml(session: &mut Session, id: FileId) -> Result<()> {
    let network = session.evaluate(id)?;

    print_warnings(session.warnings());
    print!("{}", network.to_plantuml());

    Ok(())
}

fn print_warnings(warnings: &[InterpretWarning]) {
    for w in warnings {
        eprintln!("{:?}", miette::Report::new(w.clone()));
    }
}

/// Format a port direction as a colored string.
fn format_direction(dir: PortDirection) -> String {
    match dir {
        PortDirection::In => "in".bright_blue().to_string(),
        PortDirection::Out => "out".bright_red().to_string(),
        PortDirection::Inout => "inout".bright_magenta().to_string(),
    }
}

/// Format a port list as a comma-separated string like `"in input, out output"`.
fn format_ports(ports: &[nexus_lang::Port]) -> String {
    if ports.is_empty() {
        return String::new();
    }
    ports
        .iter()
        .map(|p| format!("{} {}", format_direction(p.direction), p.name))
        .collect::<Vec<_>>()
        .join(", ")
}

/// Format node attributes as a comma-separated string like `"name=Sensor, rate=10"`.
fn format_attributes(attributes: &std::collections::BTreeMap<String, String>) -> String {
    attributes
        .iter()
        .map(|(name, value)| format!("{name}={value}"))
        .collect::<Vec<_>>()
        .join(", ")
}

/// Row type for the nodes table.
#[derive(Tabled)]
struct NodeRow {
    path: String,
    #[tabled(rename = "type")]
    schema_type: String,
    ports: String,
    attributes: String,
}

/// Row type for the groups table.
#[derive(Tabled)]
struct GroupRow {
    path: String,
    ports: String,
}

/// Row type for the connections table.
#[derive(Tabled)]
struct ConnectionRow {
    from: String,
    to: String,
}

/// Walk the node tree and collect rows for leaf nodes and group nodes separately.
fn collect_rows(network: &Network, leaf_rows: &mut Vec<NodeRow>, group_rows: &mut Vec<GroupRow>) {
    let mut segments: Vec<&str> = Vec::new();
    for event in network.walk() {
        match event {
            WalkEvent::Enter(node) => {
                segments.push(&node.name);
                group_rows.push(GroupRow {
                    path: segments.join("/"),
                    ports: format_ports(&node.ports),
                });
            }
            WalkEvent::Leaf(node) => {
                segments.push(&node.name);
                leaf_rows.push(NodeRow {
                    path: segments.join("/"),
                    schema_type: node.type_name.clone(),
                    ports: format_ports(&node.ports),
                    attributes: format_attributes(&node.attributes),
                });
                segments.pop();
            }
            WalkEvent::Leave(_) => {
                segments.pop();
            }
        }
    }
}

/// Print the network as tabled tables.
fn print_network(network: &Network) {
    let paths = network.node_paths();

    let mut leaf_rows: Vec<NodeRow> = Vec::new();
    let mut group_rows: Vec<GroupRow> = Vec::new();

    collect_rows(network, &mut leaf_rows, &mut group_rows);

    println!("{}:", "Network".white().bold());
    println!();

    // Nodes table.
    println!("Nodes:");
    if leaf_rows.is_empty() {
        println!("(none)");
    } else {
        println!(
            "{}",
            Table::new(leaf_rows)
                .with(Style::modern())
                .with(Modify::new(Rows::first()).with(Color::BOLD))
        );
    }
    println!();

    // Groups table.
    println!("Groups:");
    if group_rows.is_empty() {
        println!("(none)");
    } else {
        println!(
            "{}",
            Table::new(group_rows)
                .with(Style::modern())
                .with(Modify::new(Rows::first()).with(Color::BOLD))
        );
    }
    println!();

    // Connections table.
    println!("Connections:");
    if network.connections.is_empty() {
        println!("(none)");
    } else {
        let conn_rows: Vec<ConnectionRow> = network
            .connections
            .iter()
            .map(|c| {
                let from_path = paths.get(&c.from.node_id).cloned().unwrap_or_default();
                let to_path = paths.get(&c.to.node_id).cloned().unwrap_or_default();
                ConnectionRow {
                    from: format!("{from_path}/{}", c.from.port_name),
                    to: format!("{to_path}/{}", c.to.port_name),
                }
            })
            .collect();
        println!(
            "{}",
            Table::new(conn_rows)
                .with(Style::modern())
                .with(Modify::new(Rows::first()).with(Color::BOLD))
        );
    }
    println!();
}

fn execute_file(
    file: impl AsRef<Path>,
    search_paths: Vec<PathBuf>,
    output: OutputFormat,
    pretty: bool,
) -> Result<()> {
    let mut session = Session::new().with_search_paths(search_paths);
    // `load_file` records the canonical path so `use` imports resolve relative to it.
    let id = session.load_file(file.as_ref()).map_err(NexusError::from)?;
    run_session(&mut session, id, output, pretty)
}

fn run_repl() -> Result<()> {
    let mut rl = DefaultEditor::new().into_diagnostic()?;

    println!("{}", "Interactive Nexus REPL. Press Ctrl+D to exit.".bold());

    loop {
        match rl.readline("> ") {
            Ok(line) => {
                rl.add_history_entry(line.clone()).into_diagnostic()?;
                match execute_code(line.as_str(), Vec::new(), OutputFormat::Table, false) {
                    Ok(()) => {}
                    Err(e) => eprintln!("{e:?}"),
                }
            }
            Err(ReadlineError::Eof) => break,
            Err(ReadlineError::Interrupted) => {
                eprintln!("interrupted");
                break;
            }
            _ => {
                bail!("failed to parse input");
            }
        }
    }

    Ok(())
}
