import json
import html
import re
from datetime import datetime
from typing import Dict, Any, List

class HTMLReportGenerator:
    def __init__(self, agent_state: Dict[str, Any]):
        self.state = agent_state
        self.title = "SciLink Research Report"
        self.generated_by = "Generated by SciLink Planning Agent"
        
    def _get_css(self) -> str:
        return """
        <style>
            :root { 
                --primary: #2563eb; 
                --bg: #f8fafc; 
                --card-bg: #ffffff; 
                --code-bg: #1e293b;
                --success-bg: #f0fdf4; --success-border: #bbf7d0; --success-text: #166534;
            }
            body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif; line-height: 1.6; color: #334155; background: var(--bg); margin: 0; padding: 40px; }
            .container { max-width: 1000px; margin: 0 auto; }
            
            /* Header */
            header { background: var(--card-bg); padding: 30px; border-radius: 12px; box-shadow: 0 1px 3px rgba(0,0,0,0.1); margin-bottom: 40px; border-bottom: 4px solid var(--primary); }
            h1 { margin: 0; color: #1e293b; font-size: 2em; }
            .meta { color: #64748b; font-size: 0.9em; margin-top: 10px; font-family: monospace; }
            .objective { background: #eff6ff; border-left: 4px solid var(--primary); padding: 15px; margin-top: 20px; color: #1e40af; }
            
            /* Timeline Cards */
            .card { background: var(--card-bg); border-radius: 12px; padding: 0; margin-bottom: 40px; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1); overflow: hidden; }
            .card-header { padding: 15px 25px; background: #f1f5f9; border-bottom: 1px solid #e2e8f0; display: flex; justify-content: space-between; align-items: center; }
            .card-header h2 { margin: 0; font-size: 1.2em; color: #475569; }
            .badge { padding: 4px 10px; border-radius: 20px; font-size: 0.75em; font-weight: bold; text-transform: uppercase; letter-spacing: 0.05em; }
            .badge.plan { background: #dbeafe; color: #1e40af; }
            .badge.tea { background: #d1fae5; color: #065f46; }
            .card-body { padding: 25px; }

            /* Content Sections */
            .exp-block { margin-bottom: 30px; border-left: 3px solid #e2e8f0; padding-left: 20px; }
            .justification { background: #fffbeb; color: #92400e; padding: 10px; border-radius: 6px; font-style: italic; font-size: 0.95em; }
            
            pre { background: var(--code-bg); color: #e2e8f0; padding: 15px; border-radius: 8px; overflow-x: auto; font-family: 'Menlo', 'Monaco', monospace; font-size: 0.85em; }
            .result-box { margin-top: 20px; padding: 20px; background: var(--success-bg); border: 1px solid var(--success-border); border-radius: 8px; color: var(--success-text); }
            
            /* TEA Specific */
            .tea-list li { margin-bottom: 8px; }
            .cost { color: #b91c1c; }
            .benefit { color: #047857; }

            /* TABLE STYLING */
            .styled-table {
                border-collapse: collapse;
                margin: 25px 0;
                font-size: 0.9em;
                font-family: sans-serif;
                min-width: 400px;
                width: 100%;
                box-shadow: 0 0 20px rgba(0, 0, 0, 0.05);
                border-radius: 8px;
                overflow: hidden;
            }
            .styled-table thead tr {
                background-color: var(--primary);
                color: #ffffff;
                text-align: left;
            }
            .styled-table th, .styled-table td {
                padding: 12px 15px;
            }
            .styled-table tbody tr {
                border-bottom: 1px solid #dddddd;
            }
            .styled-table tbody tr:nth-of-type(even) {
                background-color: #f3f3f3;
            }
            .styled-table tbody tr:last-of-type {
                border-bottom: 2px solid var(--primary);
            }
            .styled-table tbody tr:hover {
                background-color: #e2e8f0;
            }
        </style>
        """

    def _markdown_to_html(self, text: str) -> str:
        """
        Parses text and converts Markdown tables to styled HTML tables.
        Also handles basic line breaks.
        """
        if not text: return ""
        
        lines = text.split('\n')
        html_output = []
        in_table = False
        table_buffer = []

        def render_table(buffer):
            if not buffer: return ""
            # Filter out divider rows (e.g., |---|---|)
            content_rows = [row for row in buffer if not re.match(r'^\s*\|?[\s\-\:|]+\|?\s*$', row)]
            
            if not content_rows: return ""

            html_table = '<table class="styled-table">'
            
            # Header
            header_cells = [c.strip() for c in content_rows[0].strip('|').split('|')]
            html_table += '<thead><tr>' + ''.join(f'<th>{html.escape(h)}</th>' for h in header_cells) + '</tr></thead>'
            
            # Body
            html_table += '<tbody>'
            for row in content_rows[1:]:
                cells = [c.strip() for c in row.strip('|').split('|')]
                # Handle cases where row might have fewer cells than header
                while len(cells) < len(header_cells): cells.append("")
                html_table += '<tr>' + ''.join(f'<td>{html.escape(c)}</td>' for c in cells) + '</tr>'
            
            html_table += '</tbody></table>'
            return html_table

        for line in lines:
            stripped = line.strip()
            # Check if line looks like a table row (starts/ends with pipe or contains pipes)
            if stripped.startswith('|') or (stripped.count('|') > 1 and '-' in lines[lines.index(line)+1] if lines.index(line)+1 < len(lines) else False):
                in_table = True
                table_buffer.append(stripped)
            else:
                if in_table:
                    # Flush table buffer
                    html_output.append(render_table(table_buffer))
                    table_buffer = []
                    in_table = False
                
                # Regular text line
                html_output.append(html.escape(line) + "<br>")

        # Flush any remaining table
        if in_table:
            html_output.append(render_table(table_buffer))

        return "\n".join(html_output)

    def _render_experiment(self, exp: Dict[str, Any], index: int) -> str:

        code_html = ""
        if "implementation_code" in exp:
            code_html = f"""
            <div style="margin-top:15px;">
                <strong>💻 Generated Code:</strong>
                <pre><code>{html.escape(exp['implementation_code'])}</code></pre>
            </div>
            """
        
        steps_html = ""
        raw_steps = exp.get('experimental_steps', [])
        
        if raw_steps:
             if isinstance(raw_steps, list):
                 merged_steps = []
                 table_buffer = []
                 
                 # Regex matches lines that have at least one internal pipe (e.g., "A | B")
                 # ^\s* : Start of string, optional whitespace
                 # \|?      : Optional starting pipe
                 # .*\|.* : Content containing at least one pipe in the middle
                 # \|?      : Optional ending pipe
                 # \s*$     : Optional whitespace, end of string
                 table_row_pattern = re.compile(r"^\s*\|?.*\|.*\|?\s*$")

                 for step in raw_steps:
                     s_str = str(step).strip()
                     
                     # Check if it looks like a table row
                     # We strictly require at least one pipe '|' to avoid merging normal text sentences
                     is_table_row = ('|' in s_str) and table_row_pattern.match(s_str)
                     
                     if is_table_row:
                         table_buffer.append(s_str)
                     else:
                         # If we hit a non-table line, flush the buffer first
                         if table_buffer:
                             merged_steps.append("\n".join(table_buffer))
                             table_buffer = []
                         merged_steps.append(s_str)
                 
                 # Flush any remaining table at the very end
                 if table_buffer:
                     merged_steps.append("\n".join(table_buffer))
                 
                 # Render the processed steps
                 steps_content = ""
                 for s in merged_steps:
                     rendered = self._markdown_to_html(s)
                     steps_content += f"<li>{rendered}</li>"
                 
                 steps_html = f"<ul>{steps_content}</ul>"
             else:
                 # Handle case where steps are just a single string block
                 steps_html = self._markdown_to_html(str(raw_steps))

        return f"""
        <div class="exp-block">
            <h4 style="color: #2563eb; margin-bottom: 5px;">Experiment {index}: {html.escape(exp.get('experiment_name', 'Unnamed'))}</h4>
            
            <div style="margin-bottom: 10px;">
                <strong>🎯 Hypothesis:</strong> 
                <div>{self._markdown_to_html(exp.get('hypothesis', ''))}</div>
            </div>
            
            <div class="justification">
                <strong>💡 Justification:</strong> 
                <div>{self._markdown_to_html(exp.get('justification', ''))}</div>
            </div>
            
            <div style="margin-top: 15px;">
                <strong>🧪 Steps:</strong>
                {steps_html}
            </div>
            {code_html}
        </div>
        """

    def _render_tea(self, plan: Dict[str, Any]) -> str:
        assess = plan.get('technoeconomic_assessment', {})
        def list_to_html(key, css_class):
            return "".join(f"<li class='{css_class}'>{html.escape(str(x))}</li>" for x in assess.get(key, []))
        
        # Summary might contain tables
        summary_html = self._markdown_to_html(assess.get('summary', ''))

        return f"""
        <div class="exp-block" style="border-left-color: #10b981;">
            <div class="justification" style="background: #ecfdf5; color: #065f46;">
                <strong>💰 Executive Summary:</strong><br>
                {summary_html}
            </div>
            <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 20px; margin-top: 20px;">
                <div><strong>💸 Key Cost Drivers:</strong><ul class="tea-list">{list_to_html('key_cost_drivers', 'cost')}</ul></div>
                <div><strong>📈 Potential Benefits:</strong><ul class="tea-list">{list_to_html('potential_benefits_or_revenue', 'benefit')}</ul></div>
            </div>
            <div style="margin-top: 15px;"><strong>⚠️ Economic Risks:</strong><ul>{list_to_html('economic_risks', '')}</ul></div>
        </div>
        """

    def generate(self, output_path: str):
        date_str = self.state.get('start_time', datetime.now().isoformat())
        
        full_history = self.state.get('plan_history', [])

        # Filter out superseded plans - they stay in JSON for transparency
        active_plans = [p for p in full_history if p.get('status') != 'superseded']

        finalized_plans = {}
        for plan in active_plans:
            iter_idx = plan.get('iteration', 0)
            finalized_plans[iter_idx] = plan

        sorted_plans = [finalized_plans[k] for k in sorted(finalized_plans.keys())]
        results = self.state.get('experimental_results', [])

        html_content = f"""
        <!DOCTYPE html><html><head><meta charset="UTF-8"><title>{self.title}</title>{self._get_css()}</head>
        <body><div class="container">
            <header>
                <h1>🔭 {self.title}</h1>
                <div class="meta">Session: {self.state.get('session_id', 'N/A')} &bull; Started: {date_str}</div>
                <div class="objective"><strong>Objective:</strong> {html.escape(self.state.get('objective', ''))}</div>
            </header>
        """

        for i, plan in enumerate(sorted_plans):
            step_num = i + 1 
            is_tea = "technoeconomic_assessment" in plan or plan.get("type") == "technoeconomic_analysis"
            badge_class = "tea" if is_tea else "plan"
            badge_text = "TECHNO-ECONOMIC ANALYSIS" if is_tea else "EXPERIMENTAL STRATEGY"
            
            if is_tea:
                content_html = self._render_tea(plan)
            else:
                content_html = "".join(self._render_experiment(exp, x+1) for x, exp in enumerate(plan.get('proposed_experiments', [])))

            html_content += f"""
            <div class="card">
                <div class="card-header">
                    <h2>Step {step_num}</h2>
                    <span class="badge {badge_class}">{badge_text}</span>
                </div>
                <div class="card-body">{content_html}"""

            if i < len(results):
                res_data = results[i].get('data_summary', '')
                try:
                    if isinstance(res_data, str) and (res_data.strip().startswith('{') or res_data.strip().startswith('[')):
                        res_data = json.dumps(json.loads(res_data), indent=2)
                except: pass
                
                # Apply markdown parser to results too, in case the result log contains a table
                formatted_result = self._markdown_to_html(str(res_data))
                
                html_content += f"""
                    <div class="result-box">
                        <div style="font-weight:bold; margin-bottom:10px;">📊 Results Received:</div>
                        <div style="font-family: monospace; font-size: 0.9em;">{formatted_result}</div>
                    </div>
                """

            html_content += "</div></div>"

        html_content += f"<footer><p>{self.generated_by}</p></footer></div></body></html>"

        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(html_content)
            print(f"  - 📄 HTML Report updated: {output_path}")
        except Exception as e:
            print(f"  - ❌ Error writing HTML report: {e}")