import os,re,json,gc,asyncio,calendar,traceback,faiss,torch,numpy as np,pdfplumber
from datetime import date,datetime,timedelta
from dateutil.relativedelta import relativedelta
from sentence_transformers import SentenceTransformer,CrossEncoder
from transformers import AutoTokenizer,AutoModelForCausalLM
from fastapi import FastAPI,HTTPException,UploadFile,File,Form
from fastapi.responses import FileResponse,StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from pypdf import PdfReader
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4
from reportlab.lib.units import mm
from reportlab.lib.enums import TA_LEFT,TA_JUSTIFY
from reportlab.lib.styles import getSampleStyleSheet,ParagraphStyle
from reportlab.platypus import SimpleDocTemplate,Paragraph,Spacer,Table,TableStyle,Image
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

BASE_PATH=os.environ.get("BASE_PATH","/app")
MODEL_NAME=os.environ.get("MODEL_NAME","Qwen/Qwen2.5-3B-Instruct")
SCHEMA_PATH=os.path.join(BASE_PATH,"json_files","schema.json")
METRIC_PATH=os.path.join(BASE_PATH,"json_files","metric_catalog.json")
COMPANY_DATA_PATH=os.path.join(BASE_PATH,"json_files","company_data.json")
PROMPTS_PATH=os.path.join(BASE_PATH,"json_files","prompts.json")
KNOWLEDGE_META_PATH=os.path.join(BASE_PATH,"report_knowledge_base","report_metadata.json")
KNOWLEDGE_INDEX_PATH=os.path.join(BASE_PATH,"report_knowledge_base","report_vectors.index")
OUTPUT_DIR=os.path.join(BASE_PATH,"output")
os.makedirs(OUTPUT_DIR,exist_ok=True)

app=FastAPI(title="Financial Intelligence API",version="1.0.0")
app.add_middleware(CORSMiddleware,allow_origins=["*"],allow_credentials=True,allow_methods=["*"],allow_headers=["*"])

def load_json(path):
    with open(path,"r",encoding="utf-8") as f:return json.load(f)

def save_json(data,path):
    directory=os.path.dirname(path)
    if directory:os.makedirs(directory,exist_ok=True)
    with open(path,"w",encoding="utf-8") as f:json.dump(data,f,indent=2,ensure_ascii=False)

print("Loading company data...");company_data=load_json(COMPANY_DATA_PATH)
print("Loading prompts...");PROMPTS=load_json(PROMPTS_PATH)
print("Loading knowledge metadata...");knowledge_metadata=load_json(KNOWLEDGE_META_PATH)
print("Loading embedding model...");embedding_model=SentenceTransformer("BAAI/bge-small-en-v1.5")
print("Loading reranker...");reranker=CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
print("Loading FAISS index...");knowledge_index=faiss.read_index(KNOWLEDGE_INDEX_PATH)
print("Loading tokenizer...");tokenizer=AutoTokenizer.from_pretrained(MODEL_NAME,trust_remote_code=True)
print("Loading Qwen model...")
model=AutoModelForCausalLM.from_pretrained(MODEL_NAME,device_map="auto",torch_dtype=torch.float16,trust_remote_code=True)
model.eval()
print("Qwen model loaded successfully.")
print("CUDA available:",torch.cuda.is_available())
if torch.cuda.is_available():print("GPU:",torch.cuda.get_device_name(0))

def safe_text(x):
    if isinstance(x,str):return x
    if isinstance(x,dict):
        for key in ["query","text","metric","name","field","value"]:
            if key in x and isinstance(x[key],(str,int,float)):return str(x[key])
        return " ".join(str(value) for value in x.values())
    if isinstance(x,list):return " ".join(safe_text(value) for value in x)
    return str(x or "")

def is_date_value(value):
    if not isinstance(value,str):return False
    for fmt in ["%Y-%m-%d","%d-%m-%Y","%m/%d/%Y","%d/%m/%Y","%Y/%m/%d","%Y-%m-%d %H:%M:%S","%Y-%m-%dT%H:%M:%S"]:
        try:datetime.strptime(value,fmt);return True
        except ValueError:pass
    return False

