import streamlit as st
import openai
import os
import re
import json
from datetime import datetime
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
from reportlab.lib.utils import simpleSplit

# === Azure OpenAI wrapper ===
class AzureTextBot:
    def __init__(self, api_key, endpoint, deployment_id, api_version):
        self.client = openai.AzureOpenAI(
            api_key=api_key,
            azure_endpoint=endpoint,
            api_version=api_version,
        )
        self.deployment_id = deployment_id

    def _chat(self, prompt):
        messages = [
            {"role": "system", "content": "You are a helpful financial analyst assistant."},
            {"role": "user", "content": prompt},
        ]
        response = self.client.chat.completions.create(
            model=self.deployment_id,
            messages=messages,
            temperature=0,
        )
        return response.choices[0].message.content.strip()

    # Extraction prompts

    def extract_earnings_elements(self, transcript):
        prompt = f"""
You are a financial data extraction assistant.

From the transcript below, extract ALL relevant earnings data needed to create the Earnings Summary.

Output a JSON object with these fields:

- companies: list of {{name, ticker, quarter_end (Dec/Jan/Feb)}}
- trading_day: {{up: [{{ticker, percent}}], down: [{{ticker, percent}}]}}
- trading_week: same format
- results: {{sales: {{beat: [tickers], miss: [tickers]}}, adj_eps: {{beat: [...], miss: [...], not_reported: [...]}}}}
- guidance: {{sales: {{provided: [...], not_provided: [...]}}, adj_eps: {{provided: [...], not_provided: [...]}}}}

Be 100% exact and output ONLY JSON.

Transcript:
\"\"\"{transcript}\"\"\"
"""
        return self._chat(prompt)

    def extract_hot_topics_elements(self, transcript):
        prompt = f"""
You are a financial analyst extraction assistant.

From the transcript below, extract ALL commentary needed for the Hot Topics summary.

Output a JSON with these keys and formats:

consumer_health: list of {{company, sentiment("Positive"/"Neutral"/"Negative"), commentary}}
qtd_performance: same format
fiscal_year_expectations: list of {{company, commentary}}
tariffs: list of {{company, commentary}}
weather: list of {{company, commentary}}
stock_initiations: list of {{company, commentary}}
largest_stock_movements: list of {{company, commentary}}

Output ONLY valid JSON.

Transcript:
\"\"\"{transcript}\"\"\"
"""
        return self._chat(prompt)

    # Validation with scoring prompt
    def validate_summary(self, transcript, summary, summary_type="earnings summary"):
        prompt = f"""
You are a financial analyst assistant.

Given the original transcript:

\"\"\"{transcript}\"\"\"

And this {summary_type} summary:

\"\"\"{summary}\"\"\"

Please:

- List any missing or incorrect companies, tickers, percentages, beats/misses, guidance or commentary compared to the transcript.
- On a scale of 0 to 100, rate the overall consistency (%) of the summary compared to the transcript.

Format your response EXACTLY as:

Consistency Score: [percentage]

Issues:
- [Issue 1]
- [Issue 2]
- [...]

If the summary is fully consistent with the transcript, respond:

Consistency Score: 100

Issues:
None
"""
        return self._chat(prompt)

    # Generation prompts with your exact prompt text
    def generate_earnings_summary(self, extracted_json):
        prompt = f"""
You are a financial data assistant. Your task is to carefully and methodically extract **all relevant earnings data** from the input context and produce a fully structured Earnings Summary.

⚠️ You must follow these instructions strictly:

1. **List every company and ticker mentioned in the text.**
2. Extract quarter-ends ONLY from December, January, or February and assign them correctly per company.
3. For the Trading Day of Print and Week Prior to Print, list all companies with exact percentage changes and whether Up or Down.
4. Extract for each company whether they beat or missed Consensus on Sales and Adj. EPS.
5. Extract whether FY guidance for Sales and Adj. EPS was provided or not, specifying companies accordingly.
6. Use the exact output format described below.
7. DO NOT omit any company or data point.
8. DO NOT infer or hallucinate any information not explicitly stated.
9. When sorting % changes, sort by the absolute value magnitude descending.
10. Preserve exact symbols (•, ), and 100% of the input data, **no rounding or summarization**.
11. If multiple companies have similar names or tickers, handle them independently.

---

Output format (use exactly):  

Earnings Summary  
• In addition to Macy’s, Inc. (M), X companies that we follow reported this week: [List all company names and tickers].  
[State the quarter-end month (Dec/Jan/Feb) for each company, based ONLY on input.]

o Trading Day of Print:  
 Up: [All companies with exact tickers and % gains, sorted largest to smallest]  
 Down: [All companies with exact tickers and % losses, sorted largest to smallest]

o Trading Week Prior to Print:  
 Up: [All companies with % weekly gains]  
 Down: [All companies with % weekly losses]

o Results vs. Consensus:  
 Sales: Beat – [Companies that beat]; Miss – [Companies that missed]  
 Adj. EPS: Beat – [Companies that beat]; Miss – [Companies that missed]  
*[Ticker] did not report adj. EPS (if applicable)*

o Annual Guidance Updates:  
 Sales: Provided by – [Tickers provided]; Not provided by – [Tickers not provided]  
 Adj EPS: Provided by – [Tickers provided]; Not provided by – [Tickers not provided]

---

Strict rules:  
- Use 100% of all relevant input data; no omission.  
- Do not hallucinate or infer anything not explicitly written.  
- Maintain exact output indentations, bullet points, and symbols.  
- Sort percentage changes strictly by magnitude, from highest to lowest.

---

Please complete each section fully before moving on.

---

Here is the extracted data (JSON):

{extracted_json}
"""
        return self._chat(prompt)

    def generate_hot_topics_summary(self, extracted_json):
        prompt = f"""
You are a professional financial analyst assistant.

Based solely on this structured Hot Topics data:

{extracted_json}

Generate a **detailed, professional Hot Topics summary** following exactly this structure and formatting:

---

Hot Topics  
• Begin with a paragraph summarizing overall tone or sentiment.

---

1. Consumer Health Commentary  
• Start with a detailed paragraph summarizing consumer behavior and demand trends.

• Then list company commentary grouped strictly under these **sentiment groups**:

**Positive**  
• **Company A:** Detailed commentary describing positive consumer health aspects.  
• **Company B:** Detailed commentary describing positive consumer health aspects.  

**Neutral**  
• **Company C:** Detailed neutral commentary.

**Negative**  
• **Company D:** Detailed negative commentary.

---

2. QTD Performance  
• Start with a detailed paragraph summarizing quarter-to-date trends.

• Then categorize company-specific comments under sentiment headings:

**Positive**  
• **Company A:** Detailed description of positive QTD performance.

**Neutral**  
• **Company B:** Detailed neutral QTD performance commentary.

**Negative**  
• **Company C:** Detailed negative QTD performance commentary.

---

3. Fiscal Year Expectations  
• Provide detailed summary of fiscal year guidance and outlook.

• Use bullet points starting with company names bolded, followed by guidance commentary.

---

4. Tariff Commentary  
• Provide a paragraph summarizing tariff or trade policy impacts.

• Then use bullet points with bolded company names and tariff-related commentary.

---

5. Weather Commentary  
• Summarize weather trends and impacts.

• List relevant companies with bolded names and commentary bullets.

---

6. Stock Initiations and Re-Ratings  
• Provide detailed summary paragraph.

• Bullet points with bolded company names and rating details.

---

7. Largest Stock Movements  
• Detailed summary paragraph of notable stock price changes.

• Bullet points with bolded company names and details.

---

Important notes:  
- Use **Positive / Neutral / Negative** tags **ONLY** in Consumer Health Commentary and QTD Performance sections.  
- Use bullet points (•) with company names in **bold** as shown.  
- Provide detailed descriptions — no one-liners.  
- Follow exact indents, bullet marks, and formatting shown.  
- Do NOT hallucinate or add information not in input.

---

Please complete each section fully before moving on.

---

Here is the extracted data (JSON):

{extracted_json}
    """
        return self._chat(prompt)
        


