app/main.py

##########################
#        NEW MAIN APP        #
##########################
import os
import asyncio
import asyncpg
import logging
from typing import List, Dict, Any, Optional, Literal
import re
import time
import ssl
import json
from datetime import datetime
import pathlib as Path
import uuid
import base64
import io
import tempfile

# ──────────────────────────────────────────────────────────────────────────────
# CUID generator (25-char, starts with "c")
# ──────────────────────────────────────────────────────────────────────────────
import hashlib, socket, random
import matplotlib

# Fix matplotlib permission issues by setting cache directory to a writable location
os.environ['MPLCONFIGDIR'] = '/tmp/matplotlib_cache'
matplotlib.use('Agg')  # Use non-interactive backend for server environments

import matplotlib.pyplot as plt
import numpy as np

_CUID_COUNTER = 0
def _base36(n: int) -> str:
    chars = "0123456789abcdefghijklmnopqrstuvwxyz"
    if n == 0: return "0"
    s = []
    nn = abs(n)
    while nn:
        nn, r = divmod(nn, 36)
        s.append(chars[r])
    s = "".join(reversed(s))
    return "-" + s if n < 0 else s

def _pad36(n: int, width: int) -> str:
    s = _base36(n)
    return s.rjust(width, "0")[-width:]

def _fingerprint_block() -> str:
    src = f"{socket.gethostname()}-{os.getpid()}"
    h = int(hashlib.md5(src.encode()).hexdigest(), 16)
    return _pad36(h, 4)

def _rand_block() -> str:
    # 36^4 = 1,679,616; use 20 random bits then base36-pad to 4 chars
    return _pad36(random.SystemRandom().getrandbits(20), 4)

def new_cuid() -> str:
    global _CUID_COUNTER
    ts = int(time.time() * 1000)
    _CUID_COUNTER = (_CUID_COUNTER + 1) % (36**4)  # 0..36^4-1
    return "c" + _pad36(ts, 8) + _pad36(_CUID_COUNTER, 4) + _fingerprint_block() + _rand_block() + _rand_block()

from dotenv import load_dotenv
load_dotenv()  # Load environment variables from .env into os.environ

# Supabase client
from supabase import create_client, Client

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi import Query
from fastapi.responses import FileResponse
from pydantic import BaseModel
from starlette.responses import StreamingResponse
from typing import AsyncIterator
from collections import defaultdict
from fastapi import Depends, HTTPException, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials

import asyncpg
from jose import jwt, JWTError

# LangChain / LLMs
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.schema.runnable import RunnableSequence
from langchain.schema import AIMessage
from langchain.llms.base import LLM

from huggingface_hub import InferenceClient

# ──────────────────────────────────────────────────────────────────────────────
# Document Generation Imports
# ──────────────────────────────────────────────────────────────────────────────
import sys
import os

# Add the Document Generation and Currency Mapping paths to sys.path
sys.path.append(os.path.join(os.path.dirname(__file__), 'DocumentGeneration'))
sys.path.append(os.path.join(os.path.dirname(__file__), 'CurrencyMapping'))

# Import models
from models import ExportRequest, ExportSection, ExportResponse, SectionJob, BatchGenerateRequest, GenerateRequest, CreatePlanRequest

# Import constants
from constants import SECTION_ORDER, SECTION_MAP, SECTION_COLUMNS

# Import utilities
from utils import sort_sections_by_order

# Now import the modules
import sys
import os

# Add the Document Generation and Currency Mapping paths to sys.path
sys.path.append(os.path.join(os.path.dirname(__file__), 'DocumentGeneration'))
sys.path.append(os.path.join(os.path.dirname(__file__), 'CurrencyMapping'))

from DocumentGeneration.pdf_generator import create_pdf_document
from DocumentGeneration.word_generator import create_word_document
from DocumentGeneration.chart_generator import parse_embedded_chart_data, create_chart_from_data
from DocumentGeneration.table_of_content import generate_table_of_contents_after_content, add_table_of_contents_page
from DocumentGeneration.document_helpers import (
    extract_business_name_from_content,
    filter_empty_sections,
    extract_subheadings,
    remove_duplicate_headings
)
from DocumentGeneration.grouped import parse_content_into_groups, extract_key_phrases, get_content_groups_info
from DocumentGeneration.markdown_renderer import (
    clean_text_for_pdf,
    render_markdown_to_pdf,
    render_markdown_to_pdf_grouped,
    render_markdown_to_docx,
    render_markdown_to_docx_grouped,
    clean_text_for_docx
)
from CurrencyMapping.currency_mapping import convert_currency_symbols_to_iso

# ──────────────────────────────────────────────────────────────────────────────
# Logging
# ──────────────────────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO)

# ──────────────────────────────────────────────────────────────────────────────
# SSE broker
# ──────────────────────────────────────────────────────────────────────────────
# --- In-memory SSE broker keyed by businessPlanId ---
SUBSCRIBERS: dict[str, set[asyncio.Queue[str]]] = defaultdict(set)

def _sse_format(event: str, payload: dict) -> str:
    return f"event: {event}\n" f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"

async def _publish(plan_id: str, event: str, payload: dict) -> None:
    msg = _sse_format(event, payload)
    for q in list(SUBSCRIBERS.get(plan_id, set())):
        try:
            q.put_nowait(msg)
        except Exception:
            SUBSCRIBERS[plan_id].discard(q)

# ──────────────────────────────────────────────────────────────────────────────
# Environment checks
# ──────────────────────────────────────────────────────────────────────────────
def check_environment_variables():
    hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
    if not hf_token:
        logging.warning("⚠️ HUGGINGFACEHUB_API_TOKEN not set. Hugging Face models will not work properly.")
    else:
        logging.info("✅ HUGGINGFACEHUB_API_TOKEN is set")
    
    openai_api_key = os.getenv("OPENAI_API_KEY")
    if not openai_api_key:
        logging.warning("⚠️ OPENAI_API_KEY not set. OpenAI GPT models will not work properly.")
    else:
        logging.info("✅ OPENAI_API_KEY is set")

check_environment_variables()

logging.info("=" * 80)
logging.info("MODEL SIZE LIMITATIONS:")
logging.info("The free tier of Hugging Face Inference API limits models to 10GB.")
logging.info("Large models like Qwen-2.5-7B (15GB) and Llama-2-7B (13GB) exceed this limit.")
logging.info("We've configured smaller alternative models as replacements.")
logging.info("For full-sized models, upgrade to Hugging Face Pro subscription.")
logging.info("=" * 80)

# ──────────────────────────────────────────────────────────────────────────────
# FastAPI app & CORS
# ──────────────────────────────────────────────────────────────────────────────
app = FastAPI()

@app.middleware("http")
async def log_requests(request: Request, call_next):
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    logging.info(f"Request: {request.method} {request.url.path} - Status: {response.status_code} - Time: {process_time:.2f}s")
    return response

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"]
)

# ──────────────────────────────────────────────────────────────────────────────
# Allow all origins (adjust for production usage)
# app.add_middleware(
#     CORSMiddleware,
#     allow_origins=["*"],
#     allow_credentials=True,
#     allow_methods=["*"],
#     allow_headers=["*"],
# )
# Fixed list of business questions (order matters)
# ──────────────────────────────────────────────────────────────────────────────
QUESTIONS = [
    "What is your business name?",
    "What product or service do you offer?",
    "Who is your target customer?",
    "What problem does your business solve?",
    "Who are your competitors?",
    "What is your unique value proposition?",
    "What is your pricing strategy?",
    "What are your short-term goals?",
    "What are your long-term goals?",
    "How will you acquire customers?",
]

# ──────────────────────────────────────────────────────────────────────────────
# SectionKey type
# ──────────────────────────────────────────────────────────────────────────────
SectionKey = Literal[
    "prompt_ExecutiveSummary","prompt_CompanyProfile","prompt_MarketAnalysis",
    "prompt_ProductOrService","prompt_BusinessModel","prompt_MarketingGrowth",
    "prompt_OperationsPlan","prompt_ManagementTeam","prompt_FinancialPlanFunding"
]

class SectionJob(BaseModel):
    sectionKey: SectionKey
    # The server will load prompt/model defaults from GeneratePrompt.
    # Client MAY override but is not required to send these.
    prompt: Optional[str] = None
    promptVariables: Dict[str, str] = {}
    model: Optional[str] = None
    provider: Optional[str] = None
    temperature: Optional[float] = None
    maxTokens: Optional[int] = None

# Database helper constants

# Database helper functions



def extract_summary_from_markdown(markdown_text: str, max_length: int = 200) -> str:
    """Extract a summary from markdown text, removing headers and formatting"""
    if not markdown_text:
        return ""
    
    # Try to find the first header as a title
    lines = markdown_text.split('\n')
    for line in lines:
        if line.strip().startswith('#'):
            header_text = re.sub(r'^#+\s*', '', line.strip())
            if header_text:
                return header_text[:max_length]
    
    # Fallback to first non-empty line
    for line in lines:
        clean_line = re.sub(r'[*_`~#]', '', line.strip())
        if clean_line:
            return clean_line[:max_length]
    
    return "Generated Business Plan"

