"""
technical_file_builder.py — EU AI Act Annex IV Technical File builder.

Orchestrates all SichGate artefacts (findings, AIBoM, attestation, differential
report) into a structured Technical File complying with EU AI Act Article 11
and Annex IV (Regulation (EU) 2024/1689).

Usage::

    builder = TechnicalFileBuilder(
        findings=annotated_findings,
        aibom=aibom_capture.to_cyclonedx(),
        differential_report=diff_report.to_dict(),
        attestation=signed_report["attestation"],
    )
    technical_file = builder.build()
"""
from __future__ import annotations

import json
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Dict, List, Optional

from sichgate.annex_iv import ANNEX_IV_SECTIONS, get_sections_for_attack, annotate_finding


# ---------------------------------------------------------------------------
# Full EU AI Act Annex IV requirement text (Regulation (EU) 2024/1689)
# ---------------------------------------------------------------------------

_SECTION_REQUIREMENTS: Dict[str, str] = {
    "1": (
        "A general description of the AI system including its intended purpose, "
        "the persons responsible for it, its version and date, and the forms "
        "in which it is placed on the market or put into service."
    ),
    "1a": (
        "The intended purpose of the AI system, the persons responsible for it, "
        "and the version and date of the system."
    ),
    "1b": (
        "A description of how the AI system interacts with hardware, software, "
        "or other AI systems, where applicable."
    ),
    "1c": (
        "The relevant software or firmware versions and any update requirement "
        "relevant to maintaining compliance."
    ),
    "1d": (
        "A description of all the forms in which the AI system is placed on "
        "the market or put into service, such as software packages embedded "
        "into hardware devices, downloads, or APIs."
    ),
    "1e": (
        "A description of the hardware on which the AI system is intended to run."
    ),
    "1f": (
        "Where applicable, photographs or illustrations showing the external "
        "features, marking, and internal layout of products."
    ),
    "1g": (
        "Instructions of use for the deployer and, where applicable, "
        "a basic description of the user interface."
    ),
    "2": (
        "A detailed description of the elements of the AI system and of the "
        "process for its development."
    ),
    "2a": (
        "The methods and steps performed for the development of the AI system, "
        "including, where relevant, recourse to pre-trained systems or tools "
        "provided by third parties."
    ),
    "2b": (
        "The design specifications of the system, namely the general logic of "
        "the AI system and of the algorithms, the key design choices including "
        "the rationale and assumptions made."
    ),
    "2c": (
        "A description of the system architecture explaining how software "
        "components build on or feed into each other and integrate into the "
        "overall processing, including, where applicable, a description of the "
        "computational resources used."
    ),
    "2d": (
        "Data requirements including datasheets describing the training "
        "methodologies and techniques and the training datasets used, including "
        "information about data provenance, scope, main characteristics, and "
        "how the data was obtained, selected, labelled and cleaned."
    ),
    "2e": (
        "Assessment of the human oversight measures needed in accordance with "
        "Article 14, including an assessment of the technical measures needed "
        "to facilitate the interpretation of the outputs of AI systems by the "
        "deployers."
    ),
    "2f": (
        "Where applicable, a description of the pre-determined changes to the "
        "AI system and its performance, together with all the relevant information "
        "relating to the technical solutions adopted to ensure continuous compliance."
    ),
    "2g": (
        "The validation and testing procedures used, including information on the "
        "validation and testing data used with their main characteristics, and the "
        "outcomes of the validation and testing, including information on the "
        "identified shortcomings and how those shortcomings were corrected."
    ),
    "2h": (
        "The cybersecurity measures implemented, including documentation of tests "
        "carried out and their results."
    ),
    "3": (
        "Detailed information about the monitoring, functioning, and control of "
        "the AI system, in particular with regard to: its capabilities and "
        "limitations in performance, including the degrees of accuracy and the "
        "expected level of accuracy in relation to the persons or groups of "
        "persons on whom the system is intended to be used."
    ),
    "4": (
        "A description of the appropriateness of the performance metrics for "
        "the specific AI system."
    ),
    "5": (
        "A detailed description of the risk management system in accordance "
        "with Article 9 of this Regulation."
    ),
    "6": (
        "A description of any change made to the system throughout its "
        "lifecycle."
    ),
    "7": (
        "A list of the harmonised standards applied in full or in part; "
        "where no such harmonised standards have been applied, a description "
        "of the solutions adopted to satisfy the requirements set out in "
        "Chapter III, Section 2, including a list of other relevant standards "
        "and technical specifications applied."
    ),
    "8": (
        "A copy of the EU declaration of conformity referred to in Article 48."
    ),
    "9": (
        "A detailed description of the system in place to evaluate the AI "
        "system performance in the post-deployment phase in accordance with "
        "Article 72, including the post-market monitoring plan referred to "
        "in Article 72(3)."
    ),
}