def create_pdf(content_dict, prefix="summary"):
    os.makedirs("temp", exist_ok=True)
    file_path = os.path.join("temp", f"{prefix}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf")
    c = canvas.Canvas(file_path, pagesize=letter)
    width, height = letter
    margin = 50
    y = height - margin
    line_height = 14
    max_width = width - 2 * margin

    for title, content in content_dict.items():
        if y < margin + 100:
            c.showPage()
            y = height - margin
        c.setFont("Helvetica-Bold", 16)
        c.drawString(margin, y, title)
        y -= line_height * 1.5

        for line in content.split("\n"):
            stripped_line = line.strip()
            if not stripped_line:
                y -= line_height * 0.5
                continue
            if stripped_line in ["Positive", "Neutral", "Negative"]:
                if y < margin + line_height:
                    c.showPage()
                    y = height - margin
                c.setFont("Helvetica-Bold", 13)
                c.drawString(margin + 10, y, stripped_line)
                y -= line_height * 1.2
                continue
            bullet_indent = 20
            if stripped_line.startswith("•") or stripped_line.startswith("-"):
                c.setFont("Helvetica", 11)
            else:
                c.setFont("Helvetica", 12)
            wrapped_lines = simpleSplit(
                stripped_line, c._fontname, c._fontsize, max_width - bullet_indent
            )
            for wrap_line in wrapped_lines:
                if y < margin + line_height:
                    c.showPage()
                    y = height - margin
                c.drawString(margin + bullet_indent, y, wrap_line)
                y -= line_height
        y -= line_height * 2
    c.save()
    return file_path