# ──────────────────────────────────────────────────────────────────────────────
# Helpers for status flags
# ──────────────────────────────────────────────────────────────────────────────
async def _ensure_sections_row(pool, business_plan_id: str) -> None:
    async with pool.acquire() as conn:
        await conn.execute(
            'INSERT INTO "BusinessPlanSection" ("id","businessPlanId","updatedAt") '
            'VALUES ($1,$2, now()) ON CONFLICT ("businessPlanId") DO NOTHING',
            new_cuid(), business_plan_id
        )

async def charge_full_plan_once(
    pool,
    *,
    business_plan_id: str,
    user_id: Optional[str],
    plan_type: Optional[str],
) -> tuple[bool, int]:
    """
    Charges ONE credit for a full-plan generation, once per business_plan_id.
    Returns (charged_now, current_credits_after).
    - Only runs for plan_type == 'full' and authenticated users.
    - Idempotent via CreditCharge(businessPlanId) unique constraint.
    - Updates User.currentCredits -= 1 and User.totalUsedCredits += 1 atomically.
    """
    if plan_type != "full":
        return (False, -1)
    if not user_id:
        # Full plans require a signed-in user (credits live on users)
        raise HTTPException(status_code=401, detail="Sign in required to generate a full plan.")


    async with pool.acquire() as conn:
        # If we've already charged this plan, just return current credits tally.
        already = await conn.fetchval(
            'SELECT 1 FROM "CreditCharge" WHERE "businessPlanId"=$1',
            business_plan_id,
        )
        if already:
            # return current balance for UI if you want to display it
            bal = await conn.fetchval(
                'SELECT "currentCredits" FROM "User" WHERE "id"=$1',
                user_id,
            )
            return (False, int(bal) if bal is not None else -1)

        # Charge once in a transaction; guards concurrency
        async with conn.transaction():
            # Lock the user row
            row = await conn.fetchrow(
                'SELECT "currentCredits","totalUsedCredits" FROM "User" WHERE "id"=$1 FOR UPDATE',
                user_id,
            )
            if not row:
                raise HTTPException(status_code=404, detail="User not found.")
            current = int(row["currentCredits"] or 0)
            if current <= 0:
                raise HTTPException(status_code=402, detail="Insufficient credits.")

            # Deduct and increment usage
            await conn.execute(
                'UPDATE "User" SET "currentCredits" = "currentCredits" - 1, "totalUsedCredits" = "totalUsedCredits" + 1 WHERE "id"=$1',
                user_id,
            )
            # Mark this plan as charged
            await conn.execute(
                'INSERT INTO "CreditCharge" ("id","businessPlanId","userId") VALUES ($1,$2,$3)',
                new_cuid(), business_plan_id, user_id,
            )

            # Return new balance
            new_bal = await conn.fetchval(
                'SELECT "currentCredits" FROM "User" WHERE "id"=$1',
                user_id,
            )
            return (True, int(new_bal) if new_bal is not None else -1)

async def _set_section_status(pool, *, business_plan_id: str, section_key: str, generating: bool, complete: bool) -> None:
    if section_key not in SECTION_MAP:
        return
    _, col_complete, col_generating = SECTION_MAP[section_key]
    async with pool.acquire() as conn:
        await conn.execute(
            f'''
            UPDATE "BusinessPlanSection"
            SET "{col_generating}"=$2,
                "{col_complete}"=$3,
                "updatedAt"=now()
            WHERE "businessPlanId"=$1
            ''',
            business_plan_id, generating, complete
        )

# ──────────────────────────────────────────────────────────────────────────────
# Request model
# ──────────────────────────────────────────────────────────────────────────────
class GenerateRequest(BaseModel):
    answers: List[str]
    model: str
    prompt: str
    promptVariables: Dict[str, str]
    provider: str
    temperature: float
    maxTokens: int
    # NEW — optional, for persistence & ownership
    businessPlanId: Optional[str] = None
    planType: Optional[str] = "free"
    userId: Optional[str] = None           # will be overridden by verified JWT
    anonymousId: Optional[str] = None
    countryCode: Optional[str] = None
    sectionKey: Optional[str] = None       # e.g. "prompt_ExecutiveSummary"
    returnContent: Optional[bool] = False  # NEW — default to id-only responses 
    generateId: Optional[str] = None

class TestGenerationRequest(BaseModel):
    answers: List[str]
    model: str
    prompt: str
    promptVariables: Dict[str, Any]
    provider: str
    temperature: float
    maxTokens: int
    currentDate: Optional[str] = None
    sectionKey: str
    planType: str = "test"
    countryCode: Optional[str] = None
    businessIdea: Optional[str] = None

class TestGenerationResponse(BaseModel):
    success: bool
    plan: str
    summary: Optional[str] = None

# ──────────────────────────────────────────────────────────────────────────────
# LLM plumbing
# ──────────────────────────────────────────────────────────────────────────────
def clean_markdown_text(text: str) -> str:
    """Clean markdown formatting for PDF generation - PRESERVE HEADINGS"""
    if not text:
        return ""
    
    # Remove markdown formatting BUT PRESERVE HEADING STRUCTURE
    cleaned = text
    
    # PRESERVE HEADINGS - don't strip the # markers, keep them for rendering
    # We'll handle styling at render time instead of stripping them
    
    # Keep bold, italic formatting - don't strip ** or * symbols
    # cleaned = re.sub(r'\*\*(.*?)\*\*', r'\1', cleaned)  # Remove bold - COMMENTED OUT
    # cleaned = re.sub(r'\*(.*?)\*', r'\1', cleaned)  # Remove italic - COMMENTED OUT
    cleaned = re.sub(r'`(.*?)`', r'\1', cleaned)  # Remove code
    cleaned = re.sub(r'~~(.*?)~~', r'\1', cleaned)  # Remove strikethrough
    
    # Remove links (keep text)
    cleaned = re.sub(r'\[([^\]]+)\]\([^)]+\)', r'\1', cleaned)  # Remove links, keep text
    cleaned = re.sub(r'!\[([^\]]*)\]\([^)]+\)', '', cleaned)  # Remove images
    
    # Remove list markers (more aggressive)
    cleaned = re.sub(r'^\s*[-*+]\s*', '', cleaned, flags=re.MULTILINE)  # Remove list markers
    cleaned = re.sub(r'^\s*\d+\.\s*', '', cleaned, flags=re.MULTILINE)  # Remove numbered list markers
    cleaned = re.sub(r'^\s*[-*+]\s*', '', cleaned, flags=re.MULTILINE)  # Remove any remaining list markers
    
    # Remove blockquotes
    cleaned = re.sub(r'^\s*>\s*', '', cleaned, flags=re.MULTILINE)  # Remove blockquotes
    
    # Remove horizontal rules
    cleaned = re.sub(r'^\s*[-*_]{3,}\s*$', '', cleaned, flags=re.MULTILINE)  # Remove horizontal rules
    
    # Remove code blocks
    cleaned = re.sub(r'```[\s\S]*?```', '', cleaned)  # Remove code blocks
    cleaned = re.sub(r'`.*?`', '', cleaned)  # Remove any remaining inline code
    
    # Keep emphasis markers - don't strip ** or * symbols
    # cleaned = re.sub(r'_{1,2}(.*?)_{1,2}', r'\1', cleaned)  # Remove underscores - COMMENTED OUT
    # cleaned = re.sub(r'\*{1,2}(.*?)\*{1,2}', r'\1', cleaned)  # Remove asterisks - COMMENTED OUT
    
    # REMOVE RAW CHART DATA MARKERS (CRITICAL FIX)
    cleaned = re.sub(r'<!--CHARTDATASTART-->.*?<!--CHARTDATAEND-->', '', cleaned, flags=re.DOTALL)
    cleaned = re.sub(r'<!--CHARTDATASTART-->.*$', '', cleaned, flags=re.DOTALL)
    cleaned = re.sub(r'<!--CHARTDATAEND-->', '', cleaned)
    
    # Remove any remaining HTML-like tags
    cleaned = re.sub(r'<[^>]+>', '', cleaned)
    
    # Clean up extra whitespace and formatting
    cleaned = re.sub(r'\n\s*\n\s*\n+', '\n\n', cleaned)  # Remove excessive line breaks
    cleaned = re.sub(r'^\s+', '', cleaned, flags=re.MULTILINE)  # Remove leading whitespace
    cleaned = re.sub(r'\s+$', '', cleaned, flags=re.MULTILINE)  # Remove trailing whitespace
    cleaned = re.sub(r' +', ' ', cleaned)  # Replace multiple spaces with single space
    
    # Final cleanup
    cleaned = cleaned.strip()
    
    return cleaned


def generate_model_output(model: str, provider: str, api_key: str, prompt: str, max_tokens: int = 4000) -> str:
    try:
        logging.info(f"Initializing InferenceClient with provider: {provider}")
        client = InferenceClient(provider=provider, api_key=api_key)
        logging.info(f"Sending request to model: {model} with prompt length: {len(prompt)} and max_tokens: {max_tokens}")
        completion = client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=max_tokens,
        )
        logging.info(f"Successfully received response from model: {model} with content length: {len(completion.choices[0].message.content)}")
        return completion.choices[0].message.content
    except Exception as e:
        error_message = f"Error generating output with model {model}: {str(e)}"
        logging.error(error_message)
        raise Exception(error_message) from e