def detect_date_fields(rows):
    if not rows:return []
    fields=[]
    for column in rows[0]:
        detected=any(is_date_value(row.get(column)) for row in rows if row.get(column)is not None)
        name_detected=any(word in column.lower() for word in ["date","time","year"])
        if detected or name_detected:fields.append(column)
    return fields

def extract_schema(data):
    output={}
    for table,rows in data.items():
        if isinstance(rows,list) and rows and isinstance(rows[0],dict):
            sample=rows[0]
            output[table]={"columns":list(sample.keys()),"column_types":{col:type(val).__name__ for col,val in sample.items()},"numeric_fields":[col for col,val in sample.items() if isinstance(val,(int,float)) and not isinstance(val,bool)],"date_fields":detect_date_fields(rows)}
        elif isinstance(rows,dict):
            output[table]={"columns":list(rows.keys()),"column_types":{col:type(val).__name__ for col,val in rows.items()},"numeric_fields":[col for col,val in rows.items() if isinstance(val,(int,float)) and not isinstance(val,bool)],"date_fields":[col for col,val in rows.items() if is_date_value(val) or any(w in col.lower() for w in ["date","time","year"])]}
    return output

schema=load_json(SCHEMA_PATH) if os.path.exists(SCHEMA_PATH) else (lambda s:save_json(s,SCHEMA_PATH) or s)(extract_schema(company_data))

def build_metric_catalog(schema):
    catalog={}
    for table,info in schema.items():
        for field in info["columns"]:
            if field in info.get("date_fields",[]):continue
            catalog[f"{table}*{field}"]={"table":table,"field":field,"display_name":field.replace("_"," ").title(),"field_type":"number" if field in info.get("numeric_fields",[]) else info["column_types"].get(field)}
    return catalog

metric_catalog=load_json(METRIC_PATH) if os.path.exists(METRIC_PATH) else (lambda c:save_json(c,METRIC_PATH) or c)(build_metric_catalog(schema))
metric_names=list(metric_catalog.keys())
metric_texts=[f"{metric_catalog[name]['table']} {metric_catalog[name]['display_name']}".lower() for name in metric_names]
print("Building metric embeddings...")
metric_embeddings=embedding_model.encode(metric_texts,normalize_embeddings=True).astype("float32")
metric_index=faiss.IndexFlatIP(metric_embeddings.shape[1])
metric_index.add(metric_embeddings)

def cleanup_gpu():
    gc.collect()
    if torch.cuda.is_available():torch.cuda.empty_cache()

llm_lock=asyncio.Lock()

def _llm_json_sync(messages,max_tokens=1500):
    inputs=None;output=None
    try:
        prompt=tokenizer.apply_chat_template(messages,tokenize=False,add_generation_prompt=True)
        inputs=tokenizer(prompt,return_tensors="pt",truncation=True,max_length=tokenizer.model_max_length)
        device=next(model.parameters()).device
        inputs={k:v.to(device) for k,v in inputs.items()}
        with torch.inference_mode():
            output=model.generate(**inputs,max_new_tokens=max_tokens,do_sample=False,use_cache=True,pad_token_id=tokenizer.eos_token_id)
        generated=output[0][inputs["input_ids"].shape[-1]:]
        text=re.sub(r"```(?:json)?","",tokenizer.decode(generated,skip_special_tokens=True),flags=re.IGNORECASE).replace("```","").strip()
        try:return json.loads(text)
        except Exception:pass
        for pat in [r"\{.*\}",r"\[.*\]"]:
            if match:=re.search(pat,text,re.DOTALL):
                try:return json.loads(match.group(0))
                except Exception:pass
        print("WARNING: Invalid LLM JSON:",text[:2000])
        return {}
    except Exception:
        traceback.print_exc();return {}
    finally:
        del inputs,output
        cleanup_gpu()