def clean_json_codeblock(text):
    """Strip markdown code fences and parse JSON."""
    cleaned = re.sub(r"```(?:json)?\n([\s\S]*?)```", r"\1", text).strip()
    return json.loads(cleaned)


def parse_validation_output(text):
    """Parse the consistency score and issues from validation prompt output."""
    score_match = re.search(r"Consistency Score:\s*(\d+)", text)
    issues = []
    score = 0
    if score_match:
        score = int(score_match.group(1))

    issues_match = re.search(r"Issues:\s*((?:- .*\n?)*)", text)
    if issues_match:
        issues_text = issues_match.group(1).strip()
        if issues_text and issues_text.lower() != "none":
            issues = [line.strip("- ").strip() for line in issues_text.splitlines() if line.strip()]
    return score, issues


def iterative_refine_summary(bot, transcript, extracted_json_str, generate_func, summary_type):
    MAX_RETRIES = 3
    last_summary = None
    last_validation = None
    last_score = 0
    notes = ""

    for attempt in range(1, MAX_RETRIES + 1):
        if attempt == 1:
            # First generation attempt
            summary = generate_func(extracted_json_str)
        else:
            # On retry, ask model to revise based on last summary + validation feedback
            prompt = f"""
The previously generated {summary_type} is inconsistent with the transcript.

Transcript:
\"\"\"{transcript}\"\"\"

Previous {summary_type}:
\"\"\"{last_summary}\"\"\"

Validation feedback:
\"\"\"{last_validation}\"\"\"

Please regenerate the complete and fully correct {summary_type}, fixing the issues highlighted.

Use the exact original formatting and instructions.
"""
            summary = bot._chat(prompt)

        # Validate summary
        validation = bot.validate_summary(transcript, summary, summary_type)
        score, issues = parse_validation_output(validation)

        # Prepare note if score >= 85 and issues present
        if 85 <= score < 100:
            notes = "\n\n⚠️ Note: The summary may be missing or inaccurately cover:\n" + "\n".join(f"• {issue}" for issue in issues)
        else:
            notes = ""

        if score >= 85:
            return summary + notes, validation, score
        else:
            last_summary = summary
            last_validation = validation
            last_score = score

    # Return last summary if max retries exceeded
    return (last_summary + (notes if last_score >= 85 else "")), last_validation, last_score


# --- Streamlit UI Setup ---
st.set_page_config(page_title="📈 Earnings & Hot Topics Summary with Consistency Check", layout="wide")
st.title("📈 Earnings and Hot Topics Summary Generator with Iterative Validation & Notes")

# Load secrets
api_key = st.secrets["AZURE_API_KEY"]
endpoint = st.secrets["AZURE_ENDPOINT"]
deployment_id = st.secrets["AZURE_DEPLOYMENT"]
api_version = st.secrets["AZURE_API_VERSION"]

bot = AzureTextBot(api_key, endpoint, deployment_id, api_version)