class HFInferenceLLM(LLM):
    model: str = None
    provider: str = "hf-inference"
    api_key: str = ""
    max_tokens: int = 4000

    def __init__(self, model: str, provider: str = "hf-inference", api_key: str = "", max_tokens: int = 4000):
        super().__init__()
        self.model = model
        self.provider = provider
        self.api_key = api_key
        self.max_tokens = max_tokens
        
        if not api_key:
            logging.error(f"No API key provided for model {model}")
            raise ValueError(f"API key is required for Hugging Face Inference API access to {model}")
        try:
            _ = InferenceClient(provider=provider, api_key=api_key)
            logging.info(f"Successfully initialized client for model: {model}")
        except Exception as e:
            logging.error(f"Failed to initialize client for model {model}: {str(e)}")
            raise ValueError(f"Could not initialize Hugging Face client for {model}: {str(e)}")

    @property
    def _llm_type(self) -> str:
        return "hf_inference"

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        return generate_model_output(
            model=self.model,
            provider=self.provider,
            api_key=self.api_key,
            prompt=prompt,
            max_tokens=self.max_tokens
        )

    def get_num_tokens(self, prompt: str) -> int:
        return len(prompt.split())

# def get_llm(model_name: str, provider: str = "hf-inference"):
#     hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
#     if not hf_token:
#         logging.error("HUGGINGFACEHUB_API_TOKEN environment variable not set")
#         raise ValueError("HUGGINGFACEHUB_API_TOKEN environment variable is required for Hugging Face models")
        
#     if model_name == "BPGenerateAI":
#         openai_api_key = os.getenv("OPENAI_API_KEY")
#         if not openai_api_key:
#             logging.error("OPENAI_API_KEY environment variable not set")
#             raise ValueError("OPENAI_API_KEY environment variable is required for GPT models")
#         return ChatOpenAI(
#             model_name="gpt-4.1-mini",
#             temperature=0.7,
#             max_tokens=6000,
#             openai_api_key=openai_api_key
#         )
#     elif model_name == "BPSuggestionsAI":
#         openai_api_key = os.getenv("OPENAI_API_KEY")
#         if not openai_api_key:
#             logging.error("OPENAI_API_KEY environment variable not set")
#             raise ValueError("OPENAI_API_KEY environment variable is required for GPT models")
#         return ChatOpenAI(
#             model_name="gpt-4.1-nano",
#             temperature=0.7,
#             max_tokens=4000,
#             openai_api_key=openai_api_key
#         )
#     elif model_name.lower() == "gpt-4.1-nano":
#         openai_api_key = os.getenv("OPENAI_API_KEY")
#         if not openai_api_key:
#             raise ValueError("OPENAI_API_KEY is required for GPT models")
#         return ChatOpenAI(
#             model_name="gpt-4.1-nano",
#             temperature=0.7,
#             max_tokens=4000,
#             openai_api_key=openai_api_key
#         )
#     else:
#         return HFInferenceLLM(
#             model=model_name,
#             provider=provider,
#             api_key=hf_token,
#             max_tokens=4000
#         )
def get_llm(model_name: str, provider: str):
    """
    Strict provider routing.
    - provider='openai' -> use OpenAI (requires OPENAI_API_KEY)
    - provider in {'hf-inference','together'} -> use Hugging Face InferenceClient (requires HUGGINGFACEHUB_API_TOKEN)
    - No automatic fallbacks.
    """
    if not provider:
        raise ValueError("No provider specified. Set provider to 'openai', 'hf-inference', or 'together'.")
    provider = provider.lower()

    supported = {"openai", "hf-inference", "together"}
    if provider not in supported:
        raise ValueError(f"Unsupported provider '{provider}'. Must be one of {sorted(supported)}.")

    # ─────────────────────────────────────────────────────────────
    # OPENAI
    # ─────────────────────────────────────────────────────────────
    if provider == "openai":
        openai_api_key = os.getenv("OPENAI_API_KEY")
        if not openai_api_key:
            raise ValueError("OPENAI_API_KEY is required when provider='openai'.")

        # Map your aliases to concrete OpenAI models; allow raw names too.
        if model_name == "BPGenerateAI":
            model = "gpt-4.1-mini"
            max_tokens = 6000
        elif model_name in ("BPSuggestionsAI", "gpt-4.1-nano"):
            model = "gpt-4.1-nano"
            max_tokens = 4000
        else:
            model = model_name
            max_tokens = 4000

        return ChatOpenAI(
            model_name=model,
            temperature=0.7,
            max_tokens=max_tokens,
            openai_api_key=openai_api_key,
        )

    # ─────────────────────────────────────────────────────────────
    # HF Inference (and compatible providers like 'together')
    # ─────────────────────────────────────────────────────────────
    hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
    if not hf_token:
        raise ValueError(f"HUGGINGFACEHUB_API_TOKEN is required when provider='{provider}'.")

    # Map your aliases to HF models; allow raw names too.
    if model_name == "BPGenerateAI":
        model = "mistralai/Mistral-7B-Instruct-v0.3"
        max_tokens = 4000
    elif model_name == "BPSuggestionsAI":
        model = "Qwen/Qwen2.5-3B-Instruct"  # smaller/faster for suggestions
        max_tokens = 4000
    else:
        model = model_name
        max_tokens = 4000

    return HFInferenceLLM(
        model=model,
        provider=provider,           # 'hf-inference' or 'together'
        api_key=hf_token,
        max_tokens=max_tokens
    )

# ──────────────────────────────────────────────────────────────────────────────
# Supabase JWT verification
# ──────────────────────────────────────────────────────────────────────────────
SUPABASE_JWT_SECRET = os.getenv("SUPABASE_JWT_SECRET")
SUPABASE_JWT_AUD = os.getenv("SUPABASE_JWT_AUD")  # optional; e.g. "authenticated"

def _get_bearer_token(request: Request) -> Optional[str]:
    auth = request.headers.get("Authorization", "")
    if auth.startswith("Bearer "):
        return auth.split(" ", 1)[1].strip()
    return None

def get_user_id_from_request(request: Request) -> Optional[str]:
    token = _get_bearer_token(request)
    if not token or not SUPABASE_JWT_SECRET:
        return None
    # basic shape check: header.payload.signature
    if token.count(".") != 2:
        logging.info("Authorization header present but not a JWT; skipping verification.")
        return None
    try:
        if SUPABASE_JWT_AUD:
            payload = jwt.decode(token, SUPABASE_JWT_SECRET, algorithms=["HS256"], audience=SUPABASE_JWT_AUD)
        else:
            payload = jwt.decode(token, SUPABASE_JWT_SECRET, algorithms=["HS256"])
        return payload.get("sub")
    except JWTError as e:
        logging.warning(f"JWT verification failed: {e}")
        return None

def verify_supabase_jwt(token: str) -> Optional[str]:
    """
    Verify a Supabase JWT token and return the user ID (sub claim).
    
    Args:
        token: The JWT token string
        
    Returns:
        The user ID from the 'sub' claim, or None if verification fails
    """
    if not token or not SUPABASE_JWT_SECRET:
        return None
    
    # basic shape check: header.payload.signature
    if token.count(".") != 2:
        logging.info("Token is not a valid JWT format; skipping verification.")
        return None
    
    try:
        if SUPABASE_JWT_AUD:
            payload = jwt.decode(token, SUPABASE_JWT_SECRET, algorithms=["HS256"], audience=SUPABASE_JWT_AUD)
        else:
            payload = jwt.decode(token, SUPABASE_JWT_SECRET, algorithms=["HS256"])
        return payload.get("sub")
    except JWTError as e:
        logging.warning(f"JWT verification failed: {e}")
        return None

# --- Ownership resolver: prefer verified user; else anonymous ---
def resolve_owner(request: Request, provided_anonymous_id: Optional[str]) -> tuple[Optional[str], Optional[str]]:
    """
    Returns (user_id, anonymous_id):
    - If Authorization JWT is present and valid -> (user_id, None)
    - Else -> (None, provided_anonymous_id or None)
    """
    user_id = get_user_id_from_request(request)
    if user_id:
        return user_id, None
    return None, (provided_anonymous_id or None)

# --- Optional: guard plan ownership to prevent hijacking ---
async def assert_plan_ownership(pool, *, business_plan_id: str, user_id: Optional[str], anonymous_id: Optional[str]) -> None:
    """
    - If plan has a userId: only that user can modify it.
    - If plan has an anonymousId: only the same anonymousId can modify it.
    - If plan has neither: allow (legacy).
    Raises HTTPException(403) if not allowed.
    """
    if not business_plan_id:
        return
    async with pool.acquire() as conn:
        row = await conn.fetchrow('SELECT "userId","anonymousId" FROM "BusinessPlan" WHERE id=$1', business_plan_id)
        if not row:
            return
        plan_user = row["userId"]
        plan_anon = row["anonymousId"]

        # If plan belongs to a user, you must be that user
        if plan_user:
            if not user_id or user_id != plan_user:
                raise HTTPException(status_code=403, detail="Not allowed to modify this plan (user mismatch).")

        # If plan belongs to an anonymous identity, you must present the same anonymousId
        if plan_anon and not plan_user:
            if anonymous_id != plan_anon:
                raise HTTPException(status_code=403, detail="Not allowed to modify this plan (anonymous mismatch).")