async def llm_json(messages,max_tokens=1500):
    async with llm_lock:return await asyncio.to_thread(_llm_json_sync,messages,max_tokens)

def parse_financial_query(query):
    system="""Return JSON only:
{"queries":[],"operations":[]}
Extract independent financial data requests as plain strings and required operations.
Do not extract database fields or answer the question.
Operations:
sum,count,average,compare,trend,percentage,group_by,status,explain_change
Do not create queries for operations.
Keep time information inside relevant queries."""
    result=_llm_json_sync([{"role":"system","content":system},{"role":"user","content":query}],300)
    if not isinstance(result,dict):result={}
    queries=result.get("queries",[])
    if not isinstance(queries,list):queries=[query]
    queries=[safe_text(item).strip() for item in queries if safe_text(item).strip()] or [query]
    operations=result.get("operations",[])
    if not isinstance(operations,list):operations=[]
    return {"queries":queries,"operations":[safe_text(item).lower().strip() for item in operations]}

def metric_embedding_rank(query,k=5):
    query=safe_text(query).strip()
    if not query:return []
    embedding=embedding_model.encode([query],normalize_embeddings=True).astype("float32")
    k=min(k,len(metric_names))
    if k<=0:return []
    scores,indices=metric_index.search(embedding,k)
    results=[]
    for score,index in zip(scores[0],indices[0]):
        if index==-1:continue
        metric_name=metric_names[index];metric=metric_catalog[metric_name]
        results.append({"metric":metric_name,"table":metric["table"],"field":metric["field"],"display_name":metric["display_name"],"field_type":metric["field_type"],"embedding_score":float(score)})
    return results

def metric_semantic_rank(query,k=5):
    candidates=metric_embedding_rank(query,k)
    if not candidates:return []
    pairs=[[query.lower(),f"{item['table']} {item['display_name']}".lower()] for item in candidates]
    scores=reranker.predict(pairs)
    for item,score in zip(candidates,scores):item["semantic_score"]=float(score)
    return sorted(candidates,key=lambda item:item["semantic_score"],reverse=True)

def retrieve_metrics(parsed):
    return {safe_text(query):metric_semantic_rank(query,5) for query in parsed.get("queries",[])}

MONTHS={name.lower():index for index,name in enumerate(calendar.month_name) if name}

def resolve_date_range(query):
    query=safe_text(query).lower();today=date.today()
    if "this month" in query:
        start=today.replace(day=1);end=today.replace(day=calendar.monthrange(today.year,today.month)[1])
        return {"start_date":start.isoformat(),"end_date":end.isoformat()}
    if "last month" in query:
        end=today.replace(day=1)-timedelta(days=1)
        return {"start_date":end.replace(day=1).isoformat(),"end_date":end.isoformat()}
    if match:=re.search(r"last\s+(\d+)\s+months?",query):
        count=max(1,int(match.group(1)));start=today.replace(day=1)-relativedelta(months=count-1);end=today.replace(day=calendar.monthrange(today.year,today.month)[1])
        return {"start_date":start.isoformat(),"end_date":end.isoformat()}
    if match:=re.search(r"\b(january|february|march|april|may|june|july|august|september|october|november|december)\s+(\d{4})\s+to\s+(january|february|march|april|may|june|july|august|september|october|november|december)\s+(\d{4})\b",query):
        sm,sy,em,ey=MONTHS[match.group(1)],int(match.group(2)),MONTHS[match.group(3)],int(match.group(4))
        return {"start_date":date(sy,sm,1).isoformat(),"end_date":date(ey,em,calendar.monthrange(ey,em)[1]).isoformat()}
    if match:=re.search(r"\b(\d{2})/(\d{4})\s+to\s+(\d{2})/(\d{4})\b",query):
        sm,sy,em,ey=map(int,match.groups())
        if 1<=sm<=12 and 1<=em<=12:return {"start_date":date(sy,sm,1).isoformat(),"end_date":date(ey,em,calendar.monthrange(ey,em)[1]).isoformat()}
    if match:=re.search(r"from\s+(\d{2})/(\d{2})/(\d{4})\s+to\s+(\d{2})/(\d{2})/(\d{4})",query):
        d1,m1,y1,d2,m2,y2=map(int,match.groups())
        return {"start_date":date(y1,m1,d1).isoformat(),"end_date":date(y2,m2,d2).isoformat()}
    if match:=re.search(r"from\s+(\d{4})\s+to\s+(\d{4})",query):return {"start_date":f"{match.group(1)}-01-01","end_date":f"{match.group(2)}-12-31"}
    if match:=re.search(r"\b(january|february|march|april|may|june|july|august|september|october|november|december)\s+(\d{4})\b",query):
        m,y=MONTHS[match.group(1)],int(match.group(2))
        return {"start_date":date(y,m,1).isoformat(),"end_date":date(y,m,calendar.monthrange(y,m)[1]).isoformat()}
    if match:=re.search(r"\b(\d{2})/(\d{4})\b",query):
        m,y=map(int,match.groups())
        if 1<=m<=12:return {"start_date":date(y,m,1).isoformat(),"end_date":date(y,m,calendar.monthrange(y,m)[1]).isoformat()}
    return None

