from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

from review_agent.models import RegulatoryInfoPackageBatch
from review_agent.regulatory_info_package.constants import GENERATED_FILE_FAILED
from review_agent.regulatory_info_package.schemas import GeneratedFileResult, MergedField, TemplateSpec
from review_agent.regulatory_info_package.services.docx_document import write_docx_from_template
from review_agent.regulatory_info_package.services.legacy_doc_document import write_legacy_doc_or_fallback
from review_agent.regulatory_info_package.services.template_repository import copy_template_to_batch, template_specs
from review_agent.regulatory_info_package.storage import ensure_batch_subdir


def generate_package_documents(
    batch: RegulatoryInfoPackageBatch,
    config: dict,
    merged_fields: dict[str, MergedField],
) -> list[GeneratedFileResult]:
    specs = template_specs(config)
    with ThreadPoolExecutor(max_workers=min(4, len(specs) or 1)) as executor:
        futures = [executor.submit(_generate_one, batch, config, spec, merged_fields) for spec in specs]
        return [future.result() for future in as_completed(futures)]


def _generate_one(
    batch: RegulatoryInfoPackageBatch,
    config: dict,
    spec: TemplateSpec,
    merged_fields: dict[str, MergedField],
) -> GeneratedFileResult:
    try:
        template_path = copy_template_to_batch(batch, config, spec)
        generated_dir = ensure_batch_subdir(batch, "generated")
        output_path = generated_dir / spec.output_name
        adapter_summary = {}
        if spec.file_format == "doc":
            actual_path, status, adapter_summary = write_legacy_doc_or_fallback(template_path, output_path, merged_fields)
            actual_format = actual_path.suffix.lower().lstrip(".")
            highlight_count = missing_count = llm_only_count = 0
        else:
            highlight_count, missing_count, llm_only_count = write_docx_from_template(
                template_path,
                output_path,
                merged_fields,
                template_code=spec.code,
            )
            actual_path = output_path
            actual_format = "docx"
            status = "success"
        return GeneratedFileResult(
            template_code=spec.code,
            file_name=actual_path.name,
            requested_format=spec.file_format,
            actual_format=actual_format,
            status=status,
            path=str(actual_path),
            highlight_count=highlight_count,
            missing_count=missing_count,
            llm_only_count=llm_only_count,
        )
    except Exception as exc:
        return GeneratedFileResult(
            template_code=spec.code,
            file_name=spec.output_name,
            requested_format=spec.file_format,
            actual_format=spec.file_format,
            status=GENERATED_FILE_FAILED,
            error_message=str(exc),
        )