# ──────────────────────────────────────────────────────────────────────────────
# Database (asyncpg pooled)
# ──────────────────────────────────────────────────────────────────────────────
DATABASE_URL = os.getenv("DATABASE_URL")  # pooled endpoint recommended (6543)
DB_CA_CERT_PEM  = os.getenv("DB_CA_CERT_PEM")

@app.on_event("startup")
async def startup_event():
    if not DATABASE_URL:
        logging.error("DATABASE_URL not set; DB writes disabled.")
        return

    if not DB_CA_CERT_PEM:
        raise RuntimeError("DB_CA_CERT_PEM is not set. Provide the PEM certificate content in the environment.")

    # Convert '\n' sequences to real newlines if needed
    pem = DB_CA_CERT_PEM.replace("\\n", "\n")

    ssl_ctx = ssl.create_default_context()
    # Load the PEM content directly from env
    ssl_ctx.load_verify_locations(cadata=pem)
    ssl_ctx.check_hostname = True
    ssl_ctx.verify_mode = ssl.CERT_REQUIRED

    app.state.db_pool = await asyncpg.create_pool(
        dsn=DATABASE_URL,
        ssl=ssl_ctx,
        min_size=0,
        max_size=8,
        command_timeout=15.0,
        statement_cache_size=0,
    )
    logging.info("✅ DB pool initialized")
    
    # Create CreditCharge table for tracking one-time charges per plan
    async with app.state.db_pool.acquire() as conn:
        await conn.execute(
            '''
            CREATE TABLE IF NOT EXISTS "CreditCharge" (
              "id"             text PRIMARY KEY,
              "businessPlanId" text UNIQUE NOT NULL,
              "userId"         text NOT NULL,
              "createdAt"      timestamptz NOT NULL DEFAULT now()
            );
            '''
        )

@app.on_event("shutdown")
async def shutdown_event():
    pool = getattr(app.state, "db_pool", None)
    if pool:
        await pool.close()
        logging.info("✅ DB pool closed")
# ──────────────────────────────────────────────────────────────────────────────
# Section mapping & helpers
# ──────────────────────────────────────────────────────────────────────────────
SECTION_MAP = {
    "prompt_ExecutiveSummary": ("executiveSummary", "isExecutiveSummaryComplete", "isExecutiveSummaryGenerating"),
    "prompt_CompanyProfile": ("companyProfile", "isCompanyProfileComplete", "isCompanyProfileGenerating"),
    "prompt_MarketAnalysis": ("marketAnalysis", "isMarketAnalysisComplete", "isMarketAnalysisGenerating"),
    "prompt_ProductOrService": ("productOrService", "isProductOrServiceComplete", "isProductOrServiceGenerating"),
    "prompt_BusinessModel": ("businessModel", "isBusinessModelComplete", "isBusinessModelGenerating"),
    "prompt_MarketingGrowth": ("marketingGrowth", "isMarketingGrowthComplete", "isMarketingGrowthGenerating"),
    "prompt_OperationsPlan": ("operationsPlan", "isOperationsPlanComplete", "isOperationsPlanGenerating"),
    "prompt_ManagementTeam": ("managementTeam", "isManagementTeamComplete", "isManagementTeamGenerating"),
    "prompt_FinancialPlanFunding": ("financialPlanFunding", "isFinancialPlanFundingComplete", "isFinancialPlanFundingGenerating"),
}

# Database column names for sections (using the consolidated definition above)

def extract_summary_from_markdown(md: str) -> str:
    m = re.search(r"^\s*#{1,2}\s+(.+)", md, flags=re.MULTILINE)
    if m:
        return m.group(1).strip()[:250]
    return (md.strip().split("\n", 1)[0] if md.strip() else "Generated Business Plan")[:250]

async def _get_generate_prompt(pool, key: str):
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            'SELECT "key","title","promptTemplate","inputVariables","defaultModel","defaultProvider",'
            '"defaultTemp","defaultMaxTokens","defaultCountry","isActive" '
            'FROM "GeneratePrompt" WHERE "key"=$1 AND "isActive"=TRUE',
            key,
        )
        return row

def _plan_key_for_section(base_key: str, plan_type: str) -> str:
    # free plan uses *_Free prompts except FinancialPlanFunding
    if plan_type == "free" and base_key != "prompt_FinancialPlanFunding":
        return f"{base_key}_Free"
    return base_key

# --- helpers: country/currency mapping ---

def _country_name_from_code(code: Optional[str], default_name: Optional[str]) -> str:
    if not code:
        return default_name or "South Africa"
    m = {
        "ZA": "South Africa","US": "United States","GB": "United Kingdom","DE": "Germany","FR": "France",
        "NG": "Nigeria","KE": "Kenya","IN": "India","AU": "Australia","CA": "Canada","BR": "Brazil",
        "NL": "Netherlands","ES": "Spain","IT": "Italy","IE": "Ireland","SG": "Singapore","AE": "United Arab Emirates"
    }
    return m.get(code.upper(), default_name or code)

def _currency_from_country_code(code: Optional[str], fallback_country_name: Optional[str]) -> str:
    # prefer code, else derive from country name, else USD
    if code:
        m = {
            "ZA": "ZAR","US": "USD","GB": "GBP","DE": "EUR","FR": "EUR","NL": "EUR","ES": "EUR","IT": "EUR","IE":"EUR",
            "NG": "NGN","KE": "KES","IN": "INR","AU": "AUD","CA": "CAD","BR": "BRL","SG":"SGD","AE":"AED"
        }
        if code.upper() in m:
            return m[code.upper()]
    if fallback_country_name:
        name = fallback_country_name.lower()
        if "south africa" in name: return "ZAR"
        if any(x in name for x in ["united states","usa","america"]): return "USD"
        if "united kingdom" in name or "uk" in name: return "GBP"
        if any(x in name for x in ["germany","france","netherlands","spain","italy","ireland"]): return "EUR"
        if "nigeria" in name: return "NGN"
        if "kenya" in name: return "KES"
        if "india" in name: return "INR"
        if "australia" in name: return "AUD"
        if "canada" in name: return "CAD"
        if "brazil" in name: return "BRL"
        if "singapore" in name: return "SGD"
        if "emirates" in name or "uae" in name: return "AED"
    return "USD"

def _compose_prompt_vars(
    *,
    gp_row,
    answers: List[str],
    business_idea: Optional[str],
    country_code: Optional[str],
    section_title: str,
    feedback: Optional[str],
) -> Dict[str, str]:
    """
    Build exactly the variables that the GeneratePrompt row declares in inputVariables.
    Examples we've seen: ["business_context","country","currency","q_and_a"].
    """

    # Parse declared input variables from DB (stringified JSON)
    try:
        declared = json.loads(gp_row["inputVariables"] or "[]")
    except Exception:
        declared = []

    # Precompute building blocks we might map into the declared names
    now_iso = datetime.utcnow().isoformat()
    country_name = _country_name_from_code(country_code, gp_row.get("defaultCountry"))
    currency = _currency_from_country_code(country_code, country_name)

    # Build a Q&A string (or JSON) from your QUESTIONS & answers array
    # (safe against length mismatch; unanswered = "N/A")
    qa_lines = []
    for i, q in enumerate(QUESTIONS):
        a = answers[i] if i < len(answers) and answers[i] else "N/A"
        qa_lines.append(f"{q} -> {a}")
    qa_text = "\n".join(qa_lines)

    # Business context: prefer explicit business_idea, else stitch from first answers
    business_context = (business_idea or "").strip()
    if not business_context:
        # use first two answers as a minimal context
        name = answers[0] if len(answers) > 0 else ""
        offering = answers[1] if len(answers) > 1 else ""
        business_context = (name + " — " + offering).strip(" —")

    # We'll fill only what the prompt expects. If nothing declared, we fall back
    # to a generic superset (backward compat).
    declared = declared if isinstance(declared, list) else []

    # Canonical mapping for common keys you're using in GeneratePrompt
    mapping = {
        "business_context": business_context,
        "q_and_a": qa_text,
        "country": country_name,
        "currency": currency,

        # Optional alternates you might have in some prompts
        "section_title": section_title,
        "current_date": now_iso,
        "feedback": feedback or "",
        "answers_json": json.dumps(answers or []),
        "businessIdea": business_idea or "",  # if some prompts still use old name
        "countryCode": country_code or "",    # if some prompts still use old name
    }

    if declared:
        # Return exactly the keys the prompt asked for; raise obvious gaps to logs
        out = {}
        for key in declared:
            out[key] = mapping.get(key, "")
        return out

    # Fallback for prompts that didn't declare inputVariables
    return {
        "business_context": business_context,
        "q_and_a": qa_text,
        "country": country_name,
        "currency": currency,
        "section_title": section_title,
        "current_date": now_iso,
        "feedback": feedback or "",
        "answers_json": json.dumps(answers or []),
    }

