//! Regenerate derived wire artifacts from `tl-core`, `tl-policy`, and
//! `tl-server` (for the OpenAPI doc).
//!
//! Run from the repo root:
//!     cargo run -p tl-codegen
//!     cargo run -p tl-codegen -- --check     # CI mode: fail if drift
//!
//! Outputs:
//!     docs/openapi.yaml
//!     policies/check-request.schema.json
//!     policies/decision.schema.json
//!     policies/policy.schema.json
//!     sdks/typescript/src/generated/   (placeholder until ts-rs is wired)

use std::fs;
use std::path::{Path, PathBuf};

use anyhow::{bail, Context, Result};
use clap::Parser;
use schemars::schema_for;
use tl_core::{AgentProfile, CheckRequest, Decision};
use tl_policy::Policy;
use tl_server::ApiDoc;
use utoipa::OpenApi;

#[derive(Parser)]
struct Args {
    /// Fail with non-zero exit if generated files differ from disk.
    /// Wire this into CI to catch missed regenerations.
    #[arg(long)]
    check: bool,
}

fn repo_root() -> PathBuf {
    let manifest = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    manifest
        .parent()
        .and_then(Path::parent)
        .map(Path::to_path_buf)
        .expect("repo root")
}

fn write_or_check(path: &Path, contents: &str, check: bool) -> Result<()> {
    if check {
        let on_disk =
            fs::read_to_string(path).with_context(|| format!("read {}", path.display()))?;
        if on_disk != contents {
            bail!(
                "drift: {} is out of sync with derived types. Run `cargo run -p tl-codegen` and commit.",
                path.display()
            );
        }
        Ok(())
    } else {
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }
        fs::write(path, contents).with_context(|| format!("write {}", path.display()))?;
        println!("wrote {}", path.display());
        Ok(())
    }
}

fn main() -> Result<()> {
    let args = Args::parse();
    let root = repo_root();

    // 1. OpenAPI YAML — sourced from tl-server's annotated handlers.
    let mut openapi = serde_yaml::to_value(ApiDoc::openapi()).context("render openapi")?;
    patch_openapi_check_request(&mut openapi);
    let openapi_yaml = serde_yaml::to_string(&openapi).context("serialize openapi")?;
    write_or_check(&root.join("docs/openapi.yaml"), &openapi_yaml, args.check)?;

    // 2. JSON Schemas for each wire type. Consumers: dashboard editor,
    //    request validator, downstream Pydantic generation.
    let mut check_schema = serde_json::to_value(schema_for!(CheckRequest))?;
    patch_json_check_request(&mut check_schema);
    let check_json = serde_json::to_string_pretty(&check_schema)? + "\n";
    write_or_check(
        &root.join("policies/check-request.schema.json"),
        &check_json,
        args.check,
    )?;

    let schemas = [
        ("policies/decision.schema.json", schema_for!(Decision)),
        ("policies/policy.schema.json", schema_for!(Policy)),
        (
            "policies/agent-profile.schema.json",
            schema_for!(AgentProfile),
        ),
    ];
    for (rel, schema) in schemas {
        let json = serde_json::to_string_pretty(&schema)? + "\n";
        write_or_check(&root.join(rel), &json, args.check)?;
    }

    // 3. TypeScript types via ts-rs. ts-rs writes files when its generated
    //    `export_bindings_*` test functions run; we invoke them here as a
    //    child process. TS_RS_EXPORT_DIR points the writer at the SDK dir.
    let ts_dir = root.join("sdks/typescript/src/generated");
    fs::create_dir_all(&ts_dir).ok();

    let status = std::process::Command::new(env!("CARGO"))
        .args([
            "test",
            "-p",
            "tl-core",
            "--features",
            "ts-export",
            "--",
            "export_bindings",
        ])
        .env("TS_RS_EXPORT_DIR", &ts_dir)
        .current_dir(&root)
        .status()
        .context("run ts-rs export via cargo test")?;
    if !status.success() {
        bail!("ts-rs export failed; check `cargo test -p tl-core --features ts-export`");
    }
    normalize_typescript(&ts_dir)?;

    let readme_path = ts_dir.join("README.md");
    let readme = "Generated by `cargo run -p tl-codegen`. Do not edit by hand.\n";
    write_or_check(&readme_path, readme, args.check)?;

    // 4. Pydantic v2 models from the OpenAPI spec via datamodel-code-generator.
    //    The OpenAPI YAML was just written above, so its components match
    //    the canonical Rust types. Skipped (with a warning) if the tool is
    //    not installed locally; CI must have it installed.
    let py_out = root.join("sdks/python/src/trustloopguard/_generated/types.py");
    let openapi_path = root.join("docs/openapi.yaml");
    if let Some(rendered) = render_pydantic(&openapi_path)? {
        write_or_check(&py_out, &rendered, args.check)?;
    } else if args.check {
        bail!(
            "datamodel-codegen not installed but --check was requested. \
             Install with: pip install datamodel-code-generator"
        );
    } else {
        eprintln!(
            "warning: datamodel-codegen not installed; skipping Python type generation. \
             Install with: pip install datamodel-code-generator"
        );
    }

    println!(
        "{}",
        if args.check {
            "ok: artifacts in sync"
        } else {
            "ok: artifacts regenerated"
        }
    );
    Ok(())
}