def build_intent(parsed,results):
    output=[]
    for query,metrics in results.items():
        if not metrics:continue
        metric=metrics[0];date_range=resolve_date_range(query)
        output.append({"query":query,"metric":{"table":metric["table"],"field":metric["field"],"name":metric["display_name"],"metric_id":metric["metric"],"semantic_score":metric.get("semantic_score",0)},"operations":parsed.get("operations",[]),"filters":{"date_range":date_range} if date_range else {}})
    return output

def get_date_field(row):
    return next((column for column in row if any(word in column.lower() for word in ["date","paid_on","approved_on","created_at","timestamp","month"])),None)

def retrieve_answer_data(intent,data):
    table=intent["metric"]["table"];field=intent["metric"]["field"];date_range=intent.get("filters",{}).get("date_range");rows=data.get(table,[])
    if isinstance(rows,dict):rows=[rows]
    if not isinstance(rows,list):return []
    result=[]
    for row in rows:
        if not isinstance(row,dict):continue
        date_field=get_date_field(row);row_date=str(row.get(date_field))[:10] if date_field and row.get(date_field)is not None else None
        if date_range and (not row_date or not(date_range["start_date"]<=row_date<=date_range["end_date"])):continue
        result.append({"date":row_date,intent["metric"]["name"]:row.get(field)})
    return result

def calculate_metric(intent,data_result):
    metric=intent["metric"]["name"];operations=intent.get("operations",[])
    values=[{"date":row.get("date"),"value":row.get(metric)} for row in data_result if isinstance(row,dict) and isinstance(row.get(metric),(int,float)) and not isinstance(row.get(metric),bool)]
    if not values:return {"metric":metric,"status":"no numeric data","data":[]}
    result={"metric":metric,"data":values}
    if "sum" in operations:result={"metric":metric,"sum":sum(x["value"] for x in values),"data":values}
    if "count" in operations:result={"metric":metric,"count":len(values),"data":values}
    if "average" in operations:result={"metric":metric,"average":round(sum(x["value"] for x in values)/len(values),2),"data":values}
    if "percentage" in operations and len(values)>=2:
        previous,current=values[-2],values[-1];result["percentage_change"]=round(((current["value"]-previous["value"])/previous["value"]*100) if previous["value"] else 0,2)
    if "compare" in operations and len(values)>=2:
        previous,current=values[-2],values[-1];difference=current["value"]-previous["value"]
        result["comparison"]={"current_period":current,"previous_period":previous,"current_value":current["value"],"previous_value":previous["value"],"difference":difference,"percentage_change":round((difference/previous["value"]*100) if previous["value"] else 0,2)}
    if "trend" in operations and len(values)>=2:result={"metric":metric,"periods":values,"overall_change":values[-1]["value"]-values[0]["value"]}
    return result