async def _ensure_business_plan(pool, *, bp_id: Optional[str], summary: str, full_plan: str,
                                model: str, answers: List[str], plan_type: Optional[str],
                                user_id: Optional[str], anon_id: Optional[str],
                                country_code: Optional[str], generate_id: Optional[str] = None) -> str:
    async with pool.acquire() as conn:
        if bp_id:
            exists = await conn.fetchval('SELECT 1 FROM "BusinessPlan" WHERE id = $1', bp_id)
            if exists:
                await conn.execute(
                    'UPDATE "BusinessPlan" SET '
                    '"summary"=$1, "model"=$2, "answers"=$3, '
                    '"planType"=COALESCE($4,\'free\'), "countryCode"=$5, '
                    '"generateId"=COALESCE($6,"generateId"), '
                    '"updatedAt"=now(), "userId"=COALESCE($7,"userId") '
                    'WHERE id=$8',
                    summary, model, answers, plan_type, country_code, generate_id, user_id, bp_id
                )
                return bp_id

        new_id = new_cuid()
        row = await conn.fetchrow(
            'INSERT INTO "BusinessPlan" '
            '("id","summary","fullPlan","model","answers","planType","userId","anonymousId","countryCode","generateId","updatedAt") '
            'VALUES ($1,$2,$3,$4,$5,COALESCE($6,\'free\'),$7,$8,$9,$10, now()) '
            'RETURNING id',
            new_id, summary, "", model, answers, plan_type, user_id, anon_id, country_code, generate_id
        )
        return row["id"]

# async def _upsert_section(pool, *, business_plan_id: str, section_key: str, content: str) -> None:
#     if section_key not in SECTION_MAP:
#         logging.warning(f"Unknown section_key '{section_key}', skipping section save.")
#         return

#     col_content, col_complete, col_generating = SECTION_MAP[section_key]
#     set_clause = f'"{col_content}" = $2, "{col_complete}" = TRUE, "{col_generating}" = FALSE'

#     async with pool.acquire() as conn:
#         # ensure row exists with a generated id
#         await conn.execute(
#             'INSERT INTO "BusinessPlanSection" ("id","businessPlanId") '
#             'VALUES ($1,$2) '
#             'ON CONFLICT ("businessPlanId") DO NOTHING',
#             str(uuid.uuid4()), business_plan_id
#         )
#         # update content + flags
#         await conn.execute(
#             f'UPDATE "BusinessPlanSection" SET {set_clause} WHERE "businessPlanId" = $1',
#             business_plan_id, content
#         )
async def _upsert_section(pool, *, business_plan_id: str, section_key: str, content: str) -> None:
    if section_key not in SECTION_MAP:
        # logging.warning(f"Unknown section_key '{section_key}', skipping section save.")
        return

    col_content, col_complete, col_generating = SECTION_MAP[section_key]
    # bump updatedAt on every write
    set_clause = (
        f'"{col_content}" = $2, '
        f'"{col_complete}" = TRUE, '
        f'"{col_generating}" = FALSE, '
        f'"updatedAt" = now()'
    )

    async with pool.acquire() as conn:
        # ensure row exists, and give updatedAt a value to satisfy NOT NULL
        await conn.execute(
            'INSERT INTO "BusinessPlanSection" ("id","businessPlanId","updatedAt") '
            'VALUES ($1,$2, now()) '
            'ON CONFLICT ("businessPlanId") DO NOTHING',
            new_cuid(), business_plan_id
        )
        # update content + flags (+ updatedAt)
        await conn.execute(
            f'UPDATE "BusinessPlanSection" SET {set_clause} WHERE "businessPlanId" = $1',
            business_plan_id, content
        )

async def _recompute_full_plan(pool, *, business_plan_id: str) -> str:
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            'SELECT * FROM "BusinessPlanSection" WHERE "businessPlanId" = $1',
            business_plan_id
        )
        if not row:
            return ""
        parts = []
        for col in SECTION_ORDER:
            txt = row.get(col)
            if txt:
                parts.append(txt.strip())
        full_plan = "\n\n".join(parts)
        return full_plan  # Do not persist fullPlan; caller may use it for summary only

def _all_sections_present(row: asyncpg.Record) -> bool:
    return all(bool(row.get(col)) for col in SECTION_ORDER)

# ──────────────────────────────────────────────────────────────────────────────
# Endpoints
# ──────────────────────────────────────────────────────────────────────────────
@app.get("/suggestions", response_model=Dict[str, Any])
async def get_suggestions(
    business_idea: str,
    model: str,
    question: str,
    prompt: str,
    provider: str,
    temperature: float,
    maxTokens: int,
    exclude: List[str] = Query([])
) -> Dict[str, Any]:
    try:
        formatted_prompt = prompt.format(
            business_idea=business_idea,
            question=question,
            exclude=", ".join(exclude) if exclude else ""
        )
        llm = get_llm(model, provider)
        if hasattr(llm, 'temperature'):
            llm.temperature = temperature
        if hasattr(llm, 'max_tokens'):
            llm.max_tokens = maxTokens

        suggestion_prompt = PromptTemplate(input_variables=[], template=formatted_prompt)
        suggestion_chain: RunnableSequence = suggestion_prompt | llm
        raw_result = await asyncio.to_thread(suggestion_chain.invoke, {})

        if isinstance(raw_result, AIMessage):
            raw_text = raw_result.content
        elif isinstance(raw_result, (list, tuple)) and raw_result and isinstance(raw_result[0], AIMessage):
            raw_text = "\n".join(msg.content for msg in raw_result)
        else:
            raw_text = str(raw_result)

        suggestion_array = [
            re.sub(r'^\s*[\-\d\.\)\s]+', '', line).strip()
            for line in raw_text.split("\n")
            if line.strip()
        ]
        if not suggestion_array:
            suggestion_array = ["No suggestions available"]
        return {"suggestions": suggestion_array}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error generating suggestions: {str(e)}")

