from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

from us_cpa.cases import CaseConflictError, CaseManager
from us_cpa.prepare import EfileExporter, PrepareEngine, render_case_forms
from us_cpa.questions import QuestionEngine, render_analysis, render_memo
from us_cpa.review import ReviewEngine, render_review_memo, render_review_summary
from us_cpa.sources import TaxYearCorpus, bootstrap_irs_catalog

COMMANDS = (
    "question",
    "prepare",
    "review",
    "fetch-year",
    "extract-docs",
    "render-forms",
    "export-efile-ready",
)


def _add_common_arguments(
    parser: argparse.ArgumentParser, *, include_tax_year: bool = True
) -> None:
    if include_tax_year:
        parser.add_argument("--tax-year", type=int, default=None)
    parser.add_argument("--case-dir", default=None)
    parser.add_argument("--format", choices=("json", "markdown"), default="json")


def _emit(payload: dict[str, Any], output_format: str) -> int:
    if output_format == "markdown":
        lines = [f"# {payload['command']}"]
        for key, value in payload.items():
            if key == "command":
                continue
            lines.append(f"- **{key}**: {value}")
        print("\n".join(lines))
    else:
        print(json.dumps(payload, indent=2))
    return 0


def _require_case_dir(args: argparse.Namespace) -> Path:
    if not args.case_dir:
        raise SystemExit("A case directory is required for this command.")
    return Path(args.case_dir).expanduser().resolve()


def _load_json_file(path_value: str | None) -> dict[str, Any]:
    if not path_value:
        return {}
    return json.loads(Path(path_value).expanduser().resolve().read_text())