def retrieve_history(query,k=5):
    query=safe_text(query).strip()
    if not query:return []
    embedding=embedding_model.encode([query],normalize_embeddings=True).astype("float32");k=min(k,len(knowledge_metadata))
    if k<=0:return []
    scores,indices=knowledge_index.search(embedding,k);candidates=[]
    for score,index in zip(scores[0],indices[0]):
        if index==-1:continue
        item=knowledge_metadata[index] if isinstance(knowledge_metadata[index],dict) else {};content=safe_text(item.get("content",""))
        if len(content.split())>=5:candidates.append({"id":int(index),"content":content,"embedding_score":float(score)})
    if not candidates:return []
    scores=reranker.predict([[query.lower(),x["content"].lower()] for x in candidates])
    for item,score in zip(candidates,scores):item["semantic_score"]=float(score)
    candidates.sort(key=lambda x:x["semantic_score"],reverse=True)
    return [{"content":x["content"],"embedding_score":x["embedding_score"],"semantic_score":x["semantic_score"]} for x in candidates[:3]]

def execute_calculations(intents):
    return [{"query":intent,"calculation":calculate_metric(intent,retrieve_answer_data(intent,company_data))} for intent in intents]

def prepare_context(parsed,intents):
    context=[]
    for intent in intents:
        rows=retrieve_answer_data(intent,company_data)
        context.append({"query":intent,"intent":intent,"current_data":rows,"calculation":calculate_metric(intent,rows),"historical_knowledge":retrieve_history(intent["query"])})
    return context

def generate_financial_response(context):
    system=PROMPTS.get("financial_report",{}).get("system_prompt","\nReturn a structured financial report as JSON with:\n\nreport_title,\nreport_type,\nexecutive_summary,\nsections,\ntables,\ncharts,\nconclusion.\n")
    result=_llm_json_sync([{"role":"system","content":system},{"role":"user","content":json.dumps(context,default=str,ensure_ascii=False)}],1800)
    if not isinstance(result,dict):result={}
    defaults={"report_title":"Financial Analysis Report","report_type":"financial_report","executive_summary":"","sections":[],"tables":[],"charts":[],"conclusion":""}
    for key,value in defaults.items():
        if key not in result:result[key]=value
    for key in ["sections","tables","charts"]:
        if not isinstance(result.get(key),list):result[key]=[]
    return result

def extract_knowledge_from_pdf(path):
    with pdfplumber.open(path) as pdf:text="\n".join(p.extract_text(layout=True) or "" for p in pdf.pages)
    rows=[]
    for i in range(0,len(text),2500):
        res=_llm_json_sync([{"role":"system","content":PROMPTS["report_extraction"]["system"]},{"role":"user","content":text[i:i+2500]}],600)
        rows.extend(res if isinstance(res,list) else next((res[k] for k in ["rows","data","results","knowledge"] if isinstance(res,dict) and isinstance(res.get(k),list)),[]))
    clean=[];seen=set()
    for r in rows:
        if isinstance(r,dict):
            c=safe_text(r.get("content","")).strip()
            if c:
                key=(r.get("report_type"),safe_text(r.get("section","")).strip(),c)
                if key not in seen:seen.add(key);clean.append(r)
    existing={(x.get("report_type"),x.get("section"),x.get("content")) for x in knowledge_metadata}
    new_rows=[r for r in clean if (r.get("report_type"),r.get("section"),r.get("content")) not in existing]
    if new_rows:
        vectors=np.array(embedding_model.encode([x["content"] for x in new_rows]),dtype="float32");faiss.normalize_L2(vectors);knowledge_index.add(vectors);knowledge_metadata.extend(new_rows);faiss.write_index(knowledge_index,KNOWLEDGE_INDEX_PATH);save_json(knowledge_metadata,KNOWLEDGE_META_PATH)
    return {"extracted":len(rows),"cleaned":len(clean),"added":len(new_rows),"total_knowledge":knowledge_index.ntotal}

def clean_text(text):return re.sub(r"[\*#`_]","",safe_text(text)).strip()