_NARRATIVE_TEMPLATES: Dict[str, str] = {
    "1": (
        "The AI system has been documented for Technical File purposes under "
        "EU AI Act Article 11. General system description captured."
    ),
    "1a": (
        "System identity, version, intended purpose, and responsible parties "
        "have been documented as required."
    ),
    "1b": (
        "System integration points, hardware dependencies, and software interfaces "
        "have been assessed and documented."
    ),
    "1c": (
        "Software version information is captured in the AI Bill of Materials (AIBoM) "
        "using CycloneDX 1.6 ML-BOM format."
    ),
    "1d": (
        "Deployment forms and distribution channels have been recorded "
        "as part of the market placement documentation."
    ),
    "1e": (
        "Hardware requirements and the execution environment are documented "
        "in the AIBoM hardware context section."
    ),
    "1f": (
        "Physical form factor documentation is not applicable for "
        "software-only AI systems deployed via API or software package."
    ),
    "1g": (
        "Deployer instructions and interface descriptions are included "
        "in the model card and API documentation."
    ),
    "2": (
        "The development process, including training methodology, architectural "
        "decisions, and testing artefacts, has been documented."
    ),
    "2a": (
        "Development methodology, pre-trained base models used, and all "
        "third-party tools are captured in the AIBoM component inventory."
    ),
    "2b": (
        "Design specifications, model architecture, and algorithmic choices "
        "are documented in the model card and architecture description."
    ),
    "2c": (
        "System architecture and computational resources at scan time are "
        "captured in the AIBoM hardware context section."
    ),
    "2d": (
        "Training data governance, provenance, and data quality are evidenced "
        "by bias fairness and membership inference test results."
    ),
    "2e": (
        "Human oversight requirements have been assessed as part of "
        "alignment and sycophancy testing."
    ),
    "2f": (
        "Lifecycle change management procedures have been established. "
        "Differential redteaming provides evidence of change impact."
    ),
    "2g": (
        "Adversarial validation and testing has been conducted using SichGate Pro. "
        "Results constitute direct evidence for this Annex IV section."
    ),
    "2h": (
        "Cybersecurity measures have been validated through prompt injection, "
        "jailbreak, and evasion robustness testing via SichGate Pro."
    ),
    "3": (
        "Monitoring and control mechanisms have been validated through "
        "functional and adversarial testing."
    ),
    "4": (
        "Performance metrics appropriate to the system's intended use have "
        "been defined and evaluated as part of the test programme."
    ),
    "5": (
        "A risk management system has been established in accordance with "
        "Article 9. Adversarial test results constitute direct evidence."
    ),
    "6": (
        "Lifecycle changes are tracked. Differential redteaming provides "
        "a quantitative comparison of vulnerability profiles across versions."
    ),
    "7": (
        "Applicable harmonised standards and technical specifications have "
        "been identified. See compliance tag mapping in findings."
    ),
    "8": (
        "EU declaration of conformity to be completed upon satisfactory "
        "compliance review and sign-off by the responsible person."
    ),
    "9": (
        "Post-market monitoring procedures have been established, "
        "including scheduled re-testing using SichGate Pro."
    ),
}


# ---------------------------------------------------------------------------
# Data Models
# ---------------------------------------------------------------------------

@dataclass
class TechnicalFileSection:
    """A single Annex IV section with its mapped findings and compliance status."""

    section_id: str
    title: str
    requirement: str
    findings: List[dict] = field(default_factory=list)
    narrative: str = ""
    compliance_status: str = "NOT_APPLICABLE"  # COMPLIANT | MITIGATED | NOT_APPLICABLE | NON_COMPLIANT

    def to_dict(self) -> dict:
        return {
            "section_id": self.section_id,
            "title": self.title,
            "requirement": self.requirement,
            "findings": self.findings,
            "narrative": self.narrative,
            "compliance_status": self.compliance_status,
        }


@dataclass
class TechnicalFile:
    """Complete EU AI Act Article 11 Technical File artefact."""

    model_metadata: dict
    sections: Dict[str, TechnicalFileSection]
    executive_summary: dict
    timestamp: str
    test_run_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    differential_report: Optional[dict] = None
    attestation: Optional[dict] = None

    def to_dict(self) -> dict:
        return {
            "test_run_id": self.test_run_id,
            "timestamp": self.timestamp,
            "model_metadata": self.model_metadata,
            "differential_report": self.differential_report,
            "sections": {
                sid: s.to_dict() for sid, s in self.sections.items()
            },
            "attestation": self.attestation,
            "executive_summary": self.executive_summary,
        }

    def to_json(self, indent: int = 2) -> str:
        return json.dumps(self.to_dict(), indent=indent)

    def save(self, path: str) -> None:
        with open(path, "w", encoding="utf-8") as fh:
            fh.write(self.to_json())


