use crate::a2ui::{BASIC_NON_MEDIA_COMPONENTS, Store, extract_lines, validate_surface_composition};
use crate::config::Config;
use crate::model::ModelChoice;
use serde_json::{Value, json};
use std::time::Duration;

const CASES: &[Case] = &[
    Case {
        name: "pie-natural",
        prompt: "Do not call tools. The file extension counts are rs 41, md 7, and toml 2. Give me a pie chart with the numbers.",
        components: &["Chart"],
        chart_type: Some("pie"),
    },
    Case {
        name: "pie-single-slice",
        prompt: "Do not call tools. There are 41 rs files and no other file extensions. Show that distribution as a one-slice pie chart.",
        components: &["Chart"],
        chart_type: Some("pie"),
    },
    Case {
        name: "donut-natural",
        prompt: "Do not call tools. Show a donut chart for 62 completed, 23 active, and 15 blocked tasks. Include the total.",
        components: &["Chart"],
        chart_type: Some("donut"),
    },
    Case {
        name: "heatmap-natural",
        prompt: "Do not call tools. Show a heatmap of pull requests reviewed in 2025 and 2026, with each year as a row and January, February, and March as columns. Use values 2, 5, 3 for 2025 and 4, 1, 6 for 2026.",
        components: &["Chart"],
        chart_type: Some("heatmap"),
    },
    Case {
        name: "form-controls",
        prompt: "Do not call tools. Build an A2UI Form for a name, multiline notes, due date, and multiple checkbox priorities, with a submit button.",
        components: &["Form"],
        chart_type: None,
    },
    Case {
        name: "filterable-choices",
        prompt: "Do not call tools. Build an A2UI surface with a filterable mutually-exclusive ChoicePicker for Rust, Python, and TypeScript.",
        components: &["ChoicePicker"],
        chart_type: None,
    },
    Case {
        name: "composed-basics",
        prompt: "Do not call tools. Build one complete A2UI project dashboard. Compose a Card and Column containing Markdown text, an avatar Image, an Icon, a Divider, a weighted Row, Tabs, a dynamic List, a Modal, a Button, TextField, CheckBox, Slider, DateTimeInput, and a filterable mutually-exclusive ChoicePicker. Bind the controls to the data model and give the button an event.",
        components: BASIC_NON_MEDIA_COMPONENTS,
        chart_type: None,
    },
    Case {
        name: "media-players",
        prompt: "Do not call tools. Build one A2UI Card containing a Video with a poster URL and an AudioPlayer with a description. Use valid HTTPS media URLs and compose both into the root tree.",
        components: &["Video", "AudioPlayer"],
        chart_type: None,
    },
];

struct Case {
    name: &'static str,
    prompt: &'static str,
    components: &'static [&'static str],
    chart_type: Option<&'static str>,
}

struct Options {
    endpoint: Option<String>,
    model: Option<String>,
    case: Option<String>,
    attempts: u32,
}

pub(crate) fn run(args: impl Iterator<Item = String>) -> Result<(), String> {
    let options = parse_options(args)?;
    let config = Config::load(&crate::app::config_path())?;
    let endpoint = options
        .endpoint
        .unwrap_or_else(|| format!("http://127.0.0.1:{}", config.endpoint.port));
    let model_id = options
        .model
        .unwrap_or_else(|| config.model.id().to_owned());
    let model =
        ModelChoice::from_id(&model_id).ok_or_else(|| format!("unknown model `{model_id}`"))?;
    let cases = CASES
        .iter()
        .filter(|case| options.case.as_deref().is_none_or(|name| name == case.name))
        .collect::<Vec<_>>();
    if cases.is_empty() {
        return Err(format!(
            "unknown case; choose one of: {}",
            CASES
                .iter()
                .map(|case| case.name)
                .collect::<Vec<_>>()
                .join(", ")
        ));
    }
    let agent: ureq::Agent = ureq::Agent::config_builder()
        .timeout_connect(Some(Duration::from_secs(5)))
        .timeout_recv_response(Some(Duration::from_secs(30 * 60)))
        .timeout_recv_body(Some(Duration::from_secs(30)))
        .build()
        .into();
    let mut system = crate::agent::system_prompt(model, &config.generation.system_prompt, false);
    system.push_str("\n\n");
    system.push_str(crate::a2ui::SYSTEM_PROMPT);
    let metadata = Store::default().client_metadata();
    let mut passed = 0;
    let total = cases.len() as u32 * options.attempts;
    println!("A2UI live validation: {model_id} at {endpoint}");
    for case in cases {
        for attempt in 1..=options.attempts {
            let user = format!("{}\n\nA2UI client metadata:\n{}", case.prompt, metadata);
            let payload = json!({
                "model": model_id,
                "messages": [
                    {"role": "system", "content": system},
                    {"role": "user", "content": user}
                ],
                "reasoning_effort": "none",
                "temperature": 0,
                "max_tokens": 4096
            });
            let content = request(&agent, &endpoint, &payload)?;
            match validate(case, &content) {
                Ok(summary) => {
                    passed += 1;
                    println!("PASS {}#{attempt}: {summary}", case.name);
                }
                Err(error) => {
                    println!("FAIL {}#{attempt}: {error}", case.name);
                    println!(
                        "--- response ---\n{}\n--- end response ---",
                        bounded(&content)
                    );
                }
            }
        }
    }
    println!("A2UI validation summary: {passed}/{total} passed");
    if passed == total {
        Ok(())
    } else {
        Err(format!("{} live validation case(s) failed", total - passed))
    }
}

