import io
import matplotlib
matplotlib.use('Agg') # Headless mode
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np

# Set seaborn theme for professional looks
sns.set_theme(style="whitegrid", palette="muted")
plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.sans-serif'] = ['Arial', 'Helvetica', 'DejaVu Sans']

def _create_bytes_img(fig):
    buf = io.BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight', dpi=150)
    buf.seek(0)
    plt.close(fig)
    return buf

def generate_revenue_net_income_chart(data: dict) -> io.BytesIO:
    years = data.get("years", [])
    if not years:
        return None
    
    rev = data.get("income_statement", {}).get("revenue", [])
    ni = data.get("income_statement", {}).get("net_income", [])
    
    # Pad to match years
    rev = (rev + [0]*len(years))[:len(years)]
    ni = (ni + [0]*len(years))[:len(years)]
    
    # Fill None with 0 for plotting
    rev = [r if r is not None else 0 for r in rev]
    ni = [n if n is not None else 0 for n in ni]
    
    fig, ax1 = plt.subplots(figsize=(8, 4.5))
    
    x = np.arange(len(years))
    width = 0.35
    
    ax1.bar(x - width/2, rev, width, label='Revenue', color='#1F4E78')
    ax1.bar(x + width/2, ni, width, label='Net Income', color='#549E39')
    
    ax1.set_ylabel(f"Amount ({data.get('currency', 'INR')} {data.get('units_label', 'Crores')})")
    ax1.set_title("Revenue & Net Income Trend")
    ax1.set_xticks(x)
    ax1.set_xticklabels(years)
    ax1.legend()
    
    return _create_bytes_img(fig)

def generate_margins_chart(data: dict) -> io.BytesIO:
    years = data.get("years", [])
    if not years:
        return None
        
    ratios = data.get("ratios", {})
    gm = ratios.get("gross_margin", [])
    om = ratios.get("operating_margin", [])
    nm = ratios.get("net_margin", [])
    
    gm = [(g*100) if g is not None else 0 for g in (gm + [0]*len(years))[:len(years)]]
    om = [(o*100) if o is not None else 0 for o in (om + [0]*len(years))[:len(years)]]
    nm = [(n*100) if n is not None else 0 for n in (nm + [0]*len(years))[:len(years)]]
    
    fig, ax = plt.subplots(figsize=(8, 4.5))
    x = np.arange(len(years))
    
    ax.plot(x, gm, marker='o', linewidth=2, label='Gross Margin %', color='#1F4E78')
    ax.plot(x, om, marker='s', linewidth=2, label='Operating Margin %', color='#455F51')
    ax.plot(x, nm, marker='^', linewidth=2, label='Net Margin %', color='#549E39')
    
    ax.set_ylabel("Margin (%)")
    ax.set_title("Profitability Margins Trend")
    ax.set_xticks(x)
    ax.set_xticklabels(years)
    
    # Add value labels
    for i in range(len(years)):
        if nm[i] != 0:
            ax.annotate(f"{nm[i]:.1f}%", (x[i], nm[i]), textcoords="offset points", xytext=(0,10), ha='center')
            
    ax.legend()
    
    return _create_bytes_img(fig)

def generate_cash_flow_chart(data: dict) -> io.BytesIO:
    years = data.get("years", [])
    if not years:
        return None
        
    cf = data.get("cash_flow", {})
    ocf = cf.get("cash_from_operations", [])
    capex = cf.get("capex", [])
    fcf = cf.get("free_cash_flow", [])
    
    ocf = [o if o is not None else 0 for o in (ocf + [0]*len(years))[:len(years)]]
    # Capex is usually negative, let's plot absolute value for magnitude comparison
    capex = [abs(c) if c is not None else 0 for c in (capex + [0]*len(years))[:len(years)]]
    fcf = [f if f is not None else 0 for f in (fcf + [0]*len(years))[:len(years)]]
    
    fig, ax = plt.subplots(figsize=(8, 4.5))
    x = np.arange(len(years))
    width = 0.25
    
    ax.bar(x - width, ocf, width, label='Operating Cash Flow', color='#1F4E78')
    ax.bar(x, capex, width, label='CAPEX (Abs)', color='#C00000')
    ax.bar(x + width, fcf, width, label='Free Cash Flow', color='#549E39')
    
    ax.set_ylabel(f"Amount ({data.get('currency', 'INR')} {data.get('units_label', 'Crores')})")
    ax.set_title("Cash Flow & CAPEX Overview")
    ax.set_xticks(x)
    ax.set_xticklabels(years)
    ax.legend()
    
    return _create_bytes_img(fig)

def generate_all_charts(data: dict) -> dict:
    """Returns a dict mapping chart names to BytesIO pngs."""
    return {
        "revenue_net_income": generate_revenue_net_income_chart(data),
        "margins": generate_margins_chart(data),
        "cash_flow": generate_cash_flow_chart(data)
    }