# ---------------------------------------------------------------------------
# Builder
# ---------------------------------------------------------------------------

# Sections populated from AIBoM even when no attack findings exist
_AIBOM_SECTIONS = {"1", "1a", "1b", "1c", "1d", "1e", "2", "2a", "2b", "2c", "2d"}

# Sections populated by testing results
_TESTING_SECTIONS = {"2g", "2h", "3", "4", "5", "9"}


class TechnicalFileBuilder:
    """
    Orchestrates SichGate artefacts into a complete EU AI Act Annex IV
    Technical File.

    Args:
        findings: List of annotated attack result dicts. May already contain
            ``annex_iv_sections`` key; if not, it will be added automatically.
        aibom: CycloneDX 1.6 BOM dict (from ``AIBoMCapture.to_cyclonedx()``).
        differential_report: Phase 2.1 DifferentialReport dict (optional).
        attestation: Attestation block from ``CryptoAttestation.sign_findings_report()``
            (optional).
    """

    def __init__(
        self,
        findings: List[dict],
        aibom: Optional[dict] = None,
        differential_report: Optional[dict] = None,
        attestation: Optional[dict] = None,
    ) -> None:
        # Ensure every finding has annex_iv_sections annotation
        self._findings: List[dict] = [
            f if "annex_iv_sections" in f else annotate_finding(f)
            for f in findings
        ]
        self._aibom: dict = aibom or {}
        self._differential_report: Optional[dict] = differential_report
        self._attestation: Optional[dict] = attestation

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def build(self) -> TechnicalFile:
        """Orchestrate the complete Technical File construction."""
        sections = self._initialize_sections()
        self._populate_sections_from_findings(sections)
        self._compute_compliance_status(sections)
        self._generate_narratives(sections)
        executive_summary = self._create_executive_summary(sections)
        model_metadata = self._extract_model_metadata()

        return TechnicalFile(
            test_run_id=str(uuid.uuid4()),
            model_metadata=model_metadata,
            differential_report=self._differential_report,
            sections=sections,
            attestation=self._attestation,
            executive_summary=executive_summary,
            timestamp=datetime.now(timezone.utc).isoformat(),
        )

    # ------------------------------------------------------------------
    # Step 1 — Initialise all Annex IV sections
    # ------------------------------------------------------------------

    def _initialize_sections(self) -> Dict[str, TechnicalFileSection]:
        """Create all Annex IV sections with their EU AI Act requirement text."""
        sections: Dict[str, TechnicalFileSection] = {}
        for section_id, title in ANNEX_IV_SECTIONS.items():
            sections[section_id] = TechnicalFileSection(
                section_id=section_id,
                title=title,
                requirement=_SECTION_REQUIREMENTS.get(section_id, title),
                findings=[],
                narrative="",
                compliance_status="NOT_APPLICABLE",
            )
        return sections

    # ------------------------------------------------------------------
    # Step 2 — Populate sections from findings
    # ------------------------------------------------------------------

    def _populate_sections_from_findings(
        self, sections: Dict[str, TechnicalFileSection]
    ) -> None:
        """Map each annotated finding to its relevant Annex IV sections."""
        for finding in self._findings:
            annex_sections = finding.get("annex_iv_sections") or get_sections_for_attack(
                finding.get("attack_type", "")
            )
            for sid in annex_sections:
                if sid in sections:
                    # Avoid duplicates if a finding maps to the same section twice
                    existing_ids = {f.get("attack_name") for f in sections[sid].findings}
                    if finding.get("attack_name") not in existing_ids:
                        sections[sid].findings.append(finding)

    # ------------------------------------------------------------------
    # Step 3 — Compute compliance status per section
    # ------------------------------------------------------------------

    def _compute_compliance_status(
        self, sections: Dict[str, TechnicalFileSection]
    ) -> None:
        """
        Compute per-section compliance status:

        - COMPLIANT: all findings passed, or no findings with AIBoM/testing coverage
        - MITIGATED: findings failed but all carry a mitigated/accepted status
        - NON_COMPLIANT: unaddressed failing findings present
        - NOT_APPLICABLE: no coverage and no relevant artefacts
        """
        for section_id, section in sections.items():
            findings = section.findings

            if not findings:
                if section_id in _AIBOM_SECTIONS and self._aibom:
                    section.compliance_status = "COMPLIANT"
                elif section_id in _TESTING_SECTIONS:
                    # Testing section with no findings means all attacks passed
                    section.compliance_status = "COMPLIANT"
                else:
                    section.compliance_status = "NOT_APPLICABLE"
                continue

            failing = [
                f for f in findings
                if f.get("status") in ("fail", "warning")
            ]

            if not failing:
                section.compliance_status = "COMPLIANT"
                continue

            # Check if all failures carry a resolution status
            all_resolved = all(
                f.get("finding_status") in ("MITIGATED", "ACCEPTED_RISK")
                for f in failing
            )
            section.compliance_status = "MITIGATED" if all_resolved else "NON_COMPLIANT"

    # ------------------------------------------------------------------
    # Step 4 — Generate narratives
    # ------------------------------------------------------------------

    def _generate_narratives(
        self, sections: Dict[str, TechnicalFileSection]
    ) -> None:
        """Generate a human-readable narrative for each section."""
        for section_id, section in sections.items():
            base = _NARRATIVE_TEMPLATES.get(section_id, "")
            total = len(section.findings)
            failing = sum(
                1 for f in section.findings if f.get("status") in ("fail", "warning")
            )
            critical = sum(
                1 for f in section.findings if f.get("severity") == "critical"
            )

            if total == 0:
                section.narrative = base
            elif failing == 0:
                section.narrative = (
                    f"{base} {total} test(s) conducted — all passed."
                )
            else:
                crit_note = (
                    f", including {critical} critical finding(s)" if critical else ""
                )
                section.narrative = (
                    f"{base} {total} test(s) conducted: "
                    f"{failing} finding(s) identified{crit_note}."
                )

    # ------------------------------------------------------------------
    # Step 5 — Executive summary
    # ------------------------------------------------------------------

    def _create_executive_summary(
        self, sections: Dict[str, TechnicalFileSection]
    ) -> dict:
        """Create a one-page executive summary of the Technical File."""
        total_findings = len(self._findings)
        fail_findings = [
            f for f in self._findings if f.get("status") in ("fail", "warning")
        ]
        critical_findings = [
            f for f in self._findings if f.get("severity") == "critical"
        ]
        high_findings = [
            f for f in self._findings if f.get("severity") == "high"
        ]

        status_counts = {
            "COMPLIANT": 0,
            "MITIGATED": 0,
            "NON_COMPLIANT": 0,
            "NOT_APPLICABLE": 0,
        }
        for s in sections.values():
            status_counts[s.compliance_status] = (
                status_counts.get(s.compliance_status, 0) + 1
            )

        if status_counts["NON_COMPLIANT"] == 0 and not critical_findings:
            overall_status = "COMPLIANT"
        elif status_counts["NON_COMPLIANT"] == 0 and critical_findings:
            overall_status = "REQUIRES_REVIEW"
        else:
            overall_status = "NON_COMPLIANT"

        model_name = self._extract_model_name()

        return {
            "model_name": model_name,
            "total_findings": total_findings,
            "fail_count": len(fail_findings),
            "critical_count": len(critical_findings),
            "high_count": len(high_findings),
            "overall_compliance_status": overall_status,
            "section_status_counts": status_counts,
            "total_sections": len(sections),
            "has_attestation": self._attestation is not None,
            "has_aibom": bool(self._aibom),
            "has_differential": self._differential_report is not None,
            "regulation": "EU AI Act (Regulation (EU) 2024/1689)",
            "article": "Article 11 — Technical Documentation",
            "annex": "Annex IV — Technical Documentation",
        }

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _extract_model_metadata(self) -> dict:
        """Extract model metadata from the AIBoM or return placeholders."""
        if not self._aibom:
            return {"name": "Unknown", "version": "unknown", "hash": None}

        metadata = self._aibom.get("metadata", {})
        component = metadata.get("component", {})

        model_hash: Optional[str] = None
        for h in component.get("hashes", []):
            if h.get("alg") == "SHA-256":
                model_hash = h.get("content")
                break

        hardware_props = [
            p for p in self._aibom.get("properties", [])
            if p.get("name", "").startswith("sichgate:hardware:")
        ]

        return {
            "name": component.get("name", "Unknown"),
            "version": component.get("version", "unknown"),
            "hash": model_hash,
            "timestamp": metadata.get("timestamp"),
            "hardware_properties": hardware_props,
            "bom_serial": self._aibom.get("serialNumber"),
        }

    def _extract_model_name(self) -> str:
        if not self._aibom:
            return "Unknown"
        component = self._aibom.get("metadata", {}).get("component", {})
        return component.get("name", "Unknown")
