# Copyright 2026 Gaofeng Fan
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
SchematicBook - Main class for generating PDF schematic documentation.
"""

import io
import os
import tempfile
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, TYPE_CHECKING

from reportlab.lib.pagesizes import letter, A4, landscape
from reportlab.lib.units import inch
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.enums import TA_CENTER, TA_LEFT
from reportlab.platypus import (
    SimpleDocTemplate, Paragraph, Spacer, PageBreak,
    Table, TableStyle, Image
)
from reportlab.lib import colors

from pypdf import PdfReader, PdfWriter

from .config import SymbolStyle, SymbolConfig, LayoutHints, BookConfig
from .symbols import SymbolRenderer, create_cell_schematic, create_testbench_block_diagram, get_default_renderer


if TYPE_CHECKING:
    from ..circuit import Circuit
    from ..testbench import Testbench


class SchematicBook:
    """
    Generates PDF schematic documentation from analogpy circuits.

    The PDF starts directly with the testbench block diagram on page 1,
    with table of contents available in the PDF sidebar (bookmarks).

    Example:
        from analogpy.visualization import SchematicBook

        book = SchematicBook(title="OLED Design")
        book.add_testbench(tb)
        book.render("schematic.pdf")
    """

    def __init__(
        self,
        title: str = "Schematic Documentation",
        author: str = "",
        default_style: SymbolStyle = SymbolStyle.DETAILED,
        page_size: str = "letter",  # "letter" or "a4"
        orientation: str = "landscape"  # "landscape" or "portrait"
    ):
        self.config = BookConfig(
            title=title,
            author=author,
            default_style=default_style,
        )
        base_size = letter if page_size == "letter" else A4
        self.page_size = landscape(base_size) if orientation == "landscape" else base_size
        self.renderer = get_default_renderer()
        self.renderer.default_style = default_style

        # Content to render
        self.testbenches: List['Testbench'] = []
        self.circuits: List['Circuit'] = []
        self.circuit_map: Dict[str, 'Circuit'] = {}  # name -> Circuit

        # Track page numbers for bookmarks
        self._page_bookmarks: List[tuple] = []  # [(title, page_num), ...]

        # Track temp files to clean up after build
        self._temp_files: List[str] = []

        # Styles
        self.styles = getSampleStyleSheet()
        self._setup_styles()

    def _setup_styles(self):
        """Setup custom paragraph styles."""
        self.styles.add(ParagraphStyle(
            name='ChapterTitle',
            parent=self.styles['Heading1'],
            fontSize=16,
            spaceBefore=6,
            spaceAfter=12,
        ))
        self.styles.add(ParagraphStyle(
            name='SectionTitle',
            parent=self.styles['Heading2'],
            fontSize=12,
            spaceBefore=8,
            spaceAfter=6,
        ))
        self.styles.add(ParagraphStyle(
            name='CellInfo',
            parent=self.styles['Normal'],
            fontSize=9,
            leftIndent=0,
            alignment=TA_LEFT,
        ))
        self.styles.add(ParagraphStyle(
            name='SmallInfo',
            parent=self.styles['Normal'],
            fontSize=8,
            textColor=colors.gray,
        ))

    def register_model_style(
        self,
        model: str,
        style: SymbolStyle = SymbolStyle.STANDARD,
        label: Optional[str] = None,
        color: str = "#E8E8E8",
        template: Optional[str] = None,
        **kwargs
    ):
        """Register a custom symbol configuration for a model."""
        config = SymbolConfig(
            style=style,
            label=label,
            color=color,
            template=template,
            **kwargs
        )
        self.renderer.register_model(model, config)
        self.config.model_configs[model] = config

    def set_layout(self, circuit_name: str, hints: LayoutHints):
        """Set layout hints for a specific circuit."""
        self.config.layout_hints[circuit_name] = hints

    def add_testbench(self, testbench: 'Testbench', depth: int = 100):
        """
        Add a testbench and its subcircuit hierarchy to the book.

        Args:
            testbench: The testbench to add
            depth: Maximum recursion depth for subcircuits (default 100 for "full")
        """
        self.testbenches.append(testbench)
        self._collect_subcircuits(testbench, depth)

    def _collect_subcircuits(self, circuit_or_tb, depth: int):
        """Recursively collect subcircuits from instances."""
        if depth <= 0:
            return

        instances = getattr(circuit_or_tb, 'subcircuit_instances', [])
        for inst in instances:
            subckt = getattr(inst, 'subcircuit', None)
            if not subckt:
                continue
            
            # Skip primitives (no internal structure or explicit flag)
            is_primitive = getattr(subckt, '_is_primitive', False)
            has_internals = hasattr(subckt, 'subcircuit_instances') and len(subckt.subcircuit_instances) > 0
            
            if is_primitive or not has_internals:
                continue

            if subckt.name not in self.circuit_map:
                self.circuit_map[subckt.name] = subckt
                self.circuits.append(subckt)
                # Recurse into this subcircuit
                self._collect_subcircuits(subckt, depth - 1)

    def add_circuit(self, circuit: 'Circuit', depth: int = 100):
        """Add a standalone circuit and its hierarchy to the book."""
        if circuit.name not in self.circuit_map:
            # Check if it's a primitive before adding
            is_primitive = getattr(circuit, '_is_primitive', False)
            has_internals = hasattr(circuit, 'subcircuit_instances') and len(circuit.subcircuit_instances) > 0
            
            if not is_primitive and has_internals:
                self.circuit_map[circuit.name] = circuit
                self.circuits.append(circuit)
                self._collect_subcircuits(circuit, depth)

    def _render_drawing_to_image(self, drawing, width_inches: float = 6.0, max_height_inches: float = 7.0) -> Optional[Image]:
        """Convert a schemdraw Drawing to a reportlab element with size constraints."""
        try:
            # Try SVG first for searchable text
            try:
                from svglib.svglib import svg2rlg
                from reportlab.graphics import renderPDF
                from reportlab.platypus import Flowable
                import math

                class SVGFlowable(Flowable):
                    def __init__(self, svg_drawing, width, height):
                        super().__init__()
                        self.svg_drawing = svg_drawing
                        self.width = width
                        self.height = height

                    def wrap(self, availWidth, availHeight):
                        return self.width, self.height

                    def draw(self):
                        dw = getattr(self.svg_drawing, "width", 0) or 1
                        dh = getattr(self.svg_drawing, "height", 0) or 1
                        sx = self.width / dw
                        sy = self.height / dh
                        s = min(sx, sy)
                        self.canv.saveState()
                        self.canv.scale(s, s)
                        renderPDF.draw(self.svg_drawing, self.canv, 0, 0)
                        self.canv.restoreState()

                fd, temp_path = tempfile.mkstemp(suffix='.svg')
                os.close(fd)
                
                # Ensure matplotlib saves text as text nodes, not paths
                import matplotlib
                matplotlib.rcParams['svg.fonttype'] = 'none'
                
                drawing.save(temp_path)
                self._temp_files.append(temp_path)

                rlg = svg2rlg(temp_path)
                if rlg:
                    img_width = getattr(rlg, "width", None) or 1
                    img_height = getattr(rlg, "height", None) or 1
                    aspect_ratio = img_height / img_width

                    target_width = width_inches * inch
                    target_height = target_width * aspect_ratio
                    if target_height > max_height_inches * inch:
                        target_height = max_height_inches * inch
                        target_width = target_height / aspect_ratio

                    return SVGFlowable(rlg, width=target_width, height=target_height)
            except Exception:
                pass

            from PIL import Image as PILImage

            # Create temp file for PNG (don't delete yet - reportlab needs it)
            fd, temp_path = tempfile.mkstemp(suffix='.png')
            os.close(fd)

            # Save drawing to PNG
            drawing.save(temp_path, dpi=150)

            # Track for cleanup later
            self._temp_files.append(temp_path)

            # Get actual image dimensions
            with PILImage.open(temp_path) as pil_img:
                img_width, img_height = pil_img.size

            # Calculate aspect ratio
            aspect_ratio = img_height / img_width

            # Calculate dimensions that fit within constraints
            target_width = width_inches * inch
            target_height = target_width * aspect_ratio

            # If too tall, scale down to fit max height
            if target_height > max_height_inches * inch:
                target_height = max_height_inches * inch
                target_width = target_height / aspect_ratio

            # Create reportlab Image with both width and height
            img = Image(temp_path, width=target_width, height=target_height)

            return img
        except Exception as e:
            print(f"Warning: Could not render drawing: {e}")
            return None

    def _cleanup_temp_files(self):
        """Clean up temporary image files."""
        for temp_path in self._temp_files:
            try:
                os.unlink(temp_path)
            except:
                pass
        self._temp_files = []

    def _create_legend_page(self, page_num: int) -> List:
        """Create a page for the color legend."""
        from .symbols import generate_legend
        elements = []

        # Record bookmark
        self._page_bookmarks.append(("Schematic Legend", page_num))

        elements.append(Paragraph("Schematic Legend", self.styles['ChapterTitle']))
        elements.append(Paragraph(
            "This legend explains the color scheme and symbols used in this document.",
            self.styles['Normal']
        ))
        elements.append(Spacer(1, 0.3 * inch))

        try:
            drawing = generate_legend()
            img = self._render_drawing_to_image(drawing, width_inches=5.0)
            if img:
                elements.append(img)
        except Exception as e:
            elements.append(Paragraph(f"[Legend error: {e}]", self.styles['Normal']))

        return elements

    def _create_testbench_page(self, testbench: 'Testbench', page_num: int) -> List:
        """Create a page for a testbench (this is page 1)."""
        elements = []

        # Record bookmark
        self._page_bookmarks.append((f"Testbench: {testbench.name}", page_num))

        # Title with testbench name prominently
        elements.append(Paragraph(f"Testbench: {testbench.name}", self.styles['ChapterTitle']))

        # Generation info (small, unobtrusive)
        gen_info = f"Generated by analogpy | {datetime.now().strftime('%Y-%m-%d %H:%M')}"
        if self.config.author:
            gen_info = f"{self.config.author} | {gen_info}"
        elements.append(Paragraph(gen_info, self.styles['SmallInfo']))
        elements.append(Spacer(1, 0.15 * inch))

        # Block diagram FIRST - this is what users want to see
        elements.append(Paragraph("Block Diagram", self.styles['SectionTitle']))

        try:
            drawing = create_testbench_block_diagram(testbench, style=SymbolStyle.DETAILED)
            img = self._render_drawing_to_image(drawing, width_inches=5.5)  # Scaled down to fit on page
            if img:
                elements.append(img)
            else:
                elements.append(Paragraph("[Block diagram rendering failed]", self.styles['Normal']))
        except Exception as e:
            elements.append(Paragraph(f"[Diagram error: {e}]", self.styles['Normal']))

        elements.append(Spacer(1, 0.3 * inch))

        # Testbench info below the diagram
        info_lines = []

        # Supplies (now identified from subcircuit_instances)
        # Collect voltage sources
        voltage_sources_info = []
        for inst in testbench.subcircuit_instances:
            if hasattr(inst.subcircuit, '_is_primitive') and inst.subcircuit._is_primitive and \
               inst.subcircuit._primitive_spectre_type == "vsource":
                p_net_name = inst.connections.get("p").name if inst.connections.get("p") else "unknown"
                dc_value = inst.params.get("dc")
                if dc_value is not None:
                    voltage_sources_info.append(f"{p_net_name}={dc_value}V")
        
        if voltage_sources_info:
            supplies_str = ", ".join(voltage_sources_info)
            info_lines.append(f"<b>Supplies:</b> {supplies_str}")

        # Instances count
        if testbench.subcircuit_instances:
            info_lines.append(f"<b>Instances:</b> {len(testbench.subcircuit_instances)} subcircuits")

        # Analyses
        if hasattr(testbench, 'analyses') and testbench.analyses:
            analyses_str = ", ".join([a.name if hasattr(a, 'name') else str(type(a).__name__)
                                      for a in testbench.analyses])
            info_lines.append(f"<b>Analyses:</b> {analyses_str}")

        for line in info_lines:
            elements.append(Paragraph(line, self.styles['CellInfo']))

        elements.append(PageBreak())
        return elements

    def _create_circuit_page(self, circuit: 'Circuit', page_num: int) -> List:
        """Create a page for a circuit/cell."""
        elements = []

        # Record bookmark
        self._page_bookmarks.append((f"Cell: {circuit.name}", page_num))

        elements.append(Paragraph(f"Cell: {circuit.name}", self.styles['ChapterTitle']))

        # Circuit info
        if circuit.ports:
            elements.append(Paragraph(
                f"<b>Ports:</b> {', '.join(circuit.ports)}",
                self.styles['CellInfo']
            ))

        # Schematic drawing FIRST
        elements.append(Spacer(1, 0.15 * inch))
        elements.append(Paragraph("Schematic", self.styles['SectionTitle']))

        try:
            drawing = create_cell_schematic(circuit, title=circuit.name, compact=True)
            img = self._render_drawing_to_image(drawing, width_inches=5.5)  # Scaled down to fit
            if img:
                elements.append(img)
            else:
                elements.append(Paragraph("[Schematic rendering failed]", self.styles['Normal']))
        except Exception as e:
            elements.append(Paragraph(f"[Diagram error: {e}]", self.styles['Normal']))

        elements.append(Spacer(1, 0.2 * inch))

        # Devices info (unified API uses subcircuit_instances)
        if circuit.subcircuit_instances:
            elements.append(Paragraph("<b>Devices:</b>", self.styles['CellInfo']))
            for inst in circuit.subcircuit_instances:
                model = inst.subcircuit.name if hasattr(inst, 'subcircuit') else 'unknown'
                elements.append(Paragraph(f"  {inst.name} ({model})", self.styles['CellInfo']))

        # Parameters from comments (added by @auto_comment_params)
        if hasattr(circuit, 'comments') and circuit.comments:
            elements.append(Spacer(1, 0.1 * inch))
            elements.append(Paragraph("<b>Parameters:</b>", self.styles['CellInfo']))
            for comment in circuit.comments[:8]:  # Limit to 8
                elements.append(Paragraph(f"  {comment}", self.styles['CellInfo']))

        elements.append(PageBreak())
        return elements

    def _add_bookmarks_to_pdf(self, input_path: str, output_path: str):
        """Add bookmarks/outline to the PDF for sidebar navigation."""
        reader = PdfReader(input_path)
        writer = PdfWriter()

        # Copy all pages
        for page in reader.pages:
            writer.add_page(page)

        # Add bookmarks (outline)
        for title, page_num in self._page_bookmarks:
            # page_num is 1-indexed, pypdf uses 0-indexed
            writer.add_outline_item(title, page_num - 1)

        # Set PDF metadata
        writer.add_metadata({
            '/Title': self.config.title,
            '/Author': self.config.author or 'analogpy',
            '/Creator': 'analogpy SchematicBook',
            '/Producer': 'reportlab + pypdf',
        })

        # Write output
        with open(output_path, 'wb') as f:
            writer.write(f)

    def render(self, output_path: str, include_legend: bool = True):
        """
        Render the schematic book to a PDF file.

        Page 1 shows the testbench block diagram immediately.
        Table of contents is in PDF sidebar (bookmarks).

        Args:
            output_path: Path to save the PDF
            include_legend: Whether to include a color legend as the last page
        """
        # Create temp file for initial PDF (without bookmarks)
        fd, temp_path = tempfile.mkstemp(suffix='.pdf')
        os.close(fd)

        try:
            doc = SimpleDocTemplate(
                temp_path,
                pagesize=self.page_size,
                leftMargin=self.config.margin * inch,
                rightMargin=self.config.margin * inch,
                topMargin=self.config.margin * inch,
                bottomMargin=self.config.margin * inch,
            )

            elements = []
            self._page_bookmarks = []
            self._temp_files = []  # Reset temp files
            page_num = 1

            # Testbench pages FIRST (page 1 = testbench)
            for tb in self.testbenches:
                elements.extend(self._create_testbench_page(tb, page_num))
                page_num += 1

            # Circuit pages
            for circuit in self.circuits:
                elements.extend(self._create_circuit_page(circuit, page_num))
                page_num += 1

            # Optional Legend page at the end
            if include_legend:
                elements.extend(self._create_legend_page(page_num))
                page_num += 1

            # Build initial PDF
            doc.build(elements)

            # Clean up image temp files after build
            self._cleanup_temp_files()

            # Add bookmarks and write final PDF
            self._add_bookmarks_to_pdf(temp_path, output_path)

            print(f"Schematic PDF generated: {output_path}")
            print(f"  - {len(self.testbenches)} testbench(es)")
            print(f"  - {len(self.circuits)} cell(s)")
            if include_legend:
                print("  - Legend page included at the end")
            print("  - Table of contents in PDF sidebar (bookmarks)")

        finally:
            # Clean up temp PDF file
            try:
                os.unlink(temp_path)
            except:
                pass
            # Clean up any remaining image temp files
            self._cleanup_temp_files()

    def render_to_bytes(self) -> bytes:
        """Render the schematic book to bytes (for in-memory use)."""
        # Create temp files
        fd1, temp_path1 = tempfile.mkstemp(suffix='.pdf')
        fd2, temp_path2 = tempfile.mkstemp(suffix='.pdf')
        os.close(fd1)
        os.close(fd2)

        try:
            doc = SimpleDocTemplate(
                temp_path1,
                pagesize=self.page_size,
                leftMargin=self.config.margin * inch,
                rightMargin=self.config.margin * inch,
                topMargin=self.config.margin * inch,
                bottomMargin=self.config.margin * inch,
            )

            elements = []
            self._page_bookmarks = []
            page_num = 1

            for tb in self.testbenches:
                elements.extend(self._create_testbench_page(tb, page_num))
                page_num += 1

            for circuit in self.circuits:
                elements.extend(self._create_circuit_page(circuit, page_num))
                page_num += 1

            doc.build(elements)
            self._add_bookmarks_to_pdf(temp_path1, temp_path2)

            with open(temp_path2, 'rb') as f:
                return f.read()
        finally:
            try:
                os.unlink(temp_path1)
                os.unlink(temp_path2)
            except:
                pass