def pdf_styles():
    styles=getSampleStyleSheet()
    return {"title":ParagraphStyle("financial_title",parent=styles["Title"],fontName="Helvetica-Bold",fontSize=21,leading=25,textColor=colors.HexColor("#1A365D"),alignment=TA_LEFT,spaceAfter=6),"sub":ParagraphStyle("financial_subtitle",parent=styles["Normal"],fontSize=9,textColor=colors.HexColor("#64748B"),spaceAfter=12),"head":ParagraphStyle("financial_heading",parent=styles["Heading2"],fontName="Helvetica-Bold",fontSize=14,leading=18,textColor=colors.HexColor("#2E6F9E"),spaceBefore=12,spaceAfter=7),"body":ParagraphStyle("financial_body",parent=styles["BodyText"],fontSize=10,leading=15,textColor=colors.HexColor("#334155"),alignment=TA_JUSTIFY,spaceAfter=8)}

def add_page_number(canvas,document):
    canvas.saveState();canvas.setFont("Helvetica",8);canvas.setFillColor(colors.HexColor("#64748B"));canvas.drawRightString(A4[0]-18*mm,10*mm,f"Page {document.page}");canvas.restoreState()

def add_pdf_table(elements,table_data,styles):
    headers=table_data.get("headers",[]);rows=table_data.get("rows",[])
    if not headers or not rows:return
    header_style=ParagraphStyle("financial_table_header",parent=styles["body"],fontName="Helvetica-Bold",textColor=colors.white,fontSize=9)
    data=[[Paragraph(clean_text(v),header_style) for v in headers]]
    for row in rows:
        if isinstance(row,list):data.append([Paragraph(clean_text(v),styles["body"]) for v in row])
    if len(data)<2:return
    table=Table(data,repeatRows=1,hAlign="LEFT")
    table.setStyle(TableStyle([("BACKGROUND",(0,0),(-1,0),colors.HexColor("#1A365D")),("GRID",(0,0),(-1,-1),0.4,colors.HexColor("#CBD5E1")),("ROWBACKGROUNDS",(0,1),(-1,-1),[colors.white,colors.HexColor("#F8FAFC")]),("VALIGN",(0,0),(-1,-1),"TOP"),("LEFTPADDING",(0,0),(-1,-1),7),("RIGHTPADDING",(0,0),(-1,-1),7),("TOPPADDING",(0,0),(-1,-1),6),("BOTTOMPADDING",(0,0),(-1,-1),6)]))
    if table_data.get("title"):elements.append(Paragraph(clean_text(table_data["title"]),styles["head"]))
    elements.extend([table,Spacer(1,10)])

def create_chart(chart,index):
    data=chart.get("data",[]) if isinstance(chart,dict) else [];labels=[];values=[]
    for item in data:
        if isinstance(item,dict):
            try:labels.append(str(item.get("label","")));values.append(float(item.get("value",0)))
            except Exception:continue
    if len(values)<2:return None
    chart_type=str(chart.get("type","line")).lower();path=os.path.join(OUTPUT_DIR,f"chart_{index}.png")
    plt.figure(figsize=(8,4.5))
    if chart_type=="bar":plt.bar(labels,values)
    elif chart_type=="pie":plt.pie(values,labels=labels,autopct="%1.1f%%")
    elif chart_type=="histogram":plt.bar(labels,values,width=1)
    else:plt.plot(labels,values,marker="o")
    plt.title(clean_text(chart.get("title","Financial Overview")));plt.xticks(rotation=30,ha="right");plt.tight_layout();plt.savefig(path,dpi=180,bbox_inches="tight");plt.close()
    return path

