use clap::{Parser, Subcommand, ValueEnum};
use cyt_indexer::{
    apply_per_tool_overrides, build_catalog_from_tools, build_process_groups_options,
    get_skill_content_retrieve_result, get_skill_document, get_skill_line_content,
    get_skill_structure, load_catalog_from_dir, load_skills_index_from_dir,
    parse_tool_policy_pair, per_tool_policies_from_value, parse_tool_policy,
    policy_context_from_values, removed_chunks, retrieve_tools_from_catalog,
    write_reconstructed_skill, DecomposedCatalog, PageIndexConfig, PolicyContext,
    ReconstructOptions, RemovedChunksOptions, RetrieveOptions, SkillsBuilder, TokenCounterKind,
    WindowMode,
};
use serde_json::Value;
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};

#[derive(Parser)]
#[command(name = "cyt-indexer")]
#[command(about = "Tool schema decomposition and skills pageindex catalog indexing")]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Build catalogs (tools or skills)
    Build {
        #[command(subcommand)]
        target: BuildTarget,
    },
    /// Retrieve from catalogs (tools or skills)
    Retrieve {
        #[command(subcommand)]
        target: RetrieveTarget,
    },
    /// List removed tool chunks
    Removed {
        #[command(subcommand)]
        target: RemovedTarget,
    },
}

#[derive(Subcommand)]
enum BuildTarget {
    /// Build a decomposed tool catalog from API tools or catalog entries JSON
    Tools {
        #[arg(long)]
        tools: PathBuf,
        #[arg(long)]
        output: PathBuf,
    },
    /// Build a skills pageindex catalog from markdown skill directories
    Skills {
        #[arg(long = "skills")]
        skill_dirs: Vec<PathBuf>,
        #[arg(long)]
        output: PathBuf,
        #[arg(long, value_enum, default_value_t = WindowModeArg::Sentence)]
        window_mode: WindowModeArg,
        #[arg(long)]
        similarity_window: Option<usize>,
        #[arg(long, default_value_t = 2048)]
        chunk_size: usize,
        #[arg(long, value_enum, default_value_t = TokenCounterArg::Approximate)]
        token_counter: TokenCounterArg,
        #[arg(long, default_value_t = 0)]
        skip_window: usize,
        #[arg(long)]
        config: Option<PathBuf>,
    },
}

#[derive(Subcommand)]
enum RetrieveTarget {
    /// Retrieve merged tool schemas from pruner output
    Tools {
        #[arg(long)]
        catalog: PathBuf,
        #[arg(long)]
        input: PathBuf,
        #[arg(long)]
        output: PathBuf,
        #[arg(long)]
        config: Option<PathBuf>,
        #[arg(long)]
        system_policy: Option<String>,
        #[arg(long)]
        mcp_policy: Option<String>,
        #[arg(long)]
        per_tool: Option<PathBuf>,
        #[arg(long = "tool-policy", value_name = "TOOL=POLICY")]
        tool_policies: Vec<String>,
        #[arg(long, value_delimiter = ',')]
        preserve: Vec<String>,
        #[arg(long)]
        score_filter: bool,
        #[arg(long, hide = true)]
        no_score_filter: bool,
        #[arg(long)]
        removed_output: Option<PathBuf>,
    },
    /// Retrieve skill document metadata, structure, or content
    Skills {
        #[arg(long)]
        catalog: PathBuf,
        #[arg(long)]
        doc_id: String,
        #[arg(long, value_enum)]
        query: SkillQuery,
        #[arg(long = "line_num")]
        line_nums: Vec<String>,
        #[arg(long = "node_id")]
        node_ids: Vec<String>,
        #[arg(long = "chunk_id")]
        chunk_ids: Vec<String>,
        #[arg(long)]
        output: Option<PathBuf>,
        #[arg(long)]
        skills_index: Option<PathBuf>,
        #[arg(long)]
        keep_all_headers: bool,
    },
}