uploaded_file = st.file_uploader("Upload transcript (.txt)", type="txt")
if uploaded_file:
    transcript_text = uploaded_file.getvalue().decode("utf-8")
    st.text_area("Uploaded Transcript", transcript_text, height=300)

    if st.button("Run Full Summary Pipeline"):

        # Earnings extraction
        with st.spinner("Extracting Earnings structured data..."):
            raw_earnings_json = bot.extract_earnings_elements(transcript_text)
            try:
                earnings_json = clean_json_codeblock(raw_earnings_json)
                st.json(earnings_json)
            except Exception as e:
                st.error(f"Parsing earnings extraction JSON failed: {e}")
                st.text_area("Raw Earnings Extraction Output", raw_earnings_json)
                earnings_json = None

        if earnings_json:
            with st.spinner("Iterative generation and validation of Earnings Summary..."):
                earnings_summary, earnings_validation, earnings_score = iterative_refine_summary(
                    bot,
                    transcript_text,
                    json.dumps(earnings_json),
                    bot.generate_earnings_summary,
                    "earnings summary"
                )
                st.text_area("Earnings Summary", earnings_summary, height=400)
                st.text_area("Earnings Summary Validation Details", earnings_validation, height=150)
                st.write(f"Consistency Score: {earnings_score}%")
            if earnings_score < 85:
                st.warning("Earnings summary consistency below threshold. Please revise input or retry.")
        else:
            st.error("No valid earnings extraction data, skipping earnings summary.")

        st.markdown("---")

        # Hot Topics extraction
        with st.spinner("Extracting Hot Topics structured data..."):
            raw_hot_topics_json = bot.extract_hot_topics_elements(transcript_text)
            try:
                hot_topics_json = clean_json_codeblock(raw_hot_topics_json)
                st.json(hot_topics_json)
            except Exception as e:
                st.error(f"Parsing hot topics extraction JSON failed: {e}")
                st.text_area("Raw Hot Topics Extraction Output", raw_hot_topics_json)
                hot_topics_json = None

        if hot_topics_json:
            with st.spinner("Iterative generation and validation of Hot Topics Summary..."):
                hot_topics_summary, hot_topics_validation, hot_topics_score = iterative_refine_summary(
                    bot,
                    transcript_text,
                    json.dumps(hot_topics_json),
                    bot.generate_hot_topics_summary,
                    "hot topics summary"
                )
                st.text_area("Hot Topics Summary", hot_topics_summary, height=400)
                st.text_area("Hot Topics Summary Validation Details", hot_topics_validation, height=150)
                st.write(f"Consistency Score: {hot_topics_score}%")
            if hot_topics_score < 85:
                st.warning("Hot Topics summary consistency below threshold. Please revise input or retry.")
        else:
            st.error("No valid hot topics extraction data, skipping hot topics summary.")

        # PDF download if both summaries generated and at least 85% consistent
        if (earnings_json and hot_topics_json and
                earnings_score >= 85 and hot_topics_score >= 85):
            pdfs_content = {
                "Earnings Summary": earnings_summary,
                "Detailed Hot Topics Summary": hot_topics_summary,
            }
            pdf_file = create_pdf(pdfs_content)
            with open(pdf_file, "rb") as f:
                st.download_button(
                    "📥 Download Combined Summary PDF",
                    data=f,
                    file_name="combined_summary.pdf",
                    mime="application/pdf"
                )
        else:
            st.info("Both summaries must be at least 85% consistent to enable PDF download.")

# --- PDF generation function ---
def create_pdf(content_dict, prefix="summary"):
    os.makedirs("temp", exist_ok=True)
    file_path = os.path.join("temp", f"{prefix}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf")
    c = canvas.Canvas(file_path, pagesize=letter)
    width, height = letter
    margin = 50
    y = height - margin
    line_height = 14
    max_width = width - 2 * margin

    for title, content in content_dict.items():
        if y < margin + 100:
            c.showPage()
            y = height - margin
        c.setFont("Helvetica-Bold", 16)
        c.drawString(margin, y, title)
        y -= line_height * 1.5

        for line in content.split("\n"):
            stripped_line = line.strip()
            if not stripped_line:
                y -= line_height * 0.5
                continue
            if stripped_line in ["Positive", "Neutral", "Negative"]:
                if y < margin + line_height:
                    c.showPage()
                    y = height - margin
                c.setFont("Helvetica-Bold", 13)
                c.drawString(margin + 10, y, stripped_line)
                y -= line_height * 1.2
                continue
            bullet_indent = 20
            if stripped_line.startswith("•") or stripped_line.startswith("-"):
                c.setFont("Helvetica", 11)
            else:
                c.setFont("Helvetica", 12)
            wrapped_lines = simpleSplit(stripped_line, c._fontname, c._fontsize, max_width - bullet_indent)
            for wrap_line in wrapped_lines:
                if y < margin + line_height:
                    c.showPage()
                    y = height - margin
                c.drawString(margin + bullet_indent, y, wrap_line)
                y -= line_height
        y -= line_height * 2
    c.save()
    return file_path

with this code when we prepare a json then we genrate a summary and validate it on the real transcript what i required is the validation gives some points until it become 0 or give validation of suumary is fully correct untill we can again send the validation point to llm telling it to prepare on these missing validation points and after that again genrate summary and check its validation