from __future__ import annotations

import json
from pathlib import Path
from typing import Optional

import typer

from core.models import CompileInput, DecisionVars, Ontology, Resident, RuleIR
from solver.builder import build_model
from nlp.parser import derive_cohorts_from_text


app = typer.Typer(help="Compile NL rules + ontology to OR-Tools decision vars (POC)")


@app.command()
def compile(
    ontology_json: Path = typer.Argument(..., help="Path to ontology.json"),
    rules_json: Optional[Path] = typer.Argument(None, help="Path to rules.json (RuleIR list)"),
    out_json: Path = typer.Option(Path("decision_vars.json"), help="Output path for variable names"),
    report_md: Path = typer.Option(Path("compile_report.md"), help="Compile report output"),
    cohorts_text: Optional[Path] = typer.Option(None, help="Optional free text to derive cohorts"),
):
    onto = Ontology.model_validate_json(ontology_json.read_text())
    rules = []
    if rules_json and rules_json.exists():
        rules = [RuleIR.model_validate(obj) for obj in json.loads(rules_json.read_text())]
    cohorts = {}
    if cohorts_text and cohorts_text.exists():
        cohorts = derive_cohorts_from_text(cohorts_text.read_text(), [r.id for r in onto.residents])

    ci = CompileInput(ontology=onto, rules=rules, cohorts=cohorts)

    model, dv, var_map, report = build_model(ci)

    out = {"x_vars": dv.x_vars, "count": len(dv.x_vars), "cohorts": cohorts}
    out_json.write_text(json.dumps(out, indent=2))
    report_md.write_text(
        f"# Compile report\n\n"
        f"Residents: {report.num_residents}\n\n"
        f"Blocks: {report.num_blocks}\n\n"
        f"Rotations: {report.num_rotations}\n\n"
        f"x-vars: {report.num_x_vars}\n\n"
        f"aux-vars: {report.num_aux_vars}\n"
    )
    typer.echo(f"Wrote {len(dv.x_vars)} decision variables to {out_json}; report to {report_md}")


if __name__ == "__main__":
    app()