#[derive(Subcommand)]
enum RemovedTarget {
    /// List decomposed tool chunks not in survivors input
    Tools {
        #[arg(long)]
        catalog: PathBuf,
        #[arg(long)]
        input: PathBuf,
        #[arg(long)]
        output: PathBuf,
        #[arg(long)]
        full: Option<PathBuf>,
        #[arg(long)]
        score_filter: bool,
    },
}

#[derive(Clone, Copy, ValueEnum)]
enum WindowModeArg {
    Sentence,
    Word,
}

#[derive(Clone, Copy, ValueEnum)]
enum TokenCounterArg {
    Approximate,
    Character,
}

impl From<WindowModeArg> for WindowMode {
    fn from(value: WindowModeArg) -> Self {
        match value {
            WindowModeArg::Sentence => Self::Sentence,
            WindowModeArg::Word => Self::Word,
        }
    }
}

impl From<TokenCounterArg> for TokenCounterKind {
    fn from(value: TokenCounterArg) -> Self {
        match value {
            TokenCounterArg::Approximate => Self::Approximate,
            TokenCounterArg::Character => Self::Character,
        }
    }
}

#[derive(Clone, Copy, clap::ValueEnum)]
enum SkillQuery {
    Metadata,
    Structure,
    Content,
}

fn write_json_pretty(path: &Path, mut json: String) -> Result<(), Box<dyn std::error::Error>> {
    json.push('\n');
    fs::write(path, json)?;
    Ok(())
}

fn load_tools_array(tools: &Path) -> Result<Vec<Value>, Box<dyn std::error::Error>> {
    let raw = fs::read_to_string(tools)?;
    let tools_val: Value = serde_json::from_str(&raw)?;
    tools_val
        .as_array()
        .cloned()
        .or_else(|| tools_val.get("tools").and_then(|v| v.as_array()).cloned())
        .ok_or_else(|| "Expected tools array in input JSON".into())
}

fn policy_context_from_cli(
    config: Option<&Path>,
    system_policy: Option<&str>,
    mcp_policy: Option<&str>,
    per_tool: Option<&Path>,
    tool_policies: &[String],
) -> Result<PolicyContext, Box<dyn std::error::Error>> {
    let mut ctx = match config {
        Some(path) => {
            let raw = fs::read_to_string(path)?;
            policy_context_from_values(&serde_json::from_str(&raw)?)
        }
        None => PolicyContext::new(),
    };

    if let Some(s) = system_policy {
        ctx.system_policy = parse_tool_policy(s)
            .ok_or_else(|| format!("invalid system policy: {s}"))?;
    }
    if let Some(m) = mcp_policy {
        ctx.mcp_policy = parse_tool_policy(m)
            .ok_or_else(|| format!("invalid mcp policy: {m}"))?;
    }

    if let Some(path) = per_tool {
        let raw = fs::read_to_string(path)?;
        let val: Value = serde_json::from_str(&raw)?;
        apply_per_tool_overrides(&mut ctx, per_tool_policies_from_value(&val)?);
    }

    let mut cli_overrides = HashMap::new();
    for spec in tool_policies {
        let (tool_id, policy) = parse_tool_policy_pair(spec)?;
        cli_overrides.insert(tool_id, policy);
    }
    apply_per_tool_overrides(&mut ctx, cli_overrides);

    Ok(ctx)
}

/// Resolve skills retrieve JSON output: default `{catalog}/skill_out.json`;
/// relative `--output` paths are written under `--catalog`.
fn resolve_skills_output_path(catalog: &Path, output: Option<&Path>) -> PathBuf {
    match output {
        None => catalog.join("skill_out.json"),
        Some(path) if path.is_absolute() => path.to_path_buf(),
        Some(path) => catalog.join(path),
    }
}

fn catalog_path_utf8(catalog: &Path) -> Result<&str, Box<dyn std::error::Error>> {
    catalog.to_str().ok_or_else(|| {
        format!("catalog path is not valid UTF-8: {}", catalog.display()).into()
    })
}

