"""
Attestation report formatting and visualization for fhenomai.

This module provides formatting capabilities for TEE attestation reports,
including terminal output, HTML, and PDF generation. It was moved from
dk-tee-attestation to fhenomai to maintain stable formatting APIs while
the core attestation library focuses on verification logic.
"""
from __future__ import annotations
import re
from typing import Optional, Dict, Any
from pathlib import Path

# Import the rust wrapper from dk_tee_attestation
from dk_tee_attestation import tee_amd_sev_snp_rust_py_wrap

# ANSI color codes for terminal output
BLACK = "\033[30m"
RED = "\033[31m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
BLUE = "\033[34m"
MAGENTA = "\033[35m"
CYAN = "\033[36m"
WHITE = "\033[37m"
RESET = "\033[0m"
BOLD = "\033[1m"

# Color mapping for report sections
SECTION_COLORS = {
    'nonce': GREEN,
    'chip_id': CYAN,
    'tcb': YELLOW,
    'cpuid': BLUE,
    'signature': MAGENTA,
}

# AMD SEV-SNP report byte offsets
REPORT_SECTIONS = {
    'report_data': {'offset': 0x50, 'length': 0x40, 'color': 'nonce'},
    'reported_tcb': {'offset': 0x180, 'length': 0x08, 'color': 'tcb'},
    'cpuid': {'offset': 0x188, 'length': 0x03, 'color': 'cpuid'},
    'chip_id': {'offset': 0x1A0, 'length': 0x40, 'color': 'chip_id'},
    'signature': {'offset': 0x2A0, 'length': 0x200, 'color': 'signature'},
}

# Specific field offsets within sections
CPUID_FAM_ID_OFF = 0x188
CPUID_MOD_ID_OFF = 0x189
CPUID_STEP_OFF = 0x18A

# Legend for colored output
LEGEND = [
    ("User challenge (nonce)", "nonce"),
    ("Hardware identifier (chip id)", "chip_id"),
    ("Trusted computing base informations (tcb)", "tcb"),
    ("CPU family, model, stepping informations (cpuid)", "cpuid"),
    ("Elliptic Curve Digital Signature Algorithm (ECDSA)", "signature"),
]

# AMD KDS constants
KDS_CERT_SITE = "https://kdsintf.amd.com"
KDS_VCEK_PATH = "/vcek/v1"
KDS_CERT_CHAIN_PATH = "cert_chain"

# HTML template for report visualization
HTML_TEMPLATE = """<!doctype html>
<html>
<head>
  <meta charset="utf-8"/>
  <title>AMD SEV-SNP Attestation Report</title>
  <style>
    /* White page background */
    body {{
      margin: 0;
      padding: 24px;
      background: #ffffff;
      color: #111;
      font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Arial, sans-serif;
    }}

    /* Header area */
    .meta {{
      max-width: 900px;
      margin: 0 auto 14px auto;
      font-size: 13px;
      line-height: 1.4;
    }}

    .meta-link {{
      text-align: center;
      margin-top: 6px;
    }}
    
    .meta a {{
      color: #0b57d0;
      text-decoration: underline;
      word-break: break-all;
    }}

    /* Terminal "photo" */
    .terminal {{
      width: fit-content;
      max-width: 900px;
      margin: 0 auto;
      background: #0C0C0C;
      border: 1px solid #222;
      border-radius: 12px;
      padding: 16px 18px;
      box-shadow: 0 10px 25px rgba(0,0,0,0.12);
    }}

    /* Pre-formatted terminal output */
    pre {{
      max-width: 100%;
      margin: 0;
      white-space: pre-wrap;
      word-break: break-word;
      font-family: "JetBrains Mono", "Fira Mono", ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
      font-size: 9pt;
      line-height: 1.15;
      color: #e6e6e6;
    }}

    /* Color classes matching ANSI codes */
    .color-green  {{ color: rgb(19, 153, 14) !important; }}
    .color-cyan   {{ color: rgb(133, 201, 255) !important; }}
    .color-yellow {{ color: rgb(255, 219, 74) !important; }}
    .color-blue   {{ color: rgb(0, 53, 209) !important; }}
    .color-magenta {{ color: rgb(133, 27, 148) !important; }}
    .color-white  {{ color: rgb(211, 215, 207) !important; }}

    /* Print optimization */
    @media print {{
      body {{
        padding: 0;
      }}
      .terminal {{
        box-shadow: none;
        border: 1px solid #222;
        page-break-inside: avoid;
      }}
    }}

    @page {{
      size: A4;
      margin: 10mm;
    }}
  </style>
</head>

<body>
  <div class="meta">
    <div><strong>AMD SEV-SNP Attestation Report Verification</strong></div>
    <div class="meta-link">
      For more information, see: <a href="https://github.com/virtee/snpguest">https://github.com/virtee/snpguest</a>
    </div>
  </div>

  <div class="terminal">
    <pre>{content}</pre>
  </div>
</body>

</html>
"""