def _ensure_question_corpus(corpus: TaxYearCorpus, tax_year: int) -> None:
    paths = corpus.paths_for_year(tax_year)
    required_slugs = {item.slug for item in bootstrap_irs_catalog(tax_year)}
    if not paths.manifest_path.exists():
        corpus.download_catalog(tax_year, bootstrap_irs_catalog(tax_year))
        return

    manifest = json.loads(paths.manifest_path.read_text())
    existing_slugs = {item["slug"] for item in manifest.get("sources", [])}
    if not required_slugs.issubset(existing_slugs):
        corpus.download_catalog(tax_year, bootstrap_irs_catalog(tax_year))


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="us-cpa",
        description="US federal individual tax workflow CLI.",
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    question = subparsers.add_parser("question", help="Answer a tax question.")
    _add_common_arguments(question)
    question.add_argument("--question", required=True)
    question.add_argument("--style", choices=("conversation", "memo"), default="conversation")

    prepare = subparsers.add_parser("prepare", help="Prepare a return case.")
    _add_common_arguments(prepare)

    review = subparsers.add_parser("review", help="Review a return case.")
    _add_common_arguments(review)
    review.add_argument("--style", choices=("conversation", "memo"), default="conversation")

    fetch_year = subparsers.add_parser(
        "fetch-year", help="Fetch tax-year forms and instructions."
    )
    _add_common_arguments(fetch_year, include_tax_year=False)
    fetch_year.add_argument("--tax-year", type=int, required=True)

    extract_docs = subparsers.add_parser(
        "extract-docs", help="Extract facts from case documents."
    )
    _add_common_arguments(extract_docs)
    extract_docs.add_argument("--create-case", action="store_true")
    extract_docs.add_argument("--case-label")
    extract_docs.add_argument("--facts-json")
    extract_docs.add_argument("--input-file", action="append", default=[])

    render_forms = subparsers.add_parser(
        "render-forms", help="Render compiled IRS forms."
    )
    _add_common_arguments(render_forms)

    export_efile = subparsers.add_parser(
        "export-efile-ready", help="Export an e-file-ready payload."
    )
    _add_common_arguments(export_efile)

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.command == "question":
        corpus = TaxYearCorpus()
        _ensure_question_corpus(corpus, args.tax_year)
        engine = QuestionEngine(corpus=corpus)
        case_facts: dict[str, Any] = {}
        if args.case_dir:
            manager = CaseManager(Path(args.case_dir))
            if manager.facts_path.exists():
                case_facts = {
                    key: value["value"]
                    for key, value in json.loads(manager.facts_path.read_text())["facts"].items()
                }
        analysis = engine.answer(
            question=args.question,
            tax_year=args.tax_year,
            case_facts=case_facts,
        )
        payload = {
            "command": "question",
            "format": args.format,
            "style": args.style,
            "taxYear": args.tax_year,
            "caseDir": args.case_dir,
            "question": args.question,
            "status": "answered",
            "analysis": analysis,
        }
        payload["rendered"] = (
            render_memo(analysis) if args.style == "memo" else render_analysis(analysis)
        )
        if args.format == "markdown":
            print(payload["rendered"])
            return 0
        return _emit(payload, args.format)

    if args.command == "extract-docs":
        case_dir = _require_case_dir(args)
        manager = CaseManager(case_dir)
        if args.create_case:
            if not args.case_label:
                raise SystemExit("--case-label is required when --create-case is used.")
            manager.create_case(case_label=args.case_label, tax_year=args.tax_year)
        elif not manager.manifest_path.exists():
            raise SystemExit("Case manifest not found. Use --create-case for a new case.")

        try:
            result = manager.intake(
                tax_year=args.tax_year,
                user_facts=_load_json_file(args.facts_json),
                document_paths=[
                    Path(path_value).expanduser().resolve() for path_value in args.input_file
                ],
            )
        except CaseConflictError as exc:
            print(json.dumps(exc.issue, indent=2))
            return 1
        payload = {
            "command": args.command,
            "format": args.format,
            **result,
        }
        return _emit(payload, args.format)

    if args.command == "prepare":
        case_dir = _require_case_dir(args)
        payload = {
            "command": args.command,
            "format": args.format,
            **PrepareEngine().prepare_case(case_dir),
        }
        return _emit(payload, args.format)

    if args.command == "render-forms":
        case_dir = _require_case_dir(args)
        manager = CaseManager(case_dir)
        manifest = manager.load_manifest()
        normalized = json.loads((case_dir / "return" / "normalized-return.json").read_text())
        artifacts = render_case_forms(case_dir, TaxYearCorpus(), normalized)
        payload = {
            "command": "render-forms",
            "format": args.format,
            "taxYear": manifest["taxYear"],
            "status": "rendered",
            **artifacts,
        }
        return _emit(payload, args.format)

    if args.command == "export-efile-ready":
        case_dir = _require_case_dir(args)
        payload = {
            "command": "export-efile-ready",
            "format": args.format,
            **EfileExporter().export_case(case_dir),
        }
        return _emit(payload, args.format)

    if args.command == "review":
        case_dir = _require_case_dir(args)
        review_payload = ReviewEngine().review_case(case_dir)
        payload = {
            "command": "review",
            "format": args.format,
            "style": args.style,
            **review_payload,
        }
        payload["rendered"] = (
            render_review_memo(review_payload)
            if args.style == "memo"
            else render_review_summary(review_payload)
        )
        if args.format == "markdown":
            print(payload["rendered"])
            return 0
        return _emit(payload, args.format)

    if args.command == "fetch-year":
        corpus = TaxYearCorpus()
        manifest = corpus.download_catalog(args.tax_year, bootstrap_irs_catalog(args.tax_year))
        payload = {
            "command": "fetch-year",
            "format": args.format,
            "taxYear": args.tax_year,
            "status": "downloaded",
            "sourceCount": manifest["sourceCount"],
            "manifestPath": corpus.paths_for_year(args.tax_year).manifest_path.as_posix(),
        }
        return _emit(payload, args.format)

    parser.error(f"Unsupported command: {args.command}")
    return 2


if __name__ == "__main__":
    sys.exit(main())