fn run_build_tools(tools: &Path, output: &Path) -> Result<(), Box<dyn std::error::Error>> {
    let tools_arr = load_tools_array(tools)?;
    let index = build_catalog_from_tools(&tools_arr);
    fs::create_dir_all(output)?;
    for (rel, content) in &index.files {
        let path = output.join(rel);
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }
        fs::write(path, content)?;
    }
    eprintln!("Wrote {} tool files to {}", index.files.len(), output.display());
    Ok(())
}

fn page_index_config_from_cli(
    window_mode: WindowModeArg,
    similarity_window: Option<usize>,
    chunk_size: usize,
    token_counter: TokenCounterArg,
    skip_window: usize,
    config_path: Option<&Path>,
) -> Result<PageIndexConfig, Box<dyn std::error::Error>> {
    let mode: WindowMode = window_mode.into();
    let mut cfg = if let Some(path) = config_path {
        let raw = fs::read_to_string(path)?;
        PageIndexConfig::from_value(&serde_json::from_str(&raw)?)
    } else {
        PageIndexConfig::default()
    };
    cfg.bm25_cohesion.apply_mode_defaults(mode);
    cfg.bm25_cohesion.window_mode = mode;
    cfg.bm25_cohesion.chunk_size = chunk_size;
    cfg.bm25_cohesion.token_counter = token_counter.into();
    cfg.bm25_cohesion.skip_window = skip_window;
    if let Some(w) = similarity_window {
        cfg.bm25_cohesion.similarity_window = w;
    }
    cfg.bm25_cohesion.validate()?;
    Ok(cfg)
}

struct BuildSkillsArgs<'a> {
    skill_dirs: &'a [PathBuf],
    output: &'a Path,
    window_mode: WindowModeArg,
    similarity_window: Option<usize>,
    chunk_size: usize,
    token_counter: TokenCounterArg,
    skip_window: usize,
    config_path: Option<&'a Path>,
}

fn run_build_skills(args: &BuildSkillsArgs<'_>) -> Result<(), Box<dyn std::error::Error>> {
    if args.skill_dirs.is_empty() {
        return Err("at least one --skills directory is required".into());
    }
    let config = page_index_config_from_cli(
        args.window_mode,
        args.similarity_window,
        args.chunk_size,
        args.token_counter,
        args.skip_window,
        args.config_path,
    )?;
    let mut builder = SkillsBuilder::new(false, Some(args.output.to_path_buf()));
    builder.build_from_dirs(args.skill_dirs, &config)?;
    builder.write_catalog()?;
    let Some(index) = builder.index() else {
        return Err("index not built after write_catalog".into());
    };
    eprintln!(
        "Wrote skills catalog ({} documents, {} files) to {}",
        index.documents.len(),
        index.files.len(),
        args.output.display()
    );
    Ok(())
}

struct RetrieveArgs<'a> {
    catalog: &'a Path,
    input: &'a Path,
    output: &'a Path,
    config: Option<&'a Path>,
    system_policy: Option<&'a str>,
    mcp_policy: Option<&'a str>,
    per_tool: Option<&'a Path>,
    tool_policies: &'a [String],
    preserve: &'a [String],
    score_filter: bool,
    no_score_filter: bool,
    removed_output: Option<&'a Path>,
}

struct RetrieveSkillsArgs<'a> {
    catalog: &'a Path,
    doc_id: &'a str,
    query: SkillQuery,
    line_nums: &'a [String],
    node_ids: &'a [String],
    chunk_ids: &'a [String],
    output: Option<&'a Path>,
    skills_index: Option<&'a Path>,
    keep_all_headers: bool,
}

fn run_retrieve_tools(args: &RetrieveArgs<'_>) -> Result<(), Box<dyn std::error::Error>> {
    let apply_score_filter = args.score_filter && !args.no_score_filter;
    let ctx = policy_context_from_cli(
        args.config,
        args.system_policy,
        args.mcp_policy,
        args.per_tool,
        args.tool_policies,
    )?;
    let catalog_dict = load_catalog_from_dir(catalog_path_utf8(args.catalog)?)?;
    let mut store = DecomposedCatalog::from_catalog_dict(&catalog_dict);
    let input_raw = fs::read_to_string(args.input)?;
    let data: Value = serde_json::from_str(&input_raw)?;
    let preserve_set = (!args.preserve.is_empty()).then(|| args.preserve.to_vec());
    let process_groups =
        build_process_groups_options(&ctx, &catalog_dict, &store, preserve_set);
    let opts = RetrieveOptions {
        apply_decomposed_score_filter: apply_score_filter,
        process_groups,
    };
    let tools = retrieve_tools_from_catalog(&ctx, &data, &catalog_dict, &mut store, &opts);
    if tools.is_empty() && apply_score_filter {
        eprintln!(
            "Warning: retrieve produced 0 tools; rerank/pruner survivors usually need \
             score filter disabled (omit --score-filter)"
        );
    }
    write_json_pretty(args.output, serde_json::to_string_pretty(&tools)?)?;
    if let Some(removed_path) = args.removed_output {
        let removed = removed_chunks(
            &catalog_dict,
            &data,
            &RemovedChunksOptions {
                apply_decomposed_score_filter: apply_score_filter,
            },
        );
        write_json_pretty(removed_path, serde_json::to_string_pretty(&removed)?)?;
    }
    Ok(())
}

fn run_retrieve_skills(args: &RetrieveSkillsArgs<'_>) -> Result<(), Box<dyn std::error::Error>> {
    let reconstruct_opts = ReconstructOptions {
        keep_all_headers: args.keep_all_headers,
    };
    let index = if let Some(path) = args.skills_index {
        let raw = fs::read_to_string(path)?;
        let val: Value = serde_json::from_str(&raw)?;
        let mut idx = cyt_indexer::SkillsIndex::from_skills_index_json(&val)?;
        cyt_indexer::load_decomposed_files_for_index(args.catalog, &mut idx)?;
        idx
    } else {
        load_skills_index_from_dir(args.catalog)?
    };

    let result = match args.query {
        SkillQuery::Metadata => get_skill_document(&index.documents, args.doc_id),
        SkillQuery::Structure => get_skill_structure(&index.documents, args.doc_id),
        SkillQuery::Content => {
            if args.line_nums.is_empty() && args.node_ids.is_empty() && args.chunk_ids.is_empty() {
                return Err(
                    "content query requires at least one --line_num, --node_id, or --chunk_id".into(),
                );
            }
            let line_num_specs: Vec<&str> = args.line_nums.iter().map(String::as_str).collect();
            let node_id_specs: Vec<&str> = args.node_ids.iter().map(String::as_str).collect();
            let chunk_id_specs: Vec<&str> = args.chunk_ids.iter().map(String::as_str).collect();
            if !chunk_id_specs.is_empty() && line_num_specs.is_empty() && node_id_specs.is_empty() {
                get_skill_line_content(
                    &index,
                    args.doc_id,
                    &line_num_specs,
                    &node_id_specs,
                    &chunk_id_specs,
                )
            } else {
                get_skill_content_retrieve_result(
                    &index,
                    args.doc_id,
                    &line_num_specs,
                    &node_id_specs,
                    &chunk_id_specs,
                    &reconstruct_opts,
                )
            }
        }
    };

    if result.get("error").is_some() {
        return Err(result
            .get("error")
            .and_then(|v| v.as_str())
            .unwrap_or("retrieve skills failed")
            .into());
    }

    let output_path = resolve_skills_output_path(args.catalog, args.output);
    write_json_pretty(&output_path, serde_json::to_string_pretty(&result)?)?;

    if matches!(args.query, SkillQuery::Content) {
        let line_num_specs: Vec<&str> = args.line_nums.iter().map(String::as_str).collect();
        let node_id_specs: Vec<&str> = args.node_ids.iter().map(String::as_str).collect();
        let chunk_id_specs: Vec<&str> = args.chunk_ids.iter().map(String::as_str).collect();
        if !line_num_specs.is_empty() || !node_id_specs.is_empty() || !chunk_id_specs.is_empty() {
            let reconstructed = write_reconstructed_skill(
                args.catalog,
                &index,
                args.doc_id,
                &line_num_specs,
                &node_id_specs,
                &chunk_id_specs,
                &reconstruct_opts,
            )?;
            eprintln!(
                "Wrote reconstructed skill to {}",
                reconstructed.display()
            );
        }
    }

    Ok(())
}

fn load_full_catalog(
    catalog: &Path,
    full: Option<&Path>,
) -> Result<Value, Box<dyn std::error::Error>> {
    if let Some(path) = full {
        let raw = fs::read_to_string(path)?;
        return Ok(serde_json::from_str(&raw)?);
    }
    load_catalog_from_dir(catalog_path_utf8(catalog)?).map_err(Into::into)
}

fn run_removed_tools(
    catalog: &Path,
    input: &Path,
    output: &Path,
    full: Option<&Path>,
    score_filter: bool,
) -> Result<(), Box<dyn std::error::Error>> {
    let full_catalog = load_full_catalog(catalog, full)?;
    let input_raw = fs::read_to_string(input)?;
    let surviving: Value = serde_json::from_str(&input_raw)?;
    let removed = removed_chunks(
        &full_catalog,
        &surviving,
        &RemovedChunksOptions {
            apply_decomposed_score_filter: score_filter,
        },
    );
    write_json_pretty(output, serde_json::to_string_pretty(&removed)?)?;
    Ok(())
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let cli = Cli::parse();
    match cli.command {
        Commands::Build { target } => match target {
            BuildTarget::Tools { tools, output } => run_build_tools(&tools, &output)?,
            BuildTarget::Skills {
                skill_dirs,
                output,
                window_mode,
                similarity_window,
                chunk_size,
                token_counter,
                skip_window,
                config,
            } => run_build_skills(&BuildSkillsArgs {
                skill_dirs: &skill_dirs,
                output: &output,
                window_mode,
                similarity_window,
                chunk_size,
                token_counter,
                skip_window,
                config_path: config.as_deref(),
            })?,
        },
        Commands::Retrieve { target } => match target {
            RetrieveTarget::Tools {
                catalog,
                input,
                output,
                config,
                system_policy,
                mcp_policy,
                per_tool,
                tool_policies,
                preserve,
                score_filter,
                no_score_filter,
                removed_output,
            } => {
                let retrieve_args = RetrieveArgs {
                    catalog: &catalog,
                    input: &input,
                    output: &output,
                    config: config.as_deref(),
                    system_policy: system_policy.as_deref(),
                    mcp_policy: mcp_policy.as_deref(),
                    per_tool: per_tool.as_deref(),
                    tool_policies: &tool_policies,
                    preserve: &preserve,
                    score_filter,
                    no_score_filter,
                    removed_output: removed_output.as_deref(),
                };
                run_retrieve_tools(&retrieve_args)?;
            }
            RetrieveTarget::Skills {
                catalog,
                doc_id,
                query,
                line_nums,
                node_ids,
                chunk_ids,
                output,
                skills_index,
                keep_all_headers,
            } => {
                let retrieve_args = RetrieveSkillsArgs {
                    catalog: &catalog,
                    doc_id: &doc_id,
                    query,
                    line_nums: &line_nums,
                    node_ids: &node_ids,
                    chunk_ids: &chunk_ids,
                    output: output.as_deref(),
                    skills_index: skills_index.as_deref(),
                    keep_all_headers,
                };
                run_retrieve_skills(&retrieve_args)?;
            }
        },
        Commands::Removed { target } => match target {
            RemovedTarget::Tools {
                catalog,
                input,
                output,
                full,
                score_filter,
            } => run_removed_tools(&catalog, &input, &output, full.as_deref(), score_filter)?,
        },
    }
    Ok(())
}
