import os
import webbrowser
import pandas as pd
from pathlib import Path

from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.models.openai_codex import OpenAICodexModel
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.providers.openai_codex import (
    OpenAICodexOAuthFlow,
    OpenAICodexProvider,
)

from .config import CodexConfig, LLMConfig, OpenAiConfig
from .credentials import FileCodexCredentionalSource
from .schemas import BaseDatasetRecord, ChatRecord, ChatMessage


class LLMModelProvider:

    @staticmethod
    def create_openai_compatible(model_name: str,
                                 api_key: str,
                                 base_url: str
                                ) -> OpenAICodexModel:

        provider = OpenAIProvider(base_url=base_url,
                                api_key=api_key)

        return OpenAIChatModel(model_name=model_name,
                            provider=provider)

    @staticmethod
    async def create_codex(model_name: str='gpt-5.6-luna',
                           credentional_path: str | Path | None=None
                           ) -> OpenAICodexModel:

        if credentional_path is None:
            credentional_path = Path.cwd() / ".config" / "codex-credentials.json"

        
        source  = FileCodexCredentionalSource(credentional_path)

        if not source.path.exists():
            flow = OpenAICodexOAuthFlow()
            webbrowser.open(flow.authorization_url())
            credentials = await flow.exchange_code_from_callback()
            await source.save(credentials)

        provider = OpenAICodexProvider(credential_source=source)

        return OpenAICodexModel(model_name=model_name,
                                provider=provider)

    @staticmethod
    async def create_from_config(config: LLMConfig):

        if isinstance(config, CodexConfig):
            return await LLMModelProvider.create_codex(
                model_name=config.model,
                credentional_path=config.credential_path
            )

        if isinstance(config, OpenAiConfig):
            api_key = os.getenv(config.api_key_env)
            base_url=os.getenv(config.base_url_env)

            if not api_key:
                raise ValueError(
                    f"Environment variable "
                    f"{config.api_key_env!r} is not set"
                )

            if not base_url:
                raise ValueError(
                    f"Environment variable "
                    f"{config.base_url_env!r} is not set"
                )

            return LLMModelProvider.create_openai_compatible(
                model_name=config.model,
                api_key=api_key,
                base_url=base_url
            )

        raise ValueError(f"Unsupported LLM config: {type(config)}")

class DatasetWriter:
    def __init__(self,
                 dataset_record: BaseDatasetRecord,
                 base_path: str | Path):


        self._dataset_record = dataset_record    
        self._path = Path(base_path)
        self._path.mkdir(parents=True, exist_ok=True)

    def __check_file_exists(self, filename :str)-> bool:
        file_path = self._path / filename
        return file_path.exists()

    def to_csv(self,
               filename: str) -> None:

        exists=self.__check_file_exists(filename)
        
        df = pd.json_normalize(
            self._dataset_record.model_dump()
        )

        df.to_csv(self._path / filename,
                  mode="a" if exists else "w",
                  header=not exists,
                  index=False,
                  encoding="utf-8")

    def to_chat_record(self) -> ChatRecord:
        response = self._dataset_record.response
        prompts = self._dataset_record.prompts

        assistant_content = (
            response.assistant_content
            or response.model_response
        )

        return ChatRecord(
            messages=[
                ChatMessage(
                    role="system",
                    content=prompts.system_prompt,
                ),
                ChatMessage(
                    role="user",
                    content=prompts.user_prompt,
                ),
                ChatMessage(
                    role="assistant",
                    content=assistant_content,
                ),
            ]
        )

    def to_jsonl(self,
                 filename: str) -> None:
        
        chat_record = self.to_chat_record()
        path = self._path/filename

        with path.open("a", encoding='utf-8') as file:
            file.write(chat_record.model_dump_json())
            file.write("\n")