fn run_context_constraints_json() -> serde_json::Value {
    serde_json::json!([
        {
            "not": {
                "required": ["run_event", "run_event_id"]
            }
        },
        {
            "anyOf": [
                {
                    "not": {
                        "required": ["run_event"]
                    }
                },
                {
                    "required": ["run_id"]
                }
            ]
        },
        {
            "anyOf": [
                {
                    "not": {
                        "required": ["run_event_id"]
                    }
                },
                {
                    "required": ["run_id"]
                }
            ]
        }
    ])
}

fn patch_json_check_request(schema: &mut serde_json::Value) {
    if let Some(properties) = schema
        .get_mut("properties")
        .and_then(serde_json::Value::as_object_mut)
    {
        for field in ["run_id", "run_event_id"] {
            if let Some(property) = properties.get_mut(field) {
                property["format"] = serde_json::Value::String("uuid".into());
            }
        }
    }
    schema["allOf"] = run_context_constraints_json();
}

fn patch_openapi_check_request(openapi: &mut serde_yaml::Value) {
    let Some(check_request) = openapi
        .get_mut("components")
        .and_then(|components| components.get_mut("schemas"))
        .and_then(|schemas| schemas.get_mut("CheckRequest"))
    else {
        return;
    };

    if let Some(properties) = check_request
        .get_mut("properties")
        .and_then(serde_yaml::Value::as_mapping_mut)
    {
        for field in ["run_id", "run_event_id"] {
            if let Some(property) = properties.get_mut(serde_yaml::Value::String(field.into())) {
                if let Some(mapping) = property.as_mapping_mut() {
                    mapping.insert(
                        serde_yaml::Value::String("format".into()),
                        serde_yaml::Value::String("uuid".into()),
                    );
                }
            }
        }
    }

    check_request["allOf"] =
        serde_yaml::to_value(run_context_constraints_json()).expect("constraints serialize");
}

fn normalize_typescript(dir: &Path) -> Result<()> {
    for entry in fs::read_dir(dir)? {
        let path = entry?.path();
        if path.extension().and_then(|value| value.to_str()) != Some("ts") {
            continue;
        }
        let body = fs::read_to_string(&path)?;
        let normalized = body
            .lines()
            .map(str::trim_end)
            .collect::<Vec<_>>()
            .join("\n")
            + "\n";
        if normalized != body {
            fs::write(&path, normalized)?;
        }
    }
    Ok(())
}

/// Run `datamodel-codegen` against the freshly written OpenAPI YAML and
/// return the generated Python source as a string. Returns `Ok(None)` if
/// the tool isn't on PATH so callers can decide whether to warn or fail.
fn render_pydantic(openapi_path: &Path) -> Result<Option<String>> {
    let tmp = tempfile::NamedTempFile::new()?;
    let output = std::process::Command::new("datamodel-codegen")
        .args([
            "--input",
            openapi_path.to_str().unwrap(),
            "--input-file-type",
            "openapi",
            "--output-model-type",
            "pydantic_v2.BaseModel",
            "--target-python-version",
            "3.10",
            "--use-standard-collections",
            "--use-union-operator",
            // Without this every run produces a fresh timestamp, breaking
            // the --check drift gate.
            "--disable-timestamp",
            "--output",
            tmp.path().to_str().unwrap(),
        ])
        .output();
    match output {
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
        Err(e) => Err(e.into()),
        Ok(out) if !out.status.success() => bail!(
            "datamodel-codegen failed:\n{}",
            String::from_utf8_lossy(&out.stderr)
        ),
        Ok(_) => {
            let body = fs::read_to_string(tmp.path())?;
            let header = "# Generated by `cargo run -p tl-codegen` from docs/openapi.yaml.\n\
                          # Do not hand-edit.\n\n";
            Ok(Some(format!("{header}{body}")))
        }
    }
}
