import os
import textwrap
import nbformat
from nbformat import v4 as nbf
import logging
import json
from datetime import datetime, timezone
import pandas as pd

logger = logging.getLogger(__name__)

class Notebook:
    def __init__(self, data):
        self.nb = nbf.new_notebook()
        
        self.nb.cells = []
        self.data = data
        self.output_path = "tmp/notebook.ipynb"
        
        self.imports_added = False


    def build_notebook(self):
        success = 0
        metrics = self.data
        logger.info(f"Start building notebook with {len(metrics)} metrics")
        for i, metric in enumerate(metrics):
            logger.info(f"Processing metric {i+1} of {len(metrics)}")
            metric_result = json.loads(metric.get("body"))
            table_data = metric_result.get("table_data")
            query_config = metric_result.get("sql_config")
            
            chart_data = metric_result.get("chart_data", {})
            chart_config = metric_result.get("chart_config", {})
            chart_type = metric_result.get("chart_config", {}).get("type", {})
            
            if not self.nb.cells:
                self._add_title()
                
            self._add_metric_header(config=chart_config, index=i)
            
            self._add_query_details(query_config=query_config)
            if not self.imports_added:
                self._import_libs()
                self.imports_added = True

            self._add_table(raw_table=table_data)

            if chart_data and chart_config:
                self._chart_code(chart_type=chart_type, chart_data=chart_data, chart_config=chart_config)

    def export_notebook(self):
        self.build_notebook()
        logger.info(f"Exporting notebook to {self.output_path}")
        os.makedirs(os.path.dirname(self.output_path), exist_ok=True)
        with open(self.output_path, 'w') as f:
            nbformat.write(self.nb, f)
        logger.info(f"Notebook exported successfully to {self.output_path}")
        return self.output_path
    
    def _import_libs(self):
        self.nb.cells.append(nbf.new_markdown_cell("### 📦 Imports"))
        self.nb.cells.append(nbf.new_code_cell("import pandas as pd\nimport plotly.express as px"))
    
    
    def _add_title(self):
        self.nb.cells.append(nbf.new_markdown_cell("# Metrics Dashboard"))
    
    
    def _add_metric_header(self, config, index):
        self.nb.cells.append(nbf.new_markdown_cell(f"## 📈 Metric {index+1} {config.get('title', '')}\n**Timestamp:** {datetime.now(timezone.utc).isoformat()}"))
    
        
    def _chart_code(self, chart_type, chart_data, chart_config):
        df_code = f"df = pd.DataFrame({chart_data.get(chart_type).get('data')})"
        chart_dispatch = {
            "bar": f"""
            fig = px.bar(df, x='x', y='y', title='{chart_config['title']}',
                   labels={{'y': '{chart_config['y_axis_label']}', 'x': '{chart_config['x_axis_label']}'}}, color='x', template='plotly_white')
            """,
            
            "line": f"""
            fig = px.line(df, x='x', y='y', title='{chart_config['title']}',
                   labels={{'y': '{chart_config['y_axis_label']}', 'x': '{chart_config['x_axis_label']}'}}, template='plotly_white')
            """,
            
            "pie": f"""
            fig = px.pie(df, values='y', names='x', title='{chart_config['title']}',
                   labels={{'y': '{chart_config['y_axis_label']}'}}, template='plotly_white', hole=0.3)
            fig.update_traces(textposition='inside', textfont_size=12)
            """,
            
            "hist": f"""
            df['bin_center'] = (df['bin_start'] + df['bin_end']) / 2
            fig = px.bar(df, x='bin_center', y='frequency', title='{chart_config['title']}',
                labels={{'frequency': '{chart_config['y_axis_label']}', 'bin_center': '{chart_config['x_axis_label']}'}},
                template='plotly_white')
            """
        }
        chart_code = f"{df_code}\n{textwrap.dedent(chart_dispatch.get(chart_type))}\nfig.update_layout(showlegend=False)\nfig.show()"
        
        self.nb.cells.append(nbf.new_markdown_cell(f"### 📊 Chart"))
        self.nb.cells.append(nbf.new_code_cell(chart_code))
        
    def _add_query_details(self, query_config):
        sql_clean = query_config['sql_query'].replace("\\n", "\n").strip()
        self.nb.cells.append(nbf.new_markdown_cell(f"**User Query:** {query_config['text_query']}\n"))
        self.nb.cells.append(nbf.new_markdown_cell(f"### 🔍 Query Details\n**SQL:**\n```sql\n{sql_clean}\n```\n**Explanation:** {query_config['explanation']}"))
        
    def _add_table(self, raw_table):
        table_data = pd.DataFrame(raw_table).fillna("NAN").to_dict(orient="records")
        raw_code = f'raw_table = pd.DataFrame({table_data})\nraw_table.head()'
        self.nb.cells.append(nbf.new_markdown_cell("### Table"))
        self.nb.cells.append(nbf.new_code_cell(raw_code))