def create_pdf(response,path):
    styles=pdf_styles();document=SimpleDocTemplate(path,pagesize=A4,leftMargin=18*mm,rightMargin=18*mm,topMargin=18*mm,bottomMargin=18*mm)
    elements=[Paragraph(clean_text(response.get("report_title","Financial Analysis Report")),styles["title"]),Paragraph(datetime.now().strftime("%B %d, %Y"),styles["sub"])]
    if summary:=response.get("executive_summary",""):elements.extend([Paragraph("Executive Summary",styles["head"]),Paragraph(clean_text(summary),styles["body"])])
    for section in response.get("sections",[]):
        if isinstance(section,dict):
            if section.get("title"):elements.append(Paragraph(clean_text(section["title"]),styles["head"]))
            if section.get("content"):elements.append(Paragraph(clean_text(section["content"]),styles["body"]))
    for table_data in response.get("tables",[])[:2]:
        if isinstance(table_data,dict):add_pdf_table(elements,table_data,styles)
    chart_paths=[]
    for index,chart in enumerate(response.get("charts",[])[:2]):
        if isinstance(chart,dict) and (chart_path:=create_chart(chart,index)):
            chart_paths.append(chart_path);elements.extend([Paragraph(clean_text(chart.get("title","Financial Overview")),styles["head"]),Image(chart_path,width=165*mm,height=90*mm),Spacer(1,8)])
    if conclusion:=response.get("conclusion",""):elements.extend([Paragraph("Conclusion",styles["head"]),Paragraph(clean_text(conclusion),styles["body"])])
    document.build(elements,onFirstPage=add_page_number,onLaterPages=add_page_number)
    for p in chart_paths:
        try:os.remove(p)
        except Exception:pass
    return path

def extract_pdf_text(path):
    reader=PdfReader(path)
    return "\n".join(page.extract_text() or "" for page in reader.pages).strip()

class PromptRequest(BaseModel):
    prompt:str

def progress_event(message,step="progress",data=None):
    payload={"step":step,"message":message}
    if isinstance(data,dict):payload.update(data)
    return f"data: {json.dumps(payload,default=str,ensure_ascii=False)}\n\n"

@app.get("/")
async def health():return {"status":"ok","model":MODEL_NAME,"gpu":torch.cuda.is_available()}

@app.get("/health")
async def health_check():
    gpu_name=torch.cuda.get_device_name(0) if torch.cuda.is_available() else None
    return {"status":"healthy","gpu":torch.cuda.is_available(),"gpu_name":gpu_name,"model":MODEL_NAME}

@app.post("/generate-report")
async def generate_report(req:PromptRequest):
    try:
        query=re.sub(r",+"," ",req.prompt).strip()
        if not query:raise HTTPException(400,"prompt is required")
        parsed=parse_financial_query(query);metrics=retrieve_metrics(parsed);intents=build_intent(parsed,metrics)
        if not intents:raise HTTPException(400,"Unable to identify a financial metric")
        context=prepare_context(parsed,intents);calculations=execute_calculations(intents)
        if not any(item["current_data"] for item in context):raise HTTPException(404,"No financial data found matching request criteria")
        response=generate_financial_response({"user_query":query,"financial_context":context,"calculations":calculations})
        if not response.get("executive_summary") and not response.get("sections"):raise HTTPException(500,"Unable to generate financial report")
        path=os.path.join(OUTPUT_DIR,"financial_report.pdf");create_pdf(response,path)
        return FileResponse(path,media_type="application/pdf",filename="financial_report.pdf")
    except HTTPException:raise
    except Exception as error:
        traceback.print_exc();raise HTTPException(500,f"{type(error).__name__}: {str(error)}")