class ParsedAttestationReport:
    """Parsed attestation report structure."""
    
    def __init__(self, chip_id: bytes, report_data: bytes, reported_tcb: Dict, cpuid: Dict,
                 signed_bytes: bytes, sig_r: bytes, sig_s: bytes):
        self.chip_id = chip_id
        self.report_data = report_data
        self.reported_tcb = type('ReportedTcb', (), reported_tcb)()
        self.cpuid = type('CpuId', (), cpuid)()
        self.signed_bytes = signed_bytes
        self.sig_r = sig_r
        self.sig_s = sig_s


class AttestationReportFormatter:
    """
    Format AMD SEV-SNP attestation reports with visual styling.
    
    Supports multiple output formats:
    - standard: Structured fields (parsed data)
    - detailed: ANSI-colored terminal output with hex dump and highlights
    - html: Standalone HTML file with embedded CSS and dark terminal theme
    - pdf: PDF document generated from HTML
    """
    
    def __init__(self, line_width: int = 64):
        """
        Initialize formatter.
        
        Args:
            line_width: Character width for formatting (default: 64)
        """
        self.line_width = line_width
        self._ansi_re = re.compile(r"\x1b\[[0-9;]*m")
    
    def _parse_report(self, report_bytes: bytes) -> ParsedAttestationReport:
        """
        Parse raw report bytes into structured data.
        
        Args:
            report_bytes: Raw attestation report bytes
            
        Returns:
            ParsedAttestationReport object
            
        Raises:
            ValueError: If parsing fails
        """
        parsed_report = tee_amd_sev_snp_rust_py_wrap.parse_report(report_bytes)
        
        return ParsedAttestationReport(
            chip_id=parsed_report["chip_id"],
            report_data=parsed_report["report_data"],
            reported_tcb=parsed_report["reported_tcb"],
            cpuid=parsed_report["cpuid"],
            signed_bytes=parsed_report["signed_bytes"],
            sig_r=parsed_report["sig_r"],
            sig_s=parsed_report["sig_s"],
        )
    
    def format_standard(
        self,
        report_bytes: bytes,
        verification_result: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Format report as structured fields (standard output).
        
        Args:
            report_bytes: Raw attestation report bytes
            verification_result: Optional verification details
            
        Returns:
            Formatted string with parsed fields
        """
        try:
            parsed = self._parse_report(report_bytes)
        except Exception as e:
            return f"{YELLOW}Warning: Could not parse report: {e}{RESET}\n"
        
        output = []
        output.append(f"\n{BOLD}{GREEN}=== Attestation Report ==={RESET}\n")
        
        # CPU Information
        output.append(f"{BOLD}CPU Information:{RESET}")
        output.append(f"  Family:    {BLUE}0x{parsed.cpuid.fam_id:02x}{RESET}")
        output.append(f"  Model:     {BLUE}0x{parsed.cpuid.mod_id:02x}{RESET}")
        output.append(f"  Stepping:  {BLUE}0x{parsed.cpuid.step:02x}{RESET}\n")
        
        # TCB
        output.append(f"{BOLD}Trusted Computing Base (TCB):{RESET}")
        output.append(f"  Bootloader: {YELLOW}{parsed.reported_tcb.bootloader}{RESET}")
        output.append(f"  TEE:        {YELLOW}{parsed.reported_tcb.tee}{RESET}")
        output.append(f"  SNP:        {YELLOW}{parsed.reported_tcb.snp}{RESET}")
        output.append(f"  Microcode:  {YELLOW}{parsed.reported_tcb.microcode}{RESET}\n")
        
        # Chip ID
        output.append(f"{BOLD}Chip ID:{RESET}")
        chip_id_hex = parsed.chip_id.hex()
        chip_id_formatted = '\n  '.join([chip_id_hex[i:i+64] for i in range(0, len(chip_id_hex), 64)])
        output.append(f"  {CYAN}{chip_id_formatted}{RESET}\n")
        
        # Report Data (Nonce)
        output.append(f"{BOLD}Report Data (Nonce):{RESET}")
        report_data_hex = parsed.report_data.hex()
        report_data_formatted = '\n  '.join([report_data_hex[i:i+64] for i in range(0, len(report_data_hex), 64)])
        output.append(f"  {GREEN}{report_data_formatted}{RESET}\n")
        
        # Signature
        output.append(f"{BOLD}ECDSA Signature:{RESET}")
        sig_r_hex = parsed.sig_r.hex()
        sig_s_hex = parsed.sig_s.hex()
        output.append(f"  R: {MAGENTA}{sig_r_hex[:64]}{RESET}")
        output.append(f"     {MAGENTA}{sig_r_hex[64:128]}{RESET}")
        if len(sig_r_hex) > 128:
            output.append(f"     {MAGENTA}{sig_r_hex[128:]}{RESET}")
        output.append(f"  S: {MAGENTA}{sig_s_hex[:64]}{RESET}")
        output.append(f"     {MAGENTA}{sig_s_hex[64:128]}{RESET}")
        if len(sig_s_hex) > 128:
            output.append(f"     {MAGENTA}{sig_s_hex[128:]}{RESET}")
        output.append("")
        
        # Verification result
        if verification_result:
            if verification_result.get('verified'):
                output.append(f"{GREEN}✓ Attestation VERIFIED{RESET}")
                output.append(f"  Platform: {verification_result.get('platform', 'AMD SEV-SNP')}")
            else:
                output.append(f"{RED}✗ Attestation FAILED{RESET}")
                if verification_result.get('error'):
                    output.append(f"  Error: {verification_result['error']}")
        
        return "\n".join(output)
    
    def format_detailed(
        self, 
        report_bytes: bytes, 
        nonce: Optional[bytes] = None,
        verification_result: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Format report with ANSI colors and detailed breakdown.
        
        Args:
            report_bytes: Raw attestation report bytes
            nonce: Optional nonce used for generation (for display)
            verification_result: Optional verification details from verify_attestation
            
        Returns:
            ANSI-colored formatted string
        """
        output = []
        
        # Legend
        output.append(self._format_legend())
        
        # Report hex dump with highlights
        output.append(self._centered_title("Attestation report bytes (hex)"))
        output.append(self._format_hex_dump(report_bytes))
        output.append("-" * self.line_width + "\n")
        
        # CPU family derivation
        cpu_info = self._extract_cpuid(report_bytes)
        output.append(self._centered_title("Deriving CPU family name"))
        cpuid_str = f"{BLUE}{cpu_info['fam_id']:02x}{cpu_info['mod_id']:02x}{cpu_info['step']:02x}{RESET}"
        output.append(self._center_text(f"{cpuid_str} ---> {BLUE}{cpu_info['family']}{RESET}"))
        output.append("-" * self.line_width + "\n")
        
        # AMD KDS URLs
        if verification_result:
            output.append(self._format_verification_steps(report_bytes, cpu_info, verification_result))
        
        return "\n".join(output)
    
    def format_html(
        self, 
        report_bytes: bytes, 
        nonce: Optional[bytes] = None,
        verification_result: Optional[Dict[str, Any]] = None
    ) -> str:
        """
        Format report as standalone HTML document.
        
        Args:
            report_bytes: Raw attestation report bytes
            nonce: Optional nonce used for generation
            verification_result: Optional verification details
            
        Returns:
            Complete HTML document as string
        """
        # Generate ANSI-colored content
        ansi_content = self.format_detailed(report_bytes, nonce, verification_result)
        
        # Convert ANSI to HTML spans
        html_content = self._ansi_to_html(ansi_content)
        
        # Wrap in template
        return HTML_TEMPLATE.format(content=html_content)
    
    def format_pdf(
        self, 
        report_bytes: bytes, 
        nonce: Optional[bytes] = None,
        verification_result: Optional[Dict[str, Any]] = None,
        output_path: Optional[str] = None
    ) -> bytes:
        """
        Format report as PDF document using ReportLab (pure Python).
        
        Args:
            report_bytes: Raw attestation report bytes
            nonce: Optional nonce used for generation
            verification_result: Optional verification details
            output_path: Optional path to save PDF (if None, returns bytes)
            
        Returns:
            PDF bytes if output_path is None, otherwise None
        """
        try:
            from reportlab.lib.pagesizes import letter
            from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
            from reportlab.lib.units import inch
            from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Preformatted
            from reportlab.lib.enums import TA_CENTER
            from io import BytesIO
        except ImportError:
            raise ImportError(
                "PDF generation requires reportlab. "
                "Install with: pip install reportlab"
            )
        
        # Get formatted text (without ANSI codes)
        text_content = self.format_standard(report_bytes, verification_result)
        # Remove ANSI color codes
        text_content = re.sub(r'\x1b\[[0-9;]*m', '', text_content)
        
        # Create PDF
        output_stream = BytesIO()
        doc = SimpleDocTemplate(
            output_stream,
            pagesize=letter,
            rightMargin=0.75*inch,
            leftMargin=0.75*inch,
            topMargin=0.75*inch,
            bottomMargin=0.75*inch
        )
        
        # Build content
        story = []
        styles = getSampleStyleSheet()
        
        # Title
        title_style = ParagraphStyle(
            'CustomTitle',
            parent=styles['Heading1'],
            fontSize=18,
            textColor='#1e90ff',
            spaceAfter=20,
            alignment=TA_CENTER
        )
        story.append(Paragraph("AMD SEV-SNP Attestation Report", title_style))
        story.append(Spacer(1, 0.3*inch))
        
        # Content in monospace
        code_style = ParagraphStyle(
            'Code',
            parent=styles['Code'],
            fontSize=9,
            fontName='Courier',
            leftIndent=0,
            rightIndent=0
        )
        
        # Split content into lines and add each
        for line in text_content.split('\n'):
            if line.strip():
                # Escape special characters for ReportLab
                line = line.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
                story.append(Preformatted(line, code_style))
            else:
                story.append(Spacer(1, 0.1*inch))
        
        # Build PDF
        doc.build(story)
        pdf_bytes = output_stream.getvalue()
        output_stream.close()
        
        if output_path:
            Path(output_path).write_bytes(pdf_bytes)
            return None
        else:
            return pdf_bytes
    
    # =========================================================================
    # Private helper methods
    # =========================================================================
    
    def _format_legend(self) -> str:
        """Generate legend for color-coded sections."""
        output = []
        output.append(self._centered_title("Legend"))
        for label, color_key in LEGEND:
            color = SECTION_COLORS.get(color_key, WHITE)
            output.append(self._center_text(f"{color}▇▇▇{RESET}  {label}"))
        output.append("-" * self.line_width + "\n")
        return "\n".join(output)
    
    def _format_hex_dump(self, report_bytes: bytes) -> str:
        """Format report as color-coded hex dump."""
        output = []
        for offset in range(0, len(report_bytes), 16):
            chunk = report_bytes[offset:offset+16]
            hex_str = ""
            for i, byte in enumerate(chunk):
                byte_offset = offset + i
                color = self._get_color_for_offset(byte_offset)
                hex_str += f"{color}{byte:02x}{RESET} "
            output.append(hex_str.rstrip())
        return "\n".join(output)
    
    def _get_color_for_offset(self, offset: int) -> str:
        """Get color for byte offset based on report structure."""
        for section_name, section_info in REPORT_SECTIONS.items():
            start = section_info['offset']
            end = start + section_info['length']
            if start <= offset < end:
                color_key = section_info['color']
                return SECTION_COLORS.get(color_key, WHITE)
        return WHITE
    
    def _extract_cpuid(self, report_bytes: bytes) -> Dict[str, Any]:
        """Extract and decode CPU ID information."""
        fam_id = report_bytes[CPUID_FAM_ID_OFF]
        mod_id = report_bytes[CPUID_MOD_ID_OFF]
        step = report_bytes[CPUID_STEP_OFF]
        
        # CPU family name mapping (simplified)
        cpuid_hex = f"{fam_id:02x}{mod_id:02x}{step:02x}"
        family_map = {
            "190101": "Milan",
            "190001": "Rome",
            "190111": "Genoa",
        }
        family = family_map.get(cpuid_hex, f"Unknown (0x{cpuid_hex})")
        
        return {
            'fam_id': fam_id,
            'mod_id': mod_id,
            'step': step,
            'family': family,
        }
    
    def _format_verification_steps(
        self, 
        report_bytes: bytes, 
        cpu_info: Dict[str, Any], 
        verification_result: Dict[str, Any]
    ) -> str:
        """Format verification steps with URLs."""
        output = []
        
        output.append(self._centered_title("AMD KDS Certificate URLs"))
        
        # Extract chip ID for URL
        chip_id_bytes = report_bytes[REPORT_SECTIONS['chip_id']['offset']:
                                    REPORT_SECTIONS['chip_id']['offset'] + 
                                    REPORT_SECTIONS['chip_id']['length']]
        chip_id_hex = chip_id_bytes.hex()
        
        # VCEK URL
        vcek_url = (f"{KDS_CERT_SITE}{KDS_VCEK_PATH}/"
                   f"{cpu_info['family']}/"
                   f"{chip_id_hex}")
        output.append(self._center_text(f"{CYAN}VCEK:{RESET}", wrap=True))
        output.append(self._center_text(vcek_url, wrap=True))
        output.append("")
        
        # Cert chain URL
        cert_url = f"{KDS_CERT_SITE}{KDS_VCEK_PATH}/{cpu_info['family']}/{KDS_CERT_CHAIN_PATH}"
        output.append(self._center_text(f"{CYAN}Cert Chain:{RESET}", wrap=True))
        output.append(self._center_text(cert_url, wrap=True))
        output.append("-" * self.line_width + "\n")
        
        # TCB info
        output.append(self._centered_title("TCB Security Version Numbers"))
        tcb_bytes = report_bytes[REPORT_SECTIONS['reported_tcb']['offset']:
                                REPORT_SECTIONS['reported_tcb']['offset'] + 
                                REPORT_SECTIONS['reported_tcb']['length']]
        
        bootloader_spl = tcb_bytes[0]
        tee_spl = tcb_bytes[1]
        snp_spl = tcb_bytes[4]
        ucode_spl = tcb_bytes[6]
        
        output.append(self._center_text(f"bootloader: {YELLOW}{bootloader_spl:02}{RESET}"))
        output.append(self._center_text(f"tee: {YELLOW}{tee_spl:02}{RESET}"))
        output.append(self._center_text(f"snp: {YELLOW}{snp_spl:02}{RESET}"))
        output.append(self._center_text(f"microcode: {YELLOW}{ucode_spl:02}{RESET}"))
        output.append("-" * self.line_width + "\n")
        
        # Verification results
        if verification_result.get('verified'):
            output.append(self._centered_title("Verification Result"))
            output.append(self._center_text(f"{GREEN}✓ Attestation verified successfully!{RESET}"))
            output.append(self._center_text(f"Platform: {verification_result.get('platform', 'AMD SEV-SNP')}"))
            output.append("-" * self.line_width + "\n")
        
        return "\n".join(output)
    
    def _ansi_to_html(self, ansi_text: str) -> str:
        """Convert ANSI color codes to HTML spans."""
        import html as html_module
        
        # First escape HTML entities in the text
        escaped = html_module.escape(ansi_text)
        
        # Then replace ANSI codes with actual HTML tags (unescape the markers)
        ansi_to_html_map = {
            '\\033[32m': '<span class="color-green">',
            '\\033[36m': '<span class="color-cyan">',
            '\\033[33m': '<span class="color-yellow">',
            '\\033[34m': '<span class="color-blue">',
            '\\033[35m': '<span class="color-magenta">',
            '\\033[37m': '<span class="color-white">',
            '\\033[0m': '</span>',
            '\\033[1m': '<strong>',
        }
        
        result = escaped
        for ansi_code, html_tag in ansi_to_html_map.items():
            # The escape() turned \033 into literal text, so match that
            result = result.replace(ansi_code.replace('\\033', '\x1b'), html_tag)
        
        # Clean up any remaining ANSI codes  
        result = self._ansi_re.sub('', result)
        
        return result
    
    def _centered_title(self, title: str, fill: str = "-") -> str:
        """Create centered title with fill characters."""
        title = f" {title} "
        if len(title) >= self.line_width:
            return title[:self.line_width]
        return title.center(self.line_width, fill)
    
    def _center_text(self, text: str, wrap: bool = False) -> str:
        """Center text, optionally wrapping long lines."""
        visible_len = len(self._ansi_re.sub('', text))
        
        if wrap and visible_len > self.line_width:
            # Simple wrap for long URLs
            lines = []
            current_line = ""
            for char in text:
                current_line += char
                if len(self._ansi_re.sub('', current_line)) >= self.line_width:
                    lines.append(current_line)
                    current_line = ""
            if current_line:
                lines.append(current_line)
            return "\n".join(self._center_text(line) for line in lines)
        
        pad = self.line_width - visible_len
        if pad <= 0:
            return text
        
        left = pad // 2
        right = pad - left
        return " " * left + text + " " * right
    
    def _strip_ansi(self, text: str) -> str:
        """Remove ANSI escape codes."""
        return self._ansi_re.sub('', text)