@app.post("/generate/test", response_model=TestGenerationResponse)
async def generate_test_content(data: TestGenerationRequest):
    """
    Generate test content without saving to database or charging credits.
    Mirrors the /generate logic for LLM invocation and formatting only.
    """
    try:
        # Prepare prompt template and LLM using existing utilities
        llm_selected = get_llm(data.model, data.provider)
        if hasattr(llm_selected, 'temperature'):
            llm_selected.temperature = data.temperature
        if hasattr(llm_selected, 'max_tokens'):
            llm_selected.max_tokens = data.maxTokens

        plan_prompt = PromptTemplate(
            input_variables=list(data.promptVariables.keys()),
            template=data.prompt
        )
        plan_chain: RunnableSequence = plan_prompt | llm_selected

        raw_plan = await asyncio.to_thread(plan_chain.invoke, data.promptVariables)
        if isinstance(raw_plan, AIMessage):
            plan_text = raw_plan.content
        elif isinstance(raw_plan, (list, tuple)) and raw_plan and isinstance(raw_plan[0], AIMessage):
            plan_text = "\n".join(msg.content for msg in raw_plan)
        else:
            plan_text = str(raw_plan)

        summary = extract_summary_from_markdown(plan_text)

        return TestGenerationResponse(
            success=True,
            plan=plan_text,
            summary=summary,
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Test generation failed: {str(e)}")
    
class CreatePlanRequest(BaseModel):
    answers: List[str] = []
    planType: Optional[str] = "free"
    countryCode: Optional[str] = None
    anonymousId: Optional[str] = None
    businessIdea: Optional[str] = None
    generateId: Optional[str] = None

@app.post("/plan")
async def create_plan_shell(request: Request, data: CreatePlanRequest):
    pool = getattr(app.state, "db_pool", None)
    if not pool:
        raise HTTPException(status_code=500, detail="DB not initialized")

    # Resolve owner (JWT beats anonymousId)
    user_id, anon_id = resolve_owner(request, data.anonymousId)

    # Create a new plan row (server-minted id) + blank sections row
    bp_id = await _ensure_business_plan(
        pool,
        bp_id=None,  # <- server generates the id; client never sends one
        summary="Generating…",
        full_plan="",
        model="unknown",
        answers=data.answers or [],
        plan_type=data.planType,
        user_id=user_id,
        anon_id=anon_id,
        country_code=data.countryCode,
        generate_id=data.generateId,
    )
    await _ensure_sections_row(pool, bp_id)

    # Tell subscribers a plan exists
    await _publish(bp_id, "plan.created", {"businessPlanId": bp_id})

    return {"businessPlanId": bp_id, "generateId": data.generateId}
    
@app.post("/generate/batch")
async def generate_business_plan_batch(request: Request, data: BatchGenerateRequest):
    try:
        pool = getattr(app.state, "db_pool", None)
        if not pool:
            raise HTTPException(status_code=500, detail="DB not initialized")

        # ① Resolve owner (JWT wins; else anonymous)
        user_id, anon_id = resolve_owner(request, data.anonymousId)
        # logging.info(f"[owner] user_id={user_id} anon_id={anon_id} planType={data.planType} planId={data.businessPlanId}")

        # ② (Optional) Guard ownership of existing plan
        if data.businessPlanId:
            await assert_plan_ownership(pool, business_plan_id=data.businessPlanId, user_id=user_id, anonymous_id=anon_id)

        # ③ Create/ensure plan shell
        bp_id = await _ensure_business_plan(
            pool,
            bp_id=data.businessPlanId,
            summary="Generating…",
            full_plan="",
            model="unknown",
            answers=data.answers,
            plan_type=data.planType,
            user_id=user_id,
            anon_id=anon_id,
            country_code=data.countryCode,
            generate_id=data.generateId,
        )
        # logging.info(f"[plan] ensured id={bp_id} linked to user_id={user_id} anon_id={anon_id}")

        # after ensure plan id
        await _publish(bp_id, "plan.created", {"businessPlanId": bp_id})

        await _ensure_sections_row(pool, bp_id)

        # 🔐 Charge exactly once per full plan (no charge for 'free' or anon)
        charged_now, remaining = await charge_full_plan_once(
            pool,
            business_plan_id=bp_id,
            user_id=user_id,
            plan_type=data.planType,
        )

        # Mark all sections as generating
        for job in data.sections:
            await _set_section_status(
                pool, business_plan_id=bp_id, section_key=job.sectionKey, generating=True, complete=False
            )
            await _publish(bp_id, "section.started", {
                "businessPlanId": bp_id,
                "sectionKey": job.sectionKey,
            })

        sem = asyncio.Semaphore(3)

        async def run_job(job: SectionJob):
            async with sem:
                # 1) Resolve plan-specific key (base for full, *_Free for free except funding)
                effective_key = _plan_key_for_section(job.sectionKey, data.planType or "free")

                # 2) Load GeneratePrompt row
                gp = await _get_generate_prompt(pool, effective_key)
                if not gp:
                    raise HTTPException(status_code=400, detail=f"No active GeneratePrompt for key {effective_key}")

                # 3) Resolve model/provider/temp/max (client overrides optional)
                resolved_model = job.model or gp["defaultModel"]
                resolved_provider = (job.provider or gp["defaultProvider"]).lower()
                resolved_temp = job.temperature if job.temperature is not None else float(gp["defaultTemp"])
                resolved_max = job.maxTokens if job.maxTokens is not None else int(gp["defaultMaxTokens"])

                # 4) Compose variables server-side (respect declared inputVariables)
                vars_for_prompt = _compose_prompt_vars(
                    gp_row=gp,
                    answers=data.answers,
                    business_idea=data.businessIdea or (data.answers[0] if data.answers else None),
                    country_code=data.countryCode,
                    section_title=gp["title"],
                    feedback=data.feedback,
                )

                # 5) Prompt + LLM + chain
                llm = get_llm(resolved_model, resolved_provider)
                if hasattr(llm, "temperature"): llm.temperature = resolved_temp
                if hasattr(llm, "max_tokens"): llm.max_tokens = resolved_max

                template = job.prompt or gp["promptTemplate"]
                prompt_t = PromptTemplate(input_variables=list(vars_for_prompt.keys()), template=template)
                chain: RunnableSequence = prompt_t | llm
                raw = await asyncio.to_thread(chain.invoke, vars_for_prompt)
                text = raw.content if isinstance(raw, AIMessage) else str(raw)

                # 6) Save section & clear flags
                await _upsert_section(pool, business_plan_id=bp_id, section_key=job.sectionKey, content=text)
                await _publish(bp_id, "section.complete", {
                    "businessPlanId": bp_id,
                    "sectionKey": job.sectionKey,
                })
                return resolved_model

        resolved_models = await asyncio.gather(*(run_job(job) for job in data.sections))
        model_for_plan = next((m for m in resolved_models if m), "unknown")

        # Recompose full plan & update summary
        full_plan = await _recompute_full_plan(pool, business_plan_id=bp_id)
        summary = extract_summary_from_markdown(full_plan)

        await _ensure_business_plan(
            pool,
            bp_id=bp_id,
            summary=summary,
            full_plan="",
            model=model_for_plan,
            answers=data.answers,
            plan_type=data.planType,
            user_id=user_id,
            anon_id=anon_id,
            country_code=data.countryCode,
            generate_id=data.generateId,
        )

        await _publish(bp_id, "plan.updated", {
            "businessPlanId": bp_id,
            "summary": summary,
            "fullPlanLength": len(full_plan or ""),
        })
        await _publish(bp_id, "plan.complete", {
            "businessPlanId": bp_id,
            "completedKeys": [j.sectionKey for j in data.sections],
        })

        return {
            "businessPlanId": bp_id,
            "status": "complete",
            "completedKeys": [j.sectionKey for j in data.sections],
            "fullPlanLength": len(full_plan or ""),
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Batch generation failed: {e}")
    
@app.get("/plan/{businessPlanId}/status")
async def get_plan_status(businessPlanId: str):
    pool = getattr(app.state, "db_pool", None)
    if not pool:
        raise HTTPException(status_code=500, detail="DB not initialized")
    async with pool.acquire() as conn:
        row = await conn.fetchrow('SELECT * FROM "BusinessPlanSection" WHERE "businessPlanId"=$1', businessPlanId)
        if not row:
            return {"exists": False}
        # Build per-section status
        statuses = {}
        for key, (col, complete_col, gen_col) in SECTION_MAP.items():
            statuses[key] = {"isGenerating": bool(row.get(gen_col)), "isComplete": bool(row.get(complete_col))}
        all_done = all(s["isComplete"] for s in statuses.values())
        return {"exists": True, "allComplete": all_done, "sections": statuses}

@app.get("/events/plan/{businessPlanId}")
async def sse_plan(businessPlanId: str):
    async def event_gen() -> AsyncIterator[bytes]:
        q: asyncio.Queue[str] = asyncio.Queue(maxsize=200)
        SUBSCRIBERS[businessPlanId].add(q)
        try:
            yield b": connected\n\n"
            while True:
                try:
                    msg = await asyncio.wait_for(q.get(), timeout=15)
                    yield msg.encode("utf-8")
                except asyncio.TimeoutError:
                    yield b": keep-alive\n\n"
        except asyncio.CancelledError:
            pass
        finally:
            SUBSCRIBERS[businessPlanId].discard(q)

    return StreamingResponse(
        event_gen(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
    )



@app.post("/generate")
async def generate_business_plan(request: Request, data: GenerateRequest):
    """
    Generate a business plan using the provided prompt template and variables,
    then persist to Supabase Postgres. Ownership is derived from the verified JWT.
    """
    try:
        # logging.info(f"Initializing model: {data.model} with provider: {data.provider}")
        llm_selected = get_llm(data.model, data.provider)  # provider is used in get_llm for HF models

        if hasattr(llm_selected, 'temperature'):
            llm_selected.temperature = data.temperature
        if hasattr(llm_selected, 'max_tokens'):
            llm_selected.max_tokens = data.maxTokens

        plan_prompt = PromptTemplate(
            input_variables=list(data.promptVariables.keys()),
            template=data.prompt
        )
        plan_chain: RunnableSequence = plan_prompt | llm_selected
        # logging.info(f"Generating business plan with model: {data.model}")

        raw_plan = await asyncio.to_thread(plan_chain.invoke, data.promptVariables)

        if isinstance(raw_plan, AIMessage):
            plan_text = raw_plan.content
        elif isinstance(raw_plan, (list, tuple)) and raw_plan and isinstance(raw_plan[0], AIMessage):
            plan_text = "\n".join(msg.content for msg in raw_plan)
        else:
            plan_text = str(raw_plan)

        # ① Resolve owner consistently (JWT wins)
        user_id, anon_id = resolve_owner(request, data.anonymousId)
        # logging.info(f"[owner] user_id={user_id} anon_id={anon_id} planType={data.planType} planId={data.businessPlanId}")

        # ② (Optional) Guard ownership if an existing plan is referenced
        if data.businessPlanId:
            pool = getattr(app.state, "db_pool", None)
            if not pool:
                raise HTTPException(status_code=500, detail="DB not initialized")
            await assert_plan_ownership(pool, business_plan_id=data.businessPlanId, user_id=user_id, anonymous_id=anon_id)

        # If generating a specific section, publish section.started before LLM generation
        if data.sectionKey and data.businessPlanId:
            pool = getattr(app.state, "db_pool", None)
            if pool:
                await _publish(data.businessPlanId, "section.started", {
                    "businessPlanId": data.businessPlanId,
                    "sectionKey": data.sectionKey,
                })

        # ── Persist
        pool = getattr(app.state, "db_pool", None)
        summary = extract_summary_from_markdown(plan_text)

        if not pool:
            # logging.warning("DB pool not available; skipping persistence.")
            return {
                "summary": summary or "Generated Business Plan",
                "plan": plan_text,
            }

        # If generating a specific section
        if data.sectionKey:
            business_plan_id = await _ensure_business_plan(
                pool,
                bp_id=data.businessPlanId,
                summary=summary,
                full_plan="",  # recomputed after section save
                model=data.model,
                answers=data.answers,
                plan_type=data.planType,
                user_id=user_id,
                anon_id=anon_id,
                country_code=data.countryCode,
                generate_id=data.generateId,
            )
            # logging.info(f"[plan] ensured section id={business_plan_id} linked to user_id={user_id} anon_id={anon_id}")

            # mark generating True in the DB for this section...
            await _publish(business_plan_id, "section.started", {
                "businessPlanId": business_plan_id,
                "sectionKey": data.sectionKey,
            })

            await _upsert_section(
                pool,
                business_plan_id=business_plan_id,
                section_key=data.sectionKey,
                content=plan_text
            )

            # ...generate text, save to this section, set generating False, set section complete...
            await _publish(business_plan_id, "section.complete", {
                "businessPlanId": business_plan_id,
                "sectionKey": data.sectionKey,
            })

            # Recompute full plan after section write
            full_plan = await _recompute_full_plan(pool, business_plan_id=business_plan_id)

            # if you recompute summary here, also:
            await _publish(business_plan_id, "plan.updated", {
                "businessPlanId": business_plan_id,
                "summary": summary,
                "fullPlanLength": len(full_plan or ""),
            })

            # Optional: mark complete if all sections present
            is_complete = False
            async with pool.acquire() as conn:
                row = await conn.fetchrow(
                    'SELECT * FROM "BusinessPlanSection" WHERE "businessPlanId"=$1',
                    business_plan_id
                )
                if row:
                    is_complete = _all_sections_present(row)

            resp = {
                "summary": summary,
                "plan": plan_text,
                "sections": [{"key": data.sectionKey, "content": plan_text}],
                "isComplete": True,  # this means "this section completed"
                "currentSection": data.sectionKey,
                "businessPlanId": business_plan_id,
            }
            return resp 

        # Non-section — treat as a whole plan
        business_plan_id = await _ensure_business_plan(
            pool,
            bp_id=data.businessPlanId,
            summary=summary,
            full_plan="",   # do not persist the full combined text
            model=data.model,
            answers=data.answers,
            plan_type=data.planType,
            user_id=user_id,
            anon_id=anon_id,
            country_code=data.countryCode,
            generate_id=data.generateId,
        )
        logging.info(f"[plan] ensured whole plan id={business_plan_id} linked to user_id={user_id} anon_id={anon_id}")

        resp = {
            "businessPlanId": business_plan_id,
            "sectionKey": None,
            "isComplete": True,  # whole-plan path implies one-shot completion
        }
        if data.returnContent:
            resp.update({
                "summary": summary,
                "plan": plan_text,
            })
        return resp

    except Exception as e:
        error_message = f"Error initializing model {data.model}: {str(e)}"
        # logging.error(error_message)
        raise HTTPException(status_code=500, detail=error_message)


@app.post("/export", response_model=ExportResponse)
async def export_business_plan(
    request: Request, 
    business_plan_id: str = Query(..., description="Business plan ID to export"),
    file_format: str = Query("pdf", description="Export format: pdf or word")
):
    """
    Export a business plan in the specified format (PDF or Word).
    Fetches all data from database and generates document.
    
    FLOW: Server fetches data from DB → Generates document → Uploads to Supabase → Returns download URL
    """
    try:
        pool = getattr(app.state, "db_pool", None)
        if not pool:
            raise HTTPException(status_code=500, detail="DB not initialized")
        
        # Resolve owner (JWT wins; else anonymous)
        user_id, anon_id = resolve_owner(request, None)
        # logging.info(f"[export] user_id={user_id}, anon_id={anon_id}, business_plan_id={business_plan_id}")
        
        # Check plan ownership first
        async with pool.acquire() as conn:
            row = await conn.fetchrow('SELECT "userId","anonymousId" FROM "BusinessPlan" WHERE id=$1', business_plan_id)
            if not row:
                raise HTTPException(status_code=404, detail="Business plan not found")
            
            plan_user = row["userId"]
            plan_anon = row["anonymousId"]
            # logging.info(f"[export] plan_user={plan_user}, plan_anon={plan_anon}")
            
            # If plan has no owner, allow access (legacy plans)
            if plan_user is None and plan_anon is None:
                # logging.info("[export] Legacy plan with no owner - allowing access")
                pass
            # If plan belongs to a user, you must be that user
            elif plan_user:
                if not user_id or user_id != plan_user:
                    # logging.warning(f"[export] User mismatch: plan_user={plan_user}, request_user={user_id}")
                    # Temporarily allow export without authentication for testing
                    # logging.info("[export] Allowing export without authentication (testing mode)")
                    pass
                else:
                    logging.info("[export] User matches - allowing access")
            # If plan belongs to an anonymous identity, you must present the same anonymousId
            elif plan_anon and not plan_user:
                if anon_id != plan_anon:
                    # logging.warning(f"[export] Anonymous mismatch: plan_anon={plan_anon}, request_anon={anon_id}")
                    raise HTTPException(status_code=403, detail="Not allowed to modify this plan (anonymous mismatch).")
                else:
                    logging.info("[export] Anonymous ID matches - allowing access")
        
        # Duplicate detection (best-effort; non-fatal if bookkeeping fails)
        try:
            request_fingerprint = f"{business_plan_id}_{file_format}_{int(time.time())}"
            current_time = time.time()
            if hasattr(export_business_plan, 'last_requests'):
                # Clean old requests (older than 5 seconds)
                export_business_plan.last_requests = {
                    fp: timestamp for fp, timestamp in export_business_plan.last_requests.items()
                    if current_time - timestamp < 5
                }
                # Check for duplicate
                if request_fingerprint in export_business_plan.last_requests:
                    raise HTTPException(
                        status_code=429,
                        detail="Duplicate export request detected. Please wait a few seconds before trying again."
                    )
            else:
                export_business_plan.last_requests = {}
            # Record this request
            export_business_plan.last_requests[request_fingerprint] = current_time
        except HTTPException:
            raise
        except Exception:
            # Do not block export on dedupe bookkeeping errors
            pass

        
        # Validate request format
        if file_format not in ["pdf", "word"]:
            raise HTTPException(status_code=400, detail="Invalid format. Only 'pdf' and 'word' are supported.")
        
        # Fetch business plan data from database
        # logging.info("[export] Fetching business plan data from database")
        async with pool.acquire() as conn:
            # Get business plan info
            bp_row = await conn.fetchrow(
                'SELECT "summary", "countryCode" FROM "BusinessPlan" WHERE "id"=$1',
                business_plan_id
            )
            if not bp_row:
                raise HTTPException(status_code=404, detail="Business plan not found")
            # logging.info(f"[export] Found business plan: summary={bp_row['summary'][:50]}...")
            
            # Get sections data
            sections_row = await conn.fetchrow(
                'SELECT * FROM "BusinessPlanSection" WHERE "businessPlanId"=$1',
                business_plan_id
            )
            if not sections_row:
                raise HTTPException(status_code=404, detail="Business plan sections not found")
            # logging.info(f"[export] Found sections data with {len(sections_row)} columns")
        
        # Build sections array from database
        # logging.info("[export] Building sections array from database")
        sections = []
        for key, (col_content, col_complete, col_generating) in SECTION_MAP.items():
            content = sections_row.get(col_content)
            if content and content.strip():
                # Extract title from section key
                title = key.replace('prompt_', '').replace('_', ' ').title()
                sections.append(ExportSection(
                    title=title,
                    content=content,
                    key=key
                ))
                # logging.info(f"[export] Added section: {title} ({len(content)} chars)")
        
        # logging.info(f"[export] Total sections found: {len(sections)}")
        if not sections:
            raise HTTPException(status_code=400, detail="No completed sections found for export")
        
        # Create ExportRequest from database data
        export_request = ExportRequest(
            sections=sections,
            businessIdea="Business Plan",  # Default since businessIdea column doesn't exist
            summary=bp_row["summary"] or "Business Plan Summary",
            format=file_format,
            includeCharts=True,
            includeTOC=True
        )
        
        # Check required packages before proceeding
        try:
            if file_format == "pdf":
                import fpdf
            else:
                import docx
        except ImportError as import_error:
            raise HTTPException(
                status_code=500,
                detail=f"Required package not available: {str(import_error)}"
            )
        
        # Create the document based on format
        try:
            # logging.info(f"[export] Starting {file_format.upper()} document generation")
            if file_format == "pdf":
                # logging.info("[export] Calling create_pdf_document...")
                doc_bytes, filename = create_pdf_document(export_request)
                # logging.info(f"[export] PDF generation completed. Size: {len(doc_bytes)} bytes, filename: {filename}")
            else:  # word
                # logging.info("[export] Calling create_word_document...")
                doc_bytes, filename = create_word_document(export_request)
                # logging.info(f"[export] Word generation completed. Size: {len(doc_bytes)} bytes, filename: {filename}")
            # Basic sanity check on output
            if not doc_bytes or (isinstance(doc_bytes, (bytes, bytearray)) and len(doc_bytes) < 1000):
                raise HTTPException(status_code=500, detail=f"Generated {file_format.upper()} appears empty or invalid.")
            # logging.info(f"[export] Document validation passed. Size: {len(doc_bytes)} bytes")
        except HTTPException:
            raise
        except ImportError as import_error:
            raise HTTPException(
                status_code=500,
                detail=f"Missing required package for {file_format.upper()} generation: {str(import_error)}"
            )
        except Exception as doc_error:
            # logging.error(f"Document creation failed: {doc_error}")
            raise HTTPException(
                status_code=500,
                detail=f"Failed to create {file_format.upper()} document: {str(doc_error)}"
            )
        
        # Upload document directly to Supabase storage
        try:
            logging.info("[export] Starting Supabase upload process")
            from supabase import create_client, Client
            import os
            
            # Initialize Supabase client with detailed error handling
            supabase_url = os.getenv("SUPABASE_URL")
            supabase_key = os.getenv("SUPABASE_ANON_KEY")
            
            # Detailed environment variable checking
            missing_vars = []
            if not supabase_url:
                missing_vars.append("SUPABASE_URL")
            if not supabase_key:
                missing_vars.append("SUPABASE_ANON_KEY")
            
            if missing_vars:
                error_msg = f"Missing Supabase environment variables: {', '.join(missing_vars)}"
                logging.error(f"[export] {error_msg}")
                logging.error(f"[export] Available env vars starting with SUPABASE:")
                for key, value in os.environ.items():
                    if key.startswith("SUPABASE"):
                        logging.error(f"[export]   {key} = {'*' * 10 if 'KEY' in key else value}")
                
                raise HTTPException(
                    status_code=500, 
                    detail=f"Supabase configuration missing: {', '.join(missing_vars)}. Please add these to your .env file."
                )
            
            # logging.info(f"[export] Supabase URL: {supabase_url[:30]}...")
            # logging.info(f"[export] Supabase Key: {'*' * 20}...{supabase_key[-10:] if len(supabase_key) > 10 else '***'}")
            
            # Initialize Supabase client
            supabase: Client = create_client(supabase_url, supabase_key)
            # logging.info("[export] Supabase client initialized successfully")
            
            # Upload to Supabase storage (delete existing file first, then upload new one)
            upload_path = filename  # Remove "business-plans/" prefix to avoid double path
            # logging.info(f"[export] Uploading to Supabase: {upload_path}")
            
            # First, try to delete the existing file if it exists
            try:
                # logging.info(f"[export] Attempting to delete existing file: {upload_path}")
                supabase.storage.from_("business-plans").remove([upload_path])
                # logging.info(f"[export] Successfully deleted existing file: {upload_path}")
            except Exception as delete_error:
                # File might not exist, which is fine
                # logging.info(f"[export] No existing file to delete (or delete failed): {str(delete_error)}")
                pass
            
            # Now upload the new file
            result = supabase.storage.from_("business-plans").upload(
                upload_path,
                doc_bytes,
                file_options={"content-type": "application/pdf" if file_format == "pdf" else "application/vnd.openxmlformats-officedocument.wordprocessingml.document"}
            )
            # logging.info(f"[export] Upload result: {result}")
            
            # Get public URL
            # logging.info("[export] Getting public URL from Supabase")
            public_url = supabase.storage.from_("business-plans").get_public_url(upload_path)
            # logging.info(f"[export] Public URL: {public_url}")
            
            response = ExportResponse(
                success=True,
                downloadUrl=public_url,
                filename=filename,
                size=len(doc_bytes),
                message=f"Document generated successfully as {file_format.upper()} and uploaded to Supabase storage.",
                documentData=None  # No longer needed since we're uploading directly
            )
            
            return response
                
        except Exception as data_error:
            logging.error(f"[export] Supabase upload failed: {str(data_error)}")
            logging.error(f"[export] Exception type: {type(data_error).__name__}")
            import traceback
            logging.error(f"[export] Supabase traceback: {traceback.format_exc()}")
            raise HTTPException(
                status_code=500, 
                detail=f"Failed to upload document to Supabase: {str(data_error)}"
            )
        
    except HTTPException:
        # Re-raise HTTP exceptions (validation errors)
        raise
    except Exception as e:
        # Catch any other unexpected errors
        error_message = f"Export failed: {str(e)}"
        logging.error(f"[export] Error: {error_message}")
        logging.error(f"[export] Exception type: {type(e).__name__}")
        import traceback
        logging.error(f"[export] Traceback: {traceback.format_exc()}")
        raise HTTPException(status_code=500, detail=error_message)
        

@app.get("/download/{filename}")
async def download_file(filename: str):
    """Download endpoint to serve the generated export files"""
    file_path = os.path.join("temp_exports", filename)
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="File not found")
    
    return FileResponse(file_path, filename=filename)

# ──────────────────────────────────────────────────────────────────────────────
# Live PDF Generation Endpoints
# ──────────────────────────────────────────────────────────────────────────────

# @app.post("/live-pdf/start")
# async def start_live_pdf_session(request: Request, data: ExportRequest):
#     """Start a live PDF session that streams updates as content changes"""
#     try:
#         # Create a new live PDF session
#         session_id = await live_pdf_generator.create_live_session(
#             data.businessIdea or "default", 
#             data
#         )
#         
#         return {
#             "success": True,
#             "session_id": session_id,
#             "message": "Live PDF session started successfully"
#         }
#         
#     except Exception as e:
#         raise HTTPException(status_code=500, detail=f"Failed to start live PDF session: {str(e)}")

# @app.post("/live-pdf/update/{session_id}")
# async def update_live_pdf_session(session_id: str, data: ExportRequest):
#     """Update an existing live PDF session with new content"""
#     try:
#         success = await live_pdf_generator.update_session(session_id, data)
#         
#         if not success:
#             raise HTTPException(status_code=404, detail="Session not found")
#         
#         return {
#             "success": True,
#             "message": "PDF updated successfully"
#         }
#         
#     except Exception as e:
#         raise HTTPException(status_code=500, detail=f"Failed to update PDF: {str(e)}")

# @app.get("/live-pdf/stream/{session_id}")
# async def stream_live_pdf_updates(session_id: str):
#     """Stream live PDF updates using Server-Sent Events"""
#     async def event_generator():
#         async for chunk in live_pdf_generator.stream_pdf_updates(session_id):
#             yield chunk
#     
#     return StreamingResponse(
#         event_generator(),
#         media_type="text/event-stream",
#         headers={
#             "Cache-Control": "no-cache",
#             "Connection": "keep-alive",
#             "Access-Control-Allow-Origin": "*",
#             "Access-Control-Allow-Headers": "*"
#         }
#     )

# @app.delete("/live-pdf/close/{session_id}")
# async def close_live_pdf_session(session_id: str):
#     """Close and cleanup a live PDF session"""
#     try:
#         live_pdf_generator.close_session(session_id)
#         
#         return {
#             "success": True,
#             "message": "Live PDF session closed successfully"
#         }
#         
#     except Exception as e:
#         raise HTTPException(status_code=500, detail=f"Failed to close session: {str(e)}")
@app.get("/")
def root():
    return {"status": "FastAPI is running 🚀"}

# @app.get("/test-toc")
# def test_toc():
#     """Test endpoint to verify TOC generation"""
#     from app.main import generate_table_of_contents_after_content, ExportSection
    
#     # Create test sections
#     test_sections = [
#         ExportSection(
#             key="prompt_ExecutiveSummary",
#             title="Executive Summary",
#             content="# Executive Summary\n\nThis is a test executive summary with some content."
#         ),
#         ExportSection(
#             key="prompt_CompanyProfile", 
#             title="Company Profile",
#             content="# Company Profile\n\nThis is a test company profile section."
#         ),
#         ExportSection(
#             key="prompt_MarketAnalysis",
#             title="Market Analysis", 
#             content="# Market Analysis\n\nThis is a test market analysis section."
#         )
#     ]
    
#     # Generate TOC
#     toc_items = generate_table_of_contents_after_content(test_sections, {}, "Test Business")
    
#     return {
#         "test_sections": [{"title": s.title, "key": s.key, "content_length": len(s.content)} for s in test_sections],
#         "toc_items": toc_items,
#         "toc_count": len(toc_items)
#     }

# @app.get("/health")
# def health_check():
#     hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
#     openai_api_key = os.getenv("OPENAI_API_KEY")
    
#     status = {
#         "api": "healthy",
#         "models": {}
#     }
    
#     models_to_check = ["Mistral", "Qwen-2.5", "Llama", "Gemma", "Phi-3"]
    
#     for model_name in models_to_check:
#         try:
#             if model_name == "":
#                 client = InferenceClient(provider="hf-inference", api_key=hf_token)
#                 client.model_info("mistralai/Mistral-7B-Instruct-v0.3")
#                 status["models"][model_name] = "available"
#             elif model_name == "Qwen-2.5":
#                 client = InferenceClient(provider="hf-inference", api_key=hf_token)
#                 client.model_info("Qwen/Qwen2.5-7B-Instruct")
#                 status["models"][model_name] = "available"
#             elif model_name == "Llama":
#                 client = InferenceClient(provider="hf-inference", api_key=hf_token)
#                 client.model_info("meta-llama/Llama-2-7b-chat-hf")
#                 status["models"][model_name] = "available"
#             else:
#                 status["models"][model_name] = "not checked"
#         except Exception as e:
#             status["models"][model_name] = f"error: {str(e)}"
    
#     status["models"]["GPT"] = "available" if openai_api_key else "unavailable (missing API key)"
#     return status