@app.post("/generate-report-stream")
async def generate_report_stream(req:PromptRequest):
    async def event_stream():
        try:
            query=re.sub(r",+"," ",req.prompt).strip()
            if not query:yield progress_event("Prompt is required","error");return
            yield progress_event("Understanding your request","parsing")
            parsed=await asyncio.to_thread(parse_financial_query,query)
            yield progress_event("Identifying the relevant financial information","metric")
            metrics=await asyncio.to_thread(retrieve_metrics,parsed);intents=await asyncio.to_thread(build_intent,parsed,metrics)
            if not intents:yield progress_event("I couldn't identify the requested financial information","error");return
            yield progress_event("Retrieving financial data","data_retrieval")
            context=await asyncio.to_thread(prepare_context,parsed,intents)
            if not any(item["current_data"] for item in context):yield progress_event("No financial data was found for your request","error");return
            yield progress_event("Calculating financial results","calculation")
            calculations=await asyncio.to_thread(execute_calculations,intents)
            yield progress_event("Analyzing the results","analysis")
            response=await asyncio.to_thread(generate_financial_response,{"user_query":query,"financial_context":context,"calculations":calculations})
            if not response.get("executive_summary") and not response.get("sections"):yield progress_event("Unable to generate the financial report","error");return
            if summary:=safe_text(response.get("executive_summary","")).strip():yield progress_event(summary,"analysis_result")
            yield progress_event("Generating PDF report","pdf_generation")
            path=os.path.join(OUTPUT_DIR,"financial_report.pdf");await asyncio.to_thread(create_pdf,response,path)
            if not os.path.exists(path):yield progress_event("PDF generation failed","error");return
            data={"filename":"financial_report.pdf","download_endpoint":"/download-pdf"}
            yield progress_event("PDF generated successfully. Starting download.","pdf_ready",data)
            yield progress_event("Download ready","complete",data)
        except Exception:
            traceback.print_exc();yield progress_event("An error occurred while generating the report","error")
    return StreamingResponse(event_stream(),media_type="text/event-stream",headers={"Cache-Control":"no-cache","Connection":"keep-alive","X-Accel-Buffering":"no"})

@app.post("/generate-report-with-pdf")
async def generate_report_with_pdf(prompt:str=Form(...),file:UploadFile=File(...)):
    try:
        query=re.sub(r",+"," ",prompt).strip()
        if not query:raise HTTPException(400,"prompt is required")
        filename=os.path.basename(file.filename);upload_path=os.path.join(OUTPUT_DIR,f"uploaded_{filename}");contents=await file.read()
        with open(upload_path,"wb") as output_file:output_file.write(contents)
        pdf_text=extract_pdf_text(upload_path)
        if not pdf_text:raise HTTPException(400,"Unable to extract text from PDF")
        pdf_text=pdf_text[:12000];parsed=parse_financial_query(query);metrics=retrieve_metrics(parsed);intents=build_intent(parsed,metrics);context=prepare_context(parsed,intents) if intents else []
        context.append({"query":query,"intent":{},"current_data":[],"calculation":{},"historical_knowledge":[{"content":pdf_text}]})
        response=generate_financial_response({"user_query":query,"financial_context":context,"calculations":execute_calculations(intents) if intents else [],"uploaded_document_text":pdf_text})
        if not response.get("executive_summary") and not response.get("sections"):raise HTTPException(500,"Unable to generate financial report")
        path=os.path.join(OUTPUT_DIR,"financial_report.pdf");create_pdf(response,path)
        return FileResponse(path,media_type="application/pdf",filename="financial_report.pdf")
    except HTTPException:raise
    except Exception as error:
        traceback.print_exc();raise HTTPException(500,f"{type(error).__name__}: {str(error)}")

@app.post("/upload-pdf")
async def upload_pdf(file:UploadFile=File(...)):
    try:
        filename=os.path.basename(file.filename);path=os.path.join(OUTPUT_DIR,f"uploaded_{filename}");contents=await file.read()
        with open(path,"wb") as output_file:output_file.write(contents)
        history=await asyncio.to_thread(extract_knowledge_from_pdf,path)
        return {"status":"success","filename":file.filename,"path":path,"knowledge_extraction":history}
    except Exception as error:
        traceback.print_exc();raise HTTPException(500,f"{type(error).__name__}: {str(error)}")

@app.get("/download-pdf")
async def download_pdf():
    path=os.path.join(OUTPUT_DIR,"financial_report.pdf")
    if not os.path.exists(path):raise HTTPException(404,"No report available")
    return FileResponse(path,media_type="application/pdf",filename="financial_report.pdf")

@app.on_event("shutdown")
async def shutdown_event():
    gc.collect()
    if torch.cuda.is_available():torch.cuda.empty_cache()