
import sys
import re
import json
import logging
from pathlib import Path
from mcp.server import Server
from .llm_analyzer import request_sampling

# Setup file logging
_log_file = Path(__file__).parent.parent.parent / "rtlviz_sampling.log"
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(_log_file, mode='a', encoding='utf-8'),
    ]
)

# STRICT PROMPT tailored for Direct DOT Generation (Total Clustering & Context Style)
DIRECT_GEN_PROMPT_TEMPLATE = """You are a Digital Design Automation Expert. Your goal is to generate a HIGH-FIDELITY, PROFESSIONAL Graphviz (DOT) visualization for the provided RTL code.

### **VISUAL DESIGN PHILOSOPHY: TOTAL CLUSTERING & WHITE-BOX CONTEXT**
1.  **Ortho-Pervasive**: Use `splines=ortho`. Lines must be clean, right-angled, and NEVER cut through nodes.
2.  **Total Clustering**: EVERY significant logic block or instance MUST be inside a `subgraph cluster_...`. grouping by semantic function or clock domain.
3.  **Pastel Palette**: Use soft, professional pastel colors for backgrounds.
    -   **Clock Domains**: Light Red/Blue backgrounds.
    -   **Logic Blocks**: Light Purple/Green/Yellow fills.
4.  **White-Box Detail**:
    -   Inside logic nodes, DO NOT just show the name.
    -   **List Internal Variables**: Show parameters, key registers, and state bits inside the node record.
    -   Example: `label="{ module_name | { state | count } }"` using record shapes.
5.  **Rich Shapes**:
    -   **MUX**: `shape=trapezium, orientation=90` (Right facing).
    -   **Flops**: `shape=box3d`.
    -   **Memory**: `shape=component`.

### **PHASE 1: STRUCTURAL ANALYSIS**
1.  Identify **Clock Domains**: Group logic driven by different clocks into separate clusters.
2.  Identify **Data Flow**: Clearly separate "Write Side" vs "Read Side" or "Control" vs "Datapath".
3.  Identify **Critical Signals**: Highlight Full/Empty flags, Pointers, and Handshakes.

### **PHASE 2: DOT GENERATION RULES**
Start with:
```dot
digraph "RTL_Systems" {
    graph [rankdir=LR, splines=ortho, nodesep=1.2, ranksep=1.5, fontname="Arial"];
    node [fontname="Arial", shape=record, style="filled,rounded", penwidth=1.5];
    edge [fontname="Arial", fontsize=10, penwidth=1.5];
    
    // Legend or Global Styles can go here
```

#### **Styling Reference**
-   **Clusters**: `style="filled,rounded"`, `color="#444444"`, `fillcolor="#[PASTEL_HEX]"`
    -   Read Domain: `#FFEBEE` (Light Red)
    -   Write Domain: `#E3F2FD` (Light Blue)
    -   Logic/Storage: `#F3E5F5` (Light Purple)
-   **Nodes**:
    -   **Registers**: `shape=record, style="filled,rounded", fillcolor="#FFF9C4"` (Yellow-ish) -> SHOW CONTENTS i.e., `q1`, `q2`.
    -   **Logic/Calculations**: `shape=record, fillcolor="#E8F5E9"` (Green) -> Show operations e.g., `Next Ptr Calculation`.
    -   **MUX**: `shape=trapezium, orientation=90, fillcolor="#E0E0E0", width=0.5`.
    -   **Ports/IO**: `shape=parallelogram, fillcolor="#B2DFDB"`.

#### **Output Requirement**
-   **NO PYTHON LOGIC**: You must generate the final, valid DOT code directly.
-   **EXPLICIT CONNECTIONS**: Label every wire with its signal name and width (e.g., `[label="wdata[7:0]"]`).
-   **NO OVERLAPS**: Use invisible edges (`style=invis`) if necessary to force ordering, but rely mostly on `ranksep/nodesep`.

### **INPUT RTL:**
```verilog
{full_source_code}
```

### **OUTPUT:**
Provide ONLY the Graphviz DOT code inside a code block.
"""

async def generate_full_dot(server: Server, full_source_code: str, line_count: int) -> str:
    """
    Generate DOT code directly from full source using LLM.
    
    Args:
        server: MCP Server instance
        full_source_code: Concatenated Verilog code
        line_count: Total lines of code
        
    Returns:
        String containing DOT code
    """
    
    logging.info(f"Direct Gen Triggered. Lines: {line_count}")
    
    # Threshold Check
    if line_count > 7000:
        logging.warning("Threshold exceeded.")
        return "" 
        
    try:
        prompt = DIRECT_GEN_PROMPT_TEMPLATE.replace("{full_source_code}", full_source_code)
        logging.info("Requesting Sampling...")
    except Exception as e:
        logging.error(f"Prompt construction failed: {e}")
        return ""
    
    try:
        # Request Sampling (High Token Limit for Graphic Generation)
        result_text = await request_sampling(server, prompt, max_tokens=4000)
    except Exception as e:
        logging.error(f"Sampling Exception: {e}")
        return ""
    
    if not result_text:
        logging.error("Empty response from LLM")
        return ""
        
    logging.info(f"LLM Response received. Length: {len(result_text)}")
    logging.debug(f"Snippet: {result_text[:200]}")
    
    # Extract DOT
        
    # Extract DOT
    dot_match = re.search(r'digraph\s+.*{', result_text, re.DOTALL)
    if not dot_match:
        # Try looser match
        dot_match = re.search(r'digraph\s*{', result_text)
        
    if dot_match:
        # Find the matching closing brace? Hard with regex.
        # Better: extract generic code block
        code_block = re.search(r'```dot\n(.*?)```', result_text, re.DOTALL)
        if code_block:
            return code_block.group(1).strip()
        
        code_block_generic = re.search(r'```\n(.*?)```', result_text, re.DOTALL)
        if code_block_generic and 'digraph' in code_block_generic.group(1):
             return code_block_generic.group(1).strip()
             
        # Fallback: Just return the whole text if it looks like DOT
        if 'digraph' in result_text and '}' in result_text:
             # Strip markdown
             clean = result_text.replace('```dot', '').replace('```', '')
             return clean.strip()
             
    return ""
