import typing
import dspy
from dspy import Signature, InputField, Module, ChainOfThought, Predict
import streamlit as st
from typing import Optional
from openinference.instrumentation.dspy import DSPyInstrumentor
from opentelemetry import trace as trace_api
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk import trace as trace_sdk
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace.export import SimpleSpanProcessor

endpoint = "http://127.0.0.1:6006/v1/traces"
resource = Resource(attributes={})
tracer_provider = trace_sdk.TracerProvider(resource=resource)
span_otlp_exporter = OTLPSpanExporter(endpoint=endpoint)
tracer_provider.add_span_processor(
    SimpleSpanProcessor(span_exporter=span_otlp_exporter)
)
trace_api.set_tracer_provider(tracer_provider=tracer_provider)
DSPyInstrumentor().instrument()


base_lm = dspy.OpenAI(
    model="gpt-4",
    max_tokens=4096,
    api_key=st.secrets["OpenAI"],
)
dspy.settings.configure(lm=base_lm)


# refin_lm = dspy.OpenAI(
#     model="gpt-3.5-turbo",
#     max_tokens=4096,
#     api_key="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",
# )

_template = """A few controversial things I believe about {Topic}:

1) {StrongOpinion1}

I’m not saying {SubtleDistinction}.

I’m just saying {ClarifyingStatement}.

- 

2) {StrongOpinion2}

Contrary to popular belief, {ClicheAction} is a myth.

- {FailureExample1}
- {FailureExample2}
- {FailureExample3}

{ClicheAction} is a fast-track to {UndesirableOutcome}.

- 

3) {StrongOpinion3}

“Most people” do what “most people” do.

Which leads to…

- {NegativeStat1}
- {NegativeStat2}
- {NegativeStat3}

Clearly what “most people do” doesn’t work.

So do something else.

- 

4) {StrongOpinion4}

99% of the time, this is a giant mistake.

Instead, {Name} has a great framework for achieving {Outcome} without falling into {Trap}.

- {Step1}
- {Step2}
- {Step3}
- 

5) {StrongOpinion5}

Pursuing {Outcome} is a silly goal.

Because {UnconventionalReason}.

{DifferentOutcome} is a better pursuit in life.

- {Benefit1}
- {Benefit2}
- {Benefit3}"""


class GenerateLinkedinArticle(Signature):
    topic = InputField(desc="Title or brief description of content")
    template = InputField(prefix="template to follow")
    purpose = InputField()
    audience = InputField()
    tone_style = InputField()
    key_points = InputField()
    num_words = InputField()
    language = InputField()
    article = dspy.OutputField(desc=f"Linkedin article of {num_words} words")


class LinkedinArticle(Module):
    def __init__(self):
        super().__init__()
        self.generate_article = ChainOfThought(GenerateLinkedinArticle)

    def forward(
        self,
        topic,
        template,
        purpose,
        audience,
        tone_style,
        key_points,
        num_words,
        language,
    ):
        article_prediction = self.generate_article(
            topic=topic,
            template=template,
            purpose=purpose,
            audience=audience,
            tone_style=tone_style,
            key_points=key_points,
            num_words=num_words,
            language=language,
        )
        return article_prediction.article


class Feedback(Module):
    def __init__(self):
        super().__init__()
        self.feedback = ChainOfThought("original_content, feedback -> refined_content")

    def forward(self, original_content, feedback):
        corrected_article = self.feedback(
            original_content=original_content, feedback=feedback
        ).refined_content
        return corrected_article


class Evaluator(dspy.Signature):
    """You are a creative writing coach, evaluate this linkedin post"""

    linkedin_post = InputField(prefix="A linkedin post to evaluate")
    output = dspy.OutputField(
        prefix="A comprehensive evaluation of the input post, styled in markdown"
    )


class SelfEval(Module):
    def __init__(self):
        super().__init__()
        self.self_eval = ChainOfThought(Evaluator)

    def forward(self, linkedin_post):
        eval = self.self_eval(linkedin_post=linkedin_post).output
        return eval


def write_article(user_inputs: dict, lm: Optional[dspy.dsp.LM] = base_lm) -> str:
    with dspy.context(lm=lm):
        article = LinkedinArticle()
        content = article.forward(**user_inputs)
    return content


def incorporate_feedback(
    original_content, feedback, lm: Optional[dspy.dsp.LM] = base_lm
):
    with dspy.context(lm=lm):
        refined_content = Feedback().forward(
            original_content=original_content, feedback=feedback
        )
    return refined_content


def evaluate_post(content: str, lm: Optional[dspy.dsp.LM] = base_lm) -> str:
    with dspy.context(lm=lm):
        eval = SelfEval().forward(linkedin_post=content)
    return eval