fn parse_options(mut args: impl Iterator<Item = String>) -> Result<Options, String> {
    let mut options = Options {
        endpoint: None,
        model: None,
        case: None,
        attempts: 1,
    };
    while let Some(argument) = args.next() {
        let value = match argument.as_str() {
            "--endpoint" | "--model" | "--case" | "--attempts" => args
                .next()
                .ok_or_else(|| format!("{argument} requires a value"))?,
            "--help" | "-h" => {
                println!(
                    "Usage: ds4-server validate-a2ui [--endpoint URL] [--model ID] [--case NAME] [--attempts N]"
                );
                std::process::exit(0);
            }
            _ => return Err(format!("unknown argument `{argument}`")),
        };
        match argument.as_str() {
            "--endpoint" => options.endpoint = Some(value.trim_end_matches('/').to_owned()),
            "--model" => options.model = Some(value),
            "--case" => options.case = Some(value),
            "--attempts" => {
                options.attempts = value
                    .parse::<u32>()
                    .ok()
                    .filter(|attempts| *attempts > 0)
                    .ok_or_else(|| "--attempts must be a positive integer".to_owned())?;
            }
            _ => unreachable!(),
        }
    }
    Ok(options)
}

fn request(agent: &ureq::Agent, endpoint: &str, payload: &Value) -> Result<String, String> {
    let mut response = agent
        .post(&format!("{endpoint}/v1/chat/completions"))
        .header("Content-Type", "application/json")
        .send(payload.to_string().as_bytes())
        .map_err(|error| format!("could not call the local endpoint: {error}"))?;
    let body = response
        .body_mut()
        .read_to_string()
        .map_err(|error| format!("could not read the local endpoint response: {error}"))?;
    let value: Value = serde_json::from_str(&body)
        .map_err(|error| format!("local endpoint returned invalid JSON: {error}"))?;
    value
        .pointer("/choices/0/message/content")
        .and_then(Value::as_str)
        .map(str::to_owned)
        .ok_or_else(|| format!("local endpoint response has no assistant content: {body}"))
}

fn validate(case: &Case, content: &str) -> Result<String, String> {
    let lines = extract_lines(content);
    if lines.is_empty() {
        return Err("the model emitted no complete fenced A2UI JSONL messages".into());
    }
    let mut store = Store::default();
    let mut errors = Vec::new();
    for (index, line) in lines.iter().enumerate() {
        if let Err(error) = &line.value {
            errors.push(format!("line {} is invalid JSON: {error}", index + 1));
        } else if let Err(error) = store.apply_raw(&line.raw, 1) {
            errors.push(format!(
                "line {} failed catalog validation: {error}",
                index + 1
            ));
        }
    }
    if !errors.is_empty() {
        return Err(errors.join("; "));
    }
    let surfaces = store.surfaces().collect::<Vec<_>>();
    let mut reachable_components = std::collections::BTreeSet::new();
    for surface in &surfaces {
        reachable_components.extend(
            validate_surface_composition(surface)
                .map_err(|error| format!("surface `{}` is not composable: {error}", surface.id))?,
        );
    }
    let components = surfaces
        .iter()
        .flat_map(|surface| surface.components.values())
        .filter_map(Value::as_object)
        .collect::<Vec<_>>();
    for expected in case.components {
        if !reachable_components.contains(*expected) {
            return Err(format!(
                "valid surface did not compose expected {expected} component into its root tree"
            ));
        }
    }
    if let Some(expected) = case.chart_type
        && !components.iter().any(|component| {
            component.get("component").and_then(Value::as_str) == Some("Chart")
                && component.get("chartType").and_then(Value::as_str) == Some(expected)
        })
    {
        return Err(format!("valid surface omitted expected {expected} chart"));
    }
    Ok(format!(
        "{} protocol message(s), {} surface(s), {} component(s)",
        lines.len(),
        surfaces.len(),
        components.len()
    ))
}

fn bounded(content: &str) -> &str {
    content
        .char_indices()
        .nth(8_000)
        .map_or(content, |(end, _)| &